chore: import upstream snapshot with attribution

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

@@ -0,0 +1,23 @@
+# isort: skip_file
+# 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.
+"""LLM support for PyTorch-like API to build IRModules."""
+
+from . import kv_cache, position_embedding
+from .position_embedding import llama_rope
+from .tree_attn import tree_attn
+from .kv_cache import PagedKVCache

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

@@ -0,0 +1,526 @@
+# 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
+# fmt: off
+
+"""Single-token decode attention kernels and attention-state merge helpers.
+
+Contents:
+- ``_attention_decode_cpu`` / ``_attention_decode`` — paged-KV decode (one Q token
+  per sequence), CPU scalar and GPU allreduce variants.
+- ``_merge_state_inplace_cpu`` / ``_merge_state_inplace`` — combine two
+  log-sum-exp attention outputs in place. Used by multi-stage decoding and by
+  the distributed KV-transfer path.
+"""
+
+# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
+import math
+from typing import Any
+
+from tvm.script import tirx as T
+from tvm.target import Target
+
+from ._kernel_common import (
+    _declare_length_info,
+    _get_kv_chunk_len,
+    _get_seq_offset,
+    _rope,
+    _var,
+    _var_cpu,
+    check_thread_limits,
+    get_max_num_threads_per_block,
+)
+
+
+def _attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_window: bool, rope_scaling: dict[str, Any], page_size: int = 16):
+    H_qo = num_qo_heads
+    H_kv = num_kv_heads
+    D = head_dim
+    group_size = num_qo_heads // num_kv_heads
+
+    global_symbol = "batch_decode_paged_kv_cpu"
+    if sliding_window:
+        global_symbol += "_sliding_window"
+
+    @T.prim_func(s_tir=True)
+    def batch_decode_paged_kv(
+        Q_handle: T.handle,
+        pages_handle: T.handle,
+        page_table_indptr_handle: T.handle,
+        page_table_values_handle: T.handle,
+        var_length_info: T.handle,  # [b] when sliding window = False, or otherwise [3, b]
+        k_rope_pos_offset_handle: T.handle,
+        q_rope_position_handle: T.handle,
+        output_handle: T.handle,
+        lse_handle: T.handle,
+        rotary_mode: T.int32,
+        rope_scale: T.float32,
+        rope_theta: T.float32,
+        sm_scale: T.float32,
+    ):
+        T.func_attr({"tirx.is_scheduled": True, "global_symbol": global_symbol})
+        B = T.int32()
+        nnz_pages = T.int32()
+        max_num_pages = T.int32()
+        page_indptr_elem_offset = T.int32()
+        page_values_elem_offset = T.int32()
+        k_rope_pos_offset_elem_offset = T.int32()
+        q_rope_position_elem_offset = T.int32()
+        length_info_elem_offset = T.int32()
+
+        Q = T.match_buffer(Q_handle, (B, H_qo, D), qkv_dtype)
+        pages = T.match_buffer(pages_handle, (max_num_pages, 2, H_kv, page_size, D), qkv_dtype)
+        page_table_indptr = T.match_buffer(page_table_indptr_handle, (B + 1,), "int32", elem_offset=page_indptr_elem_offset)
+        page_table_values = T.match_buffer(page_table_values_handle, (nnz_pages,), "int32", elem_offset=page_values_elem_offset)
+        k_rope_pos_offset = T.match_buffer(k_rope_pos_offset_handle, (B,), "int32", elem_offset=k_rope_pos_offset_elem_offset)
+        q_rope_position = T.match_buffer(q_rope_position_handle, (B,), "int32", elem_offset=q_rope_position_elem_offset)
+        output = T.match_buffer(output_handle, (B, H_qo, D), qkv_dtype)
+        lse = T.match_buffer(lse_handle, (B, H_qo), "float32")  # pylint: disable=unused-variable
+        # The length information of the sequences.
+        # - It is in shape `(3, batch_size)` when sliding window is enabled.
+        #   For a sequence "i", location
+        #   - "(0, i)" is the number of KV slots used in the last page of the seq ("last_page_len"),
+        #   - "(1, i)" is the starting offset of the sliding window in the seq,
+        #   - "(2, i)" is the attn sink length of the sequence.
+        # - It is in shape `(batch_size,)` when sliding window is disabled,
+        #   denoting the "last_page_len".
+        length_info = _declare_length_info(var_length_info, B, sliding_window, length_info_elem_offset)
+
+        for b in T.serial(B):
+            with T.sblock("attn"):
+                O_local = T.sblock_alloc_buffer((D,), "float32")
+                Q_local = T.sblock_alloc_buffer((D,), "float32")
+                K_local = T.sblock_alloc_buffer((D,), "float32")
+                V_local = T.sblock_alloc_buffer((D,), "float32")
+
+                kv_chunk_len = T.sblock_alloc_buffer((1,), "int32")
+
+                m_val = T.sblock_alloc_buffer((1,), "float32")
+                new_m = T.sblock_alloc_buffer((1,), "float32")
+                d_val = T.sblock_alloc_buffer((1,), "float32")
+                S_val = T.sblock_alloc_buffer((1,), "float32")
+                scale_O = T.sblock_alloc_buffer((1,), "float32")
+                factor = T.sblock_alloc_buffer((1,), "float32")
+
+                cur_page_indptr_begin: T.let[T.int32] = page_table_indptr[b]
+                cur_page_indptr_end: T.let[T.int32] = page_table_indptr[b + 1]
+
+                kv_chunk_len[0] = T.if_then_else(
+                    cur_page_indptr_begin != cur_page_indptr_end,
+                    _get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, b, length_info, sliding_window),
+                    0,
+                )
+
+                for h_qo in T.serial(H_qo):
+                    m_val[0] = -5e4
+                    d_val[0] = 1.0
+
+                    for d in T.serial(D):
+                        O_local[d] = 0.0
+
+                    for d in T.serial(D):
+                        Q_local[d] = T.if_then_else(
+                            rotary_mode == 1,
+                            _rope(Q, q_rope_position[b], head_dim, rope_theta, rope_scale, (b, h_qo, d), qkv_dtype, rope_scaling),
+                            Q[b, h_qo, d],
+                        )
+
+                    for row_idx in T.serial(kv_chunk_len[0]):
+                        seq_offset: T.let[T.int32()] = _get_seq_offset(row_idx, b, length_info, sliding_window)
+                        page_no: T.let[T.int32()] = page_table_values[cur_page_indptr_begin + (seq_offset // page_size)]
+                        page_offset: T.let[T.int32()] = seq_offset % page_size
+
+                        for d in T.serial(D):
+                            K_local[d] = T.if_then_else(
+                                rotary_mode == 1,
+                                _rope(pages, k_rope_pos_offset[b] + row_idx, head_dim, rope_theta, rope_scale, (page_no, 0, h_qo // group_size, page_offset, d), qkv_dtype, rope_scaling),
+                                pages[page_no, 0, h_qo // group_size, page_offset, d],
+                            )
+                        S_val[0] = 0.0
+                        for d in T.serial(D):
+                            S_val[0] += Q_local[d] * K_local[d]
+                        S_val[0] *= sm_scale * math.log2(math.exp(1))
+
+                        new_m[0] = T.max(m_val[0], S_val[0])
+                        d_val[0] = (d_val[0] * T.exp2(m_val[0] - new_m[0])) + T.exp2(S_val[0] - new_m[0])
+
+                        scale_O[0] = T.exp2(m_val[0] - new_m[0])
+
+                        for d in T.serial(D):
+                            O_local[d] = O_local[d] * scale_O[0]
+
+                        m_val[0] = new_m[0]
+                        for d in T.serial(D):
+                            V_local[d] = pages[page_no, 1, h_qo // group_size, page_offset, d]
+
+                        factor[0] = T.exp2(S_val[0] - m_val[0])
+                        for d in T.serial(D):
+                            O_local[d] = O_local[d] + V_local[d] * factor[0]
+                    for d in T.serial(D):
+                        O_local[d] = O_local[d] / d_val[0]
+                        output[b, h_qo, d] = O_local[d]
+                    lse[b, h_qo] = m_val[0] + T.log2(d_val[0])
+
+    return batch_decode_paged_kv
+
+
+def _attention_decode(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_window: bool, rope_scaling: dict[str, Any], target: Target, page_size: int = 16):
+    qkv_dtype_bytes = 2
+    H_qo = num_qo_heads
+    H_kv = num_kv_heads
+    D = head_dim
+
+    THREAD_LIMIT = 512
+    TILE_SIZE_PER_BDX = 2
+    if target.kind.name == "opencl" and (("android" in str(target.host)) or ("adreno" in str(target.attrs))):
+        # Keeping lower thread limit for this kernel on adreno target
+        # to avoid register spill
+        THREAD_LIMIT = 256
+        TILE_SIZE_PER_BDX = 1
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    thread_limit = min(max_num_threads_per_block, THREAD_LIMIT)
+
+    GROUP_SIZE = H_qo // H_kv
+    VEC_SIZE = min(max(8 // qkv_dtype_bytes, D // 32), 4)
+    bdx = D // VEC_SIZE
+    bdy = GROUP_SIZE
+    while bdx * bdy > thread_limit and bdy > 1:
+        bdy //= 2
+    gdz = GROUP_SIZE // bdy
+    threads_per_CTA = max(thread_limit, bdx * bdy)
+    bdz = threads_per_CTA // (bdx * bdy)
+    tile_size_per_bdx = TILE_SIZE_PER_BDX if GROUP_SIZE == 1 else 1
+    check_thread_limits(target, bdx=bdx, bdy=bdy, bdz=bdz, gdz=1)
+
+    global_symbol = "batch_decode_paged_kv"
+    if sliding_window:
+        global_symbol += "_sliding_window"
+
+    # pylint: disable=too-many-branches
+    @T.prim_func(s_tir=True)
+    def batch_decode_paged_kv(
+        Q_handle: T.handle,
+        pages_handle: T.handle,
+        page_table_indptr_handle: T.handle,
+        page_table_values_handle: T.handle,
+        var_length_info: T.handle, # [b] when sliding window = False, or otherwise [3, b]
+        k_rope_pos_offset_handle: T.handle,
+        q_rope_position_handle: T.handle,
+        output_handle: T.handle,
+        lse_handle: T.handle,
+        rotary_mode: T.int32,
+        rope_scale: T.float32,
+        rope_theta: T.float32,
+        sm_scale: T.float32,
+    ):
+        T.func_attr({"tirx.is_scheduled": True, "global_symbol": global_symbol})
+        B = T.int32()
+        nnz_pages = T.int32()
+        max_num_pages = T.int32()
+        pages_elem_offset = T.int64()
+        page_indptr_elem_offset = T.int32()
+        page_values_elem_offset = T.int32()
+        k_rope_pos_offset_elem_offset = T.int32()
+        q_rope_position_elem_offset = T.int32()
+        length_info_elem_offset = T.int32()
+
+        Q = T.match_buffer(Q_handle, (B, H_qo, D), qkv_dtype)
+        pages = T.match_buffer(pages_handle, (max_num_pages, 2, H_kv, page_size, D), qkv_dtype, elem_offset=pages_elem_offset)
+        page_table_indptr = T.match_buffer(page_table_indptr_handle, (B + 1,), "int32", elem_offset=page_indptr_elem_offset)
+        page_table_values = T.match_buffer(page_table_values_handle, (nnz_pages,), "int32", elem_offset=page_values_elem_offset)
+        k_rope_pos_offset = T.match_buffer(k_rope_pos_offset_handle, (B,), "int32", elem_offset=k_rope_pos_offset_elem_offset)
+        q_rope_position = T.match_buffer(q_rope_position_handle, (B,), "int32", elem_offset=q_rope_position_elem_offset)
+        output = T.match_buffer(output_handle, (B, H_qo, D), qkv_dtype)
+        lse = T.match_buffer(lse_handle, (B, H_qo), "float32")  # pylint: disable=unused-variable
+        length_info = _declare_length_info(var_length_info, B, sliding_window, length_info_elem_offset)
+
+        for bx in T.thread_binding(B, thread="blockIdx.x"):
+            for fused_by_bz in T.thread_binding(H_kv * gdz, thread="blockIdx.y"):
+                for ty in T.thread_binding(bdy, thread="threadIdx.y"):
+                    for tx in T.thread_binding(bdx, thread="threadIdx.x"):
+                        for tz in T.thread_binding(bdz, thread="threadIdx.z"):
+                            with T.sblock("attn"):
+                                Q_local = T.sblock_alloc_buffer((VEC_SIZE,), qkv_dtype, scope="local")
+                                kv_chunk_len = T.sblock_alloc_buffer((1,), "int32", scope="local")
+                                K_smem = T.sblock_alloc_buffer((bdz * bdy * tile_size_per_bdx, D), qkv_dtype, scope="shared")
+                                V_smem = T.sblock_alloc_buffer((bdz * bdy * tile_size_per_bdx, D), qkv_dtype, scope="shared")
+                                O_allreduce = T.sblock_alloc_buffer((bdz, bdy, D), "float32", scope="shared")
+                                md_allreduce = T.sblock_alloc_buffer((bdz, bdy, 2), "float32", scope="shared")
+                                S_reduce_local = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                t0 = T.sblock_alloc_buffer((1,), "float32", scope="local")
+
+                                S_local = T.sblock_alloc_buffer((bdy * tile_size_per_bdx), "float32", scope="local")
+                                QK_local = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
+                                V_local = T.sblock_alloc_buffer((VEC_SIZE,), qkv_dtype, scope="local")
+                                m_prev = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                d_prev = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                other_m = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                other_d = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                exp_mprev = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                exp_otherm = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                other_o = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
+                                st_m = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                st_d = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                O_local = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
+
+                                by: T.let[T.int32] = fused_by_bz % H_kv
+                                bz: T.let[T.int32] = fused_by_bz // H_kv
+                                batch_idx: T.let[T.int32] = bx
+                                cur_page_indptr_begin: T.let[T.int32] = page_table_indptr[batch_idx]
+                                cur_page_indptr_end: T.let[T.int32] = page_table_indptr[batch_idx + 1]
+                                kv_chunk_len[0] = T.if_then_else(
+                                    cur_page_indptr_begin != cur_page_indptr_end,
+                                    _get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, batch_idx, length_info, sliding_window),
+                                    0
+                                )
+
+                                # init states
+                                st_m[0] = -5e4
+                                st_d[0] = 1.0
+                                for vec in T.vectorized(VEC_SIZE):
+                                    O_local[vec] = 0.0
+
+                                # load q
+                                for vec in T.vectorized(VEC_SIZE):
+                                    Q_local[vec] = T.if_then_else(
+                                        rotary_mode == 1,
+                                        _rope(Q, q_rope_position[batch_idx], head_dim, rope_theta, rope_scale, (bx, by * GROUP_SIZE + bz * bdy + ty, tx * VEC_SIZE + vec), qkv_dtype, rope_scaling),
+                                        Q[bx, by * GROUP_SIZE + bz * bdy + ty, tx * VEC_SIZE + vec]
+                                    )
+
+                                for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_size_per_bdx * bdy * bdz)):
+                                    tile_start_s: T.let[T.int32()] = (tz * bdy + ty) * tile_size_per_bdx  # type: ignore
+                                    tile_start_g: T.let[T.int32()] = ((iterator * bdz + tz) * bdy + ty) * tile_size_per_bdx  # type: ignore
+                                    # load KV from global memory to shared memory
+                                    for j in T.serial(tile_size_per_bdx):
+                                        with T.sblock("KV_load"):
+                                            T.reads()
+                                            T.writes()
+                                            row_g: T.let[T.int32()] = tile_start_g + j  # type: ignore
+                                            if row_g < kv_chunk_len[0]:
+                                                seq_offset: T.let[T.int32()] = _get_seq_offset(row_g, batch_idx, length_info, sliding_window)  # type: ignore
+                                                page_no: T.let[T.int32()] = page_table_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)]  # type: ignore
+                                                page_offset: T.let[T.int32()] = T.floormod(seq_offset, page_size)  # type: ignore
+                                                for vec in T.vectorized(VEC_SIZE):
+                                                    K_smem[tile_start_s + j, tx * VEC_SIZE + vec] = T.if_then_else(
+                                                        rotary_mode == 1,
+                                                        _rope(pages, k_rope_pos_offset[batch_idx] + row_g, head_dim, rope_theta, rope_scale, (page_no, 0, by, page_offset, tx * VEC_SIZE + vec), qkv_dtype, rope_scaling),
+                                                        pages[page_no, 0, by, page_offset, tx * VEC_SIZE + vec]
+                                                    )
+                                                    V_smem[tile_start_s + j, tx * VEC_SIZE + vec] = pages[page_no, 1, by, page_offset, tx * VEC_SIZE + vec]
+                                            else:
+                                                for vec in T.vectorized(VEC_SIZE):
+                                                    K_smem[tile_start_s + j, tx * VEC_SIZE + vec] = 0.0
+                                                    V_smem[tile_start_s + j, tx * VEC_SIZE + vec] = 0.0
+                                    T.tvm_storage_sync("shared")
+                                    # compute QK
+                                    m_prev[0] = st_m[0]
+                                    for j in T.serial(bdy * tile_size_per_bdx):
+                                        # compute S = Q * K * sm_scale
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            QK_local[vec] = T.cast(Q_local[vec], "float32") * T.cast(K_smem[tz * bdy * tile_size_per_bdx + j, tx * VEC_SIZE + vec], "float32") * sm_scale * math.log2(math.exp(1))
+                                        S_reduce_local[0] = 0
+                                        for vec in T.unroll(VEC_SIZE):
+                                            S_reduce_local[0] += QK_local[vec]
+
+                                        with T.sblock("block_cross_thread"):
+                                            T.reads(S_reduce_local[0])
+                                            T.writes(t0[0])
+                                            T.attr(
+                                                T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                                                "reduce_scope",
+                                                T.int32(0),
+                                            )
+                                            T.tvm_thread_allreduce(T.uint32(1), S_reduce_local[0], True, t0[0], tx, dtype="void")
+
+                                        S_local[j] = -5e4
+                                        if (iterator * bdz + tz) * bdy * tile_size_per_bdx + j < kv_chunk_len[0]:
+                                            S_local[j] = t0[0]
+                                        # update st_m
+                                        st_m[0] = T.max(st_m[0], S_local[j])
+
+                                    # update st_d, st_O
+                                    o_scale: T.let[T.float32] = T.exp2(m_prev[0] - st_m[0])
+                                    st_d[0] *= o_scale
+                                    for j in T.serial(bdy * tile_size_per_bdx):
+                                        S_local[j] = T.exp2(S_local[j] - st_m[0])
+                                        st_d[0] += S_local[j]
+                                    for j in T.vectorized(VEC_SIZE):
+                                        O_local[j] *= o_scale
+
+                                    # load V from shared memory to local memory
+                                    # compute O
+                                    for j in T.serial(bdy * tile_size_per_bdx):
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            V_local[vec] = V_smem[tz * bdy * tile_size_per_bdx + j, tx * VEC_SIZE + vec]
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            O_local[vec] += T.cast(V_local[vec], "float32") * S_local[j]
+
+                                if bdz > 1:
+                                    # allreduce over bdz
+                                    for vec in T.vectorized(VEC_SIZE):
+                                        O_allreduce[tz, ty, tx * VEC_SIZE + vec] = O_local[vec]
+                                    md_allreduce[tz, ty, 0] = st_m[0]
+                                    md_allreduce[tz, ty, 1] = st_d[0]
+                                    T.tvm_storage_sync("shared")
+
+                                    st_m[0] = -5e4
+                                    st_d[0] = 1.0
+                                    for vec in T.vectorized(VEC_SIZE):
+                                        O_local[vec] = 0.0
+
+                                    for j in T.serial(bdz):
+                                        m_prev[0] = st_m[0]
+                                        d_prev[0] = st_d[0]
+                                        other_m[0] = md_allreduce[j, ty, 0]
+                                        other_d[0] = md_allreduce[j, ty, 1]
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            other_o[vec] = O_allreduce[j, ty, tx * VEC_SIZE + vec]
+                                        st_m[0] = T.max(st_m[0], other_m[0])
+                                        st_d[0] = d_prev[0] * T.exp2(m_prev[0] - st_m[0]) + other_d[0] * T.exp2(other_m[0] - st_m[0])
+                                        exp_mprev[0] = T.exp2(m_prev[0] - st_m[0])
+                                        exp_otherm[0] = T.exp2(other_m[0] - st_m[0])
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            O_local[vec] = O_local[vec] * exp_mprev[0] + other_o[vec] * exp_otherm[0]
+
+                                # normalize O
+                                for vec in T.vectorized(VEC_SIZE):
+                                    O_local[vec] /= st_d[0]
+
+                                # store O to global memory
+                                for vec in T.vectorized(VEC_SIZE):
+                                    output[batch_idx, by * GROUP_SIZE + bz * bdy + ty, tx * VEC_SIZE + vec] = O_local[vec]
+
+                                # store lse to global memory
+                                lse[batch_idx, by * GROUP_SIZE + bz * bdy + ty] = st_m[0] + T.log2(st_d[0])
+    # pylint: enable=too-many-branches
+    return batch_decode_paged_kv
+
+
+def _merge_state_inplace_cpu(v_dtype):
+    @T.prim_func(s_tir=True)
+    def merge_state_inplace_cpu(
+        v: T.handle,
+        s: T.handle,
+        v_other: T.handle,
+        s_other: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": True})
+        N = T.int32()
+        H = T.int32()
+        D = T.int32()
+
+        V = T.match_buffer(v, (N, H, D), v_dtype)
+        S = T.match_buffer(s, (N, H), "float32")
+        V_other = T.match_buffer(v_other, (N, H, D), v_dtype)
+        S_other = T.match_buffer(s_other, (N, H), "float32")
+
+        for n in T.serial(N):
+            for h in T.serial(H):
+                with T.sblock("merge"):
+                    s_val = _var_cpu("float32")
+                    s_other_val = _var_cpu("float32")
+                    s_max = _var_cpu("float32")
+                    scale = _var_cpu("float32")
+                    other_scale = _var_cpu("float32")
+
+                    s_val[0] = S[n, h]
+                    s_other_val[0] = S_other[n, h]
+                    s_max[0] = T.max(s_val[0], s_other_val[0])
+                    s_val[0] = T.exp2(s_val[0] - s_max[0])
+                    s_other_val[0] = T.exp2(s_other_val[0] - s_max[0])
+                    scale[0] = s_val[0] / (s_val[0] + s_other_val[0])
+                    other_scale[0] = s_other_val[0] / (s_val[0] + s_other_val[0])
+                    for d in T.serial(D):
+                        V[n, h, d] = V[n, h, d] * scale[0] + V_other[n, h, d] * other_scale[0]
+                    S[n, h] = T.log2(s_val[0] + s_other_val[0]) + s_max[0]
+
+    return merge_state_inplace_cpu
+
+
+def _merge_state_inplace(num_heads, head_dim, v_dtype, target: Target, global_symbol: str | None = None):
+    v_dtype_bytes = 2
+    VEC_SIZE = min(max(8 // v_dtype_bytes, head_dim // 32), 4)
+    bdx = head_dim // VEC_SIZE
+    bdy = num_heads
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    while bdx * bdy > max_num_threads_per_block and bdy > 1:
+        bdy //= 2
+    gdy = num_heads // bdy
+    check_thread_limits(target, bdx=bdx, bdy=bdy, bdz=1, gdz=1)
+
+    @T.prim_func(s_tir=True)
+    def merge_state_inplace(
+        v: T.handle,
+        s: T.handle,
+        v_other: T.handle,
+        s_other: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": True})
+        N = T.int32()
+        H = T.int32()
+        D = T.int32()
+
+        V = T.match_buffer(v, (N, H, D), v_dtype)
+        S = T.match_buffer(s, (N, H), "float32")
+        V_other = T.match_buffer(v_other, (N, H, D), v_dtype)
+        S_other = T.match_buffer(s_other, (N, H), "float32")
+
+        for bx in T.thread_binding(N, thread="blockIdx.x"):
+            for by in T.thread_binding(gdy, thread="blockIdx.y"):
+                for ty in T.thread_binding(bdy, thread="threadIdx.y"):
+                    for tx in T.thread_binding(bdx, thread="threadIdx.x"):
+                        with T.sblock("merge"):
+                            s_val = _var("float32")
+                            s_other_val = _var("float32")
+                            s_max = _var("float32")
+                            scale = _var("float32")
+                            other_scale = _var("float32")
+
+                            v_vec = T.sblock_alloc_buffer((VEC_SIZE,), v_dtype, scope="local")
+                            v_other_vec = T.sblock_alloc_buffer((VEC_SIZE,), v_dtype, scope="local")
+
+                            s_val[0] = S[bx, ty + by * bdy]
+                            s_other_val[0] = S_other[bx, ty + by * bdy]
+                            s_max[0] = T.max(s_val[0], s_other_val[0])
+                            s_val[0] = T.exp2(s_val[0] - s_max[0])
+                            s_other_val[0] = T.exp2(s_other_val[0] - s_max[0])
+                            scale[0] = s_val[0] / (s_val[0] + s_other_val[0])
+                            other_scale[0] = s_other_val[0] / (s_val[0] + s_other_val[0])
+
+                            # load v
+                            for vec in T.vectorized(VEC_SIZE):
+                                v_vec[vec] = V[bx, ty + by * bdy, tx * VEC_SIZE + vec]
+                            # load v_other
+                            for vec in T.vectorized(VEC_SIZE):
+                                v_other_vec[vec] = V_other[bx, ty + by * bdy, tx * VEC_SIZE + vec]
+
+                            # merge
+                            for vec in T.serial(VEC_SIZE):
+                                v_vec[vec] = v_vec[vec] * scale[0] + v_other_vec[vec] * other_scale[0]
+
+                            # store v
+                            for vec in T.vectorized(VEC_SIZE):
+                                V[bx, ty + by * bdy, tx * VEC_SIZE + vec] = v_vec[vec]
+
+                            # store s
+                            S[bx, ty + by * bdy] = T.log2(s_val[0] + s_other_val[0]) + s_max[0]
+
+    func = merge_state_inplace
+    if global_symbol:
+        func = func.with_attr("global_symbol", global_symbol)
+    return func

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

@@ -0,0 +1,569 @@
+# 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, E731, RUF005
+# fmt: off
+
+"""Shared TIR helpers used by KV-cache / attention kernels in this package.
+
+This module consolidates constructs reused by the prefill/decode/paged/tree
+attention kernels so each kernel file can focus on its own specialised logic.
+
+Contents:
+- Thread-limit checks (``get_max_num_threads_per_block``, ``check_thread_limits``)
+- KV-cache enums (``AttnKind``, ``RopeMode``)
+- Small TVMScript helpers (``_var``, ``_var_cpu``, ``_causal_mask``, ``_rope``)
+- Length-info accessors for sliding-window-aware indexing
+- Buffer allocators for the tiled online-softmax state used by every prefill kernel
+- ``_make_prefill_macros`` — the ``@T.macro`` bundle invoked by the prefill kernels
+- Tiling config (``_get_prefill_kernel_config``) and scheduling (``_schedule_prefill_kernel``)
+"""
+
+# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
+import enum
+import math
+from typing import Any
+
+import tvm
+from tvm import s_tir, tirx
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.target import Target
+
+from .position_embedding import switch_rope_freq_func
+
+
+def _var(dtype):
+    return T.sblock_alloc_buffer((1,), dtype, scope="local")
+
+
+def _var_cpu(dtype):
+    return T.sblock_alloc_buffer((1,), dtype)
+
+
+def get_max_num_threads_per_block(target: Target) -> int:
+    """
+    max(max_num_threads, max_threads_per_block); if latter does not exist, return max_num_threads.
+    We add this method since some targets have both fields and `max_threads_per_block` is larger.
+    """
+    max_num_threads = int(target.attrs["max_num_threads"])
+    max_threads_per_block = target.attrs.get("max_threads_per_block", None)
+    if max_threads_per_block is None:
+        return max_num_threads
+    return max(max_num_threads, max_threads_per_block)
+
+
+def check_thread_limits(target: Target, bdx: int, bdy: int, bdz: int, gdz: int):
+    """
+    Check whether max num threads exceeded given a target.
+
+    Parameters
+    ----------
+    bdx: threadIdx.x
+    bdy: threadIdx.y
+    bdz: threadIdx.z
+    gdz: blockIdx.z
+    """
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+
+    assert bdx * bdy * bdz <= max_num_threads_per_block, (
+        f"{target.kind} max num threads exceeded: {bdx}*{bdy}*{bdz}>{max_num_threads_per_block}"
+    )
+
+    if target.kind.name == "webgpu":
+        # https://gpuweb.github.io/gpuweb/#dom-supported-limits-maxcomputeworkgroupsizez
+        assert bdz <= 64, f"webgpu's threadIdx.z cannot exceed 64, but got bdz={bdz}"
+        assert gdz == 1, f"webgpu's blockIdx.z should be 1, but got gdz={gdz}"
+
+
+class AttnKind(enum.IntEnum):
+    """The attention kind class.
+    MHA denotes multi-head attention, multi-query attention or grouped query attention.
+    MLA denotes multi-head latent attention.
+    """
+
+    MHA = 0
+    MLA = 1
+    MHA_SLIDING = 3
+
+
+class RopeMode(enum.IntEnum):
+    """The RoPE mode of the Paged KV cache.
+    If it is none, the KV cache will not apply RoPE to q and k.
+    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.
+    """
+
+    NONE = 0
+    NORMAL = 1
+    INLINE = 2
+
+
+def _rope(buffer: T.Buffer, offset: tirx.Var, rotary_dim: int, theta: tirx.Var, scale: tirx.Var, indices: tuple[tirx.Var, ...], qkv_dtype: str, rope_scaling: dict[str, Any]):
+    d = indices[-1]
+    cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(offset * scale, d, rotary_dim, theta, "float32")
+    cos = cos_freq * buffer[indices].astype("float32")
+    sin = sin_freq * tirx.if_then_else(
+        d < rotary_dim // 2,
+        -buffer[indices[:-1] + (d + rotary_dim // 2,)],
+        buffer[indices[:-1] + (d - rotary_dim // 2,)],
+    ).astype("float32")
+    expr = (cos + sin).astype(qkv_dtype)
+    for var, value in var_map.items():
+        expr = tirx.Let(var, value, expr)
+    return expr
+
+
+def _causal_mask(causal, row, col, kv_len, qo_len):
+    return T.if_then_else(
+        causal > 0,
+        col < kv_len - qo_len + row + 1,
+        col < kv_len,
+    )
+
+
+def _declare_length_info(var_length_info, batch_size, sliding_window, elem_offset):
+    return (
+        T.match_buffer(var_length_info, (3, batch_size), "int32", elem_offset=elem_offset)
+        if sliding_window
+        else T.match_buffer(var_length_info, (batch_size,), "int32", elem_offset=elem_offset)
+    )
+
+
+def _get_kv_chunk_len(num_pages, page_size, seq_id, length_info, sliding_window):
+    if not sliding_window:
+        return (num_pages - 1) * page_size + length_info[seq_id]
+    # ((num_pages - 1) * page_size + last_page_len) - sliding_window_offset + sink_size
+    return (num_pages - 1) * page_size + length_info[0, seq_id] - length_info[1, seq_id] + length_info[2, seq_id]
+
+
+def _get_seq_offset(pos, seq_id, length_info, sliding_window):
+    if not sliding_window:
+        return pos
+    # pos if pos < sink_size else pos - sink_size + sliding_window_offset
+    return T.if_then_else(
+        pos < length_info[2, seq_id],
+        pos,
+        pos - length_info[2, seq_id] + length_info[1, seq_id],
+    )
+
+
+def _alloc_softmax_state_buffers(tile_x, tile_z, bdx, num_warps):
+    """Allocate the shared/local online-softmax working state used by every tiled prefill kernel.
+
+    Returns ``(S_smem, S_local, m_smem, m_prev_smem, d_smem, m_new, m_prev, d_new)``.
+    """
+    S_smem = T.sblock_alloc_buffer((tile_x, tile_z), "float32", scope="shared")
+    S_local = T.sblock_alloc_buffer((tile_x, tile_z), "float32", scope="local")
+    m_smem = T.sblock_alloc_buffer((tile_x,), "float32", scope="shared")
+    m_prev_smem = T.sblock_alloc_buffer((tile_x,), "float32", scope="shared")
+    d_smem = T.sblock_alloc_buffer((tile_x,), "float32", scope="shared")
+    md_shape = (math.ceil(tile_x / (bdx * num_warps)),)
+    m_new = T.sblock_alloc_buffer(md_shape, "float32", scope="local")
+    m_prev = T.sblock_alloc_buffer(md_shape, "float32", scope="local")
+    d_new = T.sblock_alloc_buffer(md_shape, "float32", scope="local")
+    return S_smem, S_local, m_smem, m_prev_smem, d_smem, m_new, m_prev, d_new
+
+
+def _alloc_mha_qkvo_buffers(tile_x, tile_z, d_qk, d_v, dtype):
+    """Allocate Q/K/V shared + O local buffers for standard MHA/GQA prefill kernels."""
+    Q_smem = T.sblock_alloc_buffer((tile_x, d_qk), dtype, scope="shared")
+    K_smem = T.sblock_alloc_buffer((tile_z, d_qk), dtype, scope="shared")
+    V_smem = T.sblock_alloc_buffer((tile_z, d_v), dtype, scope="shared")
+    O_local = T.sblock_alloc_buffer((tile_x, d_v), "float32", scope="local")
+    return Q_smem, K_smem, V_smem, O_local
+
+
+def _alloc_mla_qkvo_buffers(tile_x, tile_z, d_qk, d_latent, dtype):
+    """Allocate Q + combined KV shared + O local for MLA prefill (V reuses the KV buffer)."""
+    Q_smem = T.sblock_alloc_buffer((tile_x, d_qk), dtype, scope="shared")
+    KV_smem = T.sblock_alloc_buffer((tile_z, d_qk), dtype, scope="shared")
+    O_local = T.sblock_alloc_buffer((tile_x, d_latent), "float32", scope="local")
+    return Q_smem, KV_smem, O_local
+
+
+def _alloc_tile_walk_state():
+    """Return (tile_id, batch_idx, batch_tiles, batch_rows, iterator, kv_chunk_len) int32 scalars for the paged/ragged/MLA tile-walk state machine."""
+    return _var("int32"), _var("int32"), _var("int32"), _var("int32"), _var("int32"), _var("int32")
+
+
+def _make_prefill_macros(tile_x, tile_y, tile_z, tile_o, bdx, num_warps, group_size):
+    """Build @T.macro helpers shared across tiled online-softmax prefill kernels.
+
+    Parameters
+    ----------
+    tile_x : int  # query/output row tile
+    tile_y : int  # QK reduction dim (head_dim for MHA, d_qk for MLA/ragged)
+    tile_z : int  # key/value column tile
+    tile_o : int  # output/V column dim (d for MHA/sequence, d_v for ragged, d_latent for MLA)
+    """
+    @T.macro
+    def init_states(
+        m_smem: T.Buffer, d_smem: T.Buffer, O_local: T.Buffer, ty: T.int32, tx: T.int32,
+    ):
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                m_smem[row] = -5e4
+                d_smem[row] = 1.0
+        for li, lj in T.grid(tile_x, tile_o):
+            with T.sblock("O_init"):
+                i, j = T.axis.remap("SS", [li, lj])
+                O_local[i, j] = 0.0
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def compute_s_gemm(
+        Q_smem: T.Buffer, K_smem: T.Buffer, S_local: T.Buffer, S_smem: T.Buffer, sm_scale: T.float32,
+    ):
+        with T.sblock():
+            for li, lj, lk in T.grid(tile_x, tile_z, tile_y):
+                with T.sblock("S_gemm"):
+                    i, j, k = T.axis.remap("SSR", [li, lj, lk])
+                    with T.init():
+                        S_local[i, j] = 0.0
+                    S_local[i, j] += T.cast(Q_smem[i, k], "float32") * T.cast(K_smem[j, k], "float32") * sm_scale * math.log2(math.exp(1))
+        T.tvm_storage_sync("shared")
+        for li, lj in T.grid(tile_x, tile_z):
+            with T.sblock("S_store"):
+                i, j = T.axis.remap("SS", [li, lj])
+                S_smem[i, j] = S_local[i, j]
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def softmax_update_causal(
+        S_smem: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer, m_prev_smem: T.Buffer,
+        m_new: T.Buffer, m_prev: T.Buffer, d_new: T.Buffer,
+        ty: T.int32, tx: T.int32, LH_start: T.int32, L_kv_start: T.int32,
+        causal: T.int32, kv_len: T.int32, qo_len: T.int32,
+    ):
+        # Phase 1: compute m_new = max(masked S over kv tile), d_new = d_prev * exp2(m_prev - m_new)
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update1"):
+                    m_prev[i] = m_smem[row]
+                    m_new[i] = m_smem[row]
+                    row_: T.let[T.int32] = (LH_start + row) // group_size
+                    for j in T.serial(tile_z):
+                        if _causal_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len):
+                            m_new[i] = T.max(m_new[i], S_smem[row, j])
+                    d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
+        # Phase 2: exp-and-scale S_smem; masked-out entries use -inf
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            with T.sblock("update"):
+                for j in T.serial(tile_z):
+                    # predicate sits inside loop so sync stays outside conditional branches
+                    if row < tile_x:
+                        row_: T.let[T.int32] = (LH_start + row) // group_size
+                        if _causal_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len):
+                            S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
+                        else:
+                            S_smem[row, j] = T.exp2(-5e4 - m_new[i])
+        # Phase 3: d_new += sum(S_smem[row, :]); write m/d/m_prev back to smem
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update"):
+                    for j in T.serial(tile_z):
+                        d_new[i] += S_smem[row, j]
+                    m_smem[row] = m_new[i]
+                    d_smem[row] = d_new[i]
+                    m_prev_smem[row] = m_prev[i]
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def compute_o_gemm(
+        S_smem: T.Buffer, V_smem: T.Buffer, O_local: T.Buffer,
+        m_prev_smem: T.Buffer, m_smem: T.Buffer,
+    ):
+        with T.sblock():
+            for li, lj, lk in T.grid(tile_x, tile_o, tile_z):
+                with T.sblock("O_gemm"):
+                    i, j, k = T.axis.remap("SSR", [li, lj, lk])
+                    with T.init():
+                        O_local[i, j] *= T.exp2(m_prev_smem[i] - m_smem[i])
+                    O_local[i, j] += S_smem[i, k] * T.cast(V_smem[k, j], "float32")
+
+    @T.macro
+    def paged_store_output_lse(
+        output: T.Buffer, lse: T.Buffer, O_local: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer,
+        q_indptr: T.Buffer, b_idx: T.int32, by: T.int32, LH_start: T.int32,
+    ):
+        """Paged-style (q_indptr-based) O_store + lse_store epilogue.
+
+        Used by paged prefill, ragged prefill and MLA prefill. MLA passes ``by=0`` so
+        the ``by * group_size`` term drops to zero at compile time.
+        """
+        for li, lj in T.grid(tile_x, tile_o):
+            with T.sblock("O_store"):
+                i, j = T.axis.remap("SS", [li, lj])
+                cur_L: T.let[T.int32] = q_indptr[b_idx] + (LH_start + i) // group_size
+                cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
+                if cur_L < q_indptr[b_idx + 1]:
+                    output[cur_L, cur_H_qo, j] = O_local[i, j] / d_smem[i]
+        for li in T.grid(tile_x):
+            with T.sblock("lse_store"):
+                i = T.axis.remap("S", [li])
+                cur_L: T.let[T.int32] = q_indptr[b_idx] + (LH_start + i) // group_size
+                cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
+                if cur_L < q_indptr[b_idx + 1]:
+                    lse[cur_L, cur_H_qo] = m_smem[i] + T.log2(d_smem[i])
+
+    @T.macro
+    def advance_tile_batch(
+        tile_id: T.Buffer, batch_idx: T.Buffer, batch_tiles: T.Buffer, batch_rows: T.Buffer,
+        q_indptr: T.Buffer, batch_size: T.int32,
+    ):
+        """Advance tile_id/batch_idx past exhausted batches.
+
+        After the loop, either batch_idx[0] >= batch_size (all tiles consumed) or
+        tile_id[0] < batch_tiles[0] (the current batch still has work to do).
+        """
+        while tile_id[0] >= batch_tiles[0] and batch_idx[0] < batch_size:
+            tile_id[0] -= batch_tiles[0]
+            batch_idx[0] += 1
+            if batch_idx[0] < batch_size:
+                b_idx: T.let[T.int32] = batch_idx[0]
+                batch_rows[0] = (q_indptr[b_idx + 1] - q_indptr[b_idx]) * group_size
+                batch_tiles[0] = T.ceildiv(batch_rows[0], tile_x)
+
+    @T.macro
+    def softmax_update_valid_length(
+        S_smem: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer, m_prev_smem: T.Buffer,
+        m_new: T.Buffer, m_prev: T.Buffer, d_new: T.Buffer,
+        ty: T.int32, tx: T.int32, LH_start: T.int32, L_kv_start: T.int32,
+        valid_len: T.int32, qo_len: T.int32, kv_len: T.int32,
+    ):
+        # Same three-phase online softmax as softmax_update_causal but with a
+        # per-batch right-padding mask in place of causal masking.
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update1"):
+                    m_prev[i] = m_smem[row]
+                    m_new[i] = m_smem[row]
+                    row_: T.let[T.int32] = (LH_start + row) // group_size
+                    for j in T.serial(tile_z):
+                        if tirx.And(tirx.And(row_ < qo_len, row_ < valid_len), L_kv_start + j < valid_len):
+                            m_new[i] = T.max(m_new[i], S_smem[row, j])
+                    d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            with T.sblock("update"):
+                for j in T.serial(tile_z):
+                    if row < tile_x:
+                        row_: T.let[T.int32] = (LH_start + row) // group_size
+                        if tirx.And(tirx.And(row_ < qo_len, row_ < valid_len), L_kv_start + j < valid_len):
+                            S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
+                        else:
+                            S_smem[row, j] = T.exp2(-5e4 - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update"):
+                    for j in T.serial(tile_z):
+                        d_new[i] += S_smem[row, j]
+                    m_smem[row] = m_new[i]
+                    d_smem[row] = d_new[i]
+                    m_prev_smem[row] = m_prev[i]
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def softmax_update_causal_padded_left(
+        S_smem: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer, m_prev_smem: T.Buffer,
+        m_new: T.Buffer, m_prev: T.Buffer, d_new: T.Buffer,
+        ty: T.int32, tx: T.int32, LH_start: T.int32, L_kv_start: T.int32,
+        valid_len: T.int32, qo_len: T.int32, kv_len: T.int32,
+    ):
+        # Three-phase online softmax with left-padding + causal mask. Real
+        # queries occupy [qo_len - valid_len, qo_len); real keys occupy
+        # [kv_len - valid_len, kv_len). Causal keeps
+        # col <= row + (kv_len - qo_len) within those valid suffixes.
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update1"):
+                    m_prev[i] = m_smem[row]
+                    m_new[i] = m_smem[row]
+                    row_: T.let[T.int32] = (LH_start + row) // group_size
+                    pad_q: T.let[T.int32] = qo_len - valid_len
+                    pad_kv: T.let[T.int32] = kv_len - valid_len
+                    for j in T.serial(tile_z):
+                        col_: T.let[T.int32] = L_kv_start + j
+                        if tirx.And(tirx.And(row_ < qo_len, row_ >= pad_q), tirx.And(col_ >= pad_kv, col_ < kv_len - qo_len + row_ + 1)):
+                            m_new[i] = T.max(m_new[i], S_smem[row, j])
+                    d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            with T.sblock("update"):
+                for j in T.serial(tile_z):
+                    if row < tile_x:
+                        row_: T.let[T.int32] = (LH_start + row) // group_size
+                        pad_q: T.let[T.int32] = qo_len - valid_len
+                        pad_kv: T.let[T.int32] = kv_len - valid_len
+                        col_: T.let[T.int32] = L_kv_start + j
+                        if tirx.And(tirx.And(row_ < qo_len, row_ >= pad_q), tirx.And(col_ >= pad_kv, col_ < kv_len - qo_len + row_ + 1)):
+                            S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
+                        else:
+                            S_smem[row, j] = T.exp2(-5e4 - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update"):
+                    for j in T.serial(tile_z):
+                        d_new[i] += S_smem[row, j]
+                    m_smem[row] = m_new[i]
+                    d_smem[row] = d_new[i]
+                    m_prev_smem[row] = m_prev[i]
+        T.tvm_storage_sync("shared")
+
+    return init_states, compute_s_gemm, softmax_update_causal, compute_o_gemm, softmax_update_valid_length, advance_tile_batch, paged_store_output_lse, softmax_update_causal_padded_left
+
+
+def _get_prefill_kernel_config(h_kv, h_q, d, dtype, target: Target):
+    NUM_BLKS = 16
+    LOAD_VEC = 8 // ((DataType(dtype).bits + 7) // 8)  # 8 bytes
+    group_size = h_q // h_kv
+
+    bdx = 32
+    num_warps = 4
+    tile_x, tile_y, tile_z = (
+        64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
+        d,
+        64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
+    )
+    original_tile_y = tile_y
+    original_tile_z = tile_z
+    while (tile_x * tile_z) % (bdx * num_warps) != 0:
+        tile_z += original_tile_z
+    while (tile_x * tile_y) % (bdx * num_warps) != 0:
+        tile_y += original_tile_y
+
+    # Otherwise we would exceed maxComputeWorkgroupStorageSize
+    if (
+        target.kind.name == "webgpu"
+        and ((d + 127) // 128) * ((DataType(dtype).bits + 15) // 16) >= 4
+    ):
+        tile_z = 8
+        num_warps = 2
+    if target.kind.name == "opencl" and (
+        ("android" in str(target.host)) or ("adreno" in str(target.attrs))
+    ):
+        LOAD_VEC = 16 // ((DataType(dtype).bits + 7) // 8)  # 16 bytes
+        NUM_BLKS = group_size * 8
+
+    check_thread_limits(target, bdx=bdx, bdy=num_warps, bdz=1, gdz=1)
+
+    return NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z
+
+
+def _schedule_prefill_kernel(sch: s_tir.Schedule, load_vec, bdx, num_warps, tile_x, tile_y, tile_z, transform_k_load: bool, merged_qk_load: bool) -> tvm.s_tir.Schedule:
+    get_extent = lambda *lps: [int(sch.get(lp).extent) for lp in lps]
+
+    def get_vecsize(extent):
+        return min(load_vec, (extent & ~(extent - 1)))
+
+    def getxy_vecsize(x, y, t):
+        assert (x * y) % t == 0
+        return min(get_vecsize(y), get_vecsize(x * y // t))
+
+    def get_tile_size(x, y, t):
+        cnt = (x * y) // t
+        assert (x * y) % t == 0
+        tile_y = math.ceil(math.sqrt(cnt))
+        while (cnt % tile_y != 0 or y % tile_y != 0 or x % (cnt // tile_y) != 0) and tile_y <= cnt:
+            tile_y += 1
+        assert tile_y <= cnt
+        tile_x = cnt // tile_y
+        return tile_x, tile_y
+
+    def apply_to_qkv_load(sch: s_tir.Schedule, block):
+        loop_x, loop_y = sch.get_loops(block)[-2:]
+        x_extent, y_extent = get_extent(loop_x, loop_y)
+        vec_size = getxy_vecsize(x_extent, y_extent, bdx * num_warps)
+        yo, yv = sch.split(loop_y, [None, vec_size])
+        yo_extent = y_extent // vec_size
+        tile_x, tile_y = get_tile_size(x_extent, yo_extent, (bdx * num_warps))
+        xo, xi = sch.split(loop_x, [tile_x, None])
+        yo, yi = sch.split(yo, [tile_y, None])
+        sch.reorder(xi, yi, xo, yo)
+        t = sch.fuse(xi, yi)
+        ty, tx = sch.split(t, [num_warps, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+        sch.vectorize(yv)
+
+    def apply_to_so_ewise(sch: s_tir.Schedule, block, tile):
+        loop_x, loop_y = sch.get_loops(block)[-2:]
+        xo, xi = sch.split(loop_x, factors=[None, tile[0]])
+        yo, yi = sch.split(loop_y, factors=[None, tile[1]])
+        sch.reorder(xo, yo, xi, yi)
+        yiv_extent = get_vecsize(tile[1])
+        yio, yiv = sch.split(yi, [None, yiv_extent])
+        sch.unroll(yio)
+        sch.vectorize(yiv)
+        t = sch.fuse(xo, yo)
+        ty, tx = sch.split(t, factors=[None, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+    def apply_to_gemm(sch: s_tir.Schedule, block, tile, r_len=16, k_major=False):
+        loop_x, loop_y, loop_z = sch.get_loops(block)[-3:]
+        xo, xi = sch.split(loop_x, factors=[None, tile[0]])
+        yo, yi = sch.split(loop_y, factors=[None, tile[1]])
+        sch.reorder(xo, yo, xi, yi)
+        t = sch.fuse(xo, yo)
+        ty, tx = sch.split(t, factors=[None, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+        ko, ki = sch.split(loop_z, factors=[None, r_len])
+        if k_major:
+            sch.reorder(ko, xi, yi, ki)
+        else:
+            sch.reorder(ko, ki, xi, yi)
+        yiv_extent = get_vecsize(tile[1])
+        yio, yiv = sch.split(yi, [None, yiv_extent])
+        sch.unroll(yio)
+        sch.vectorize(yiv)
+        sch.unroll(xi)
+        sch.decompose_reduction(block, ty)
+
+    def apply_to_md(sch, block):
+        loop = sch.get_loops(block)[-1]
+        _, ty, tx = sch.split(loop, factors=[None, num_warps, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+    if transform_k_load and not merged_qk_load:
+        sch.transform_layout("K_load", ("write", 0), lambda i, j: (j, i))
+    tile_s = get_tile_size(tile_x, tile_z, bdx * num_warps)
+    tile_o = get_tile_size(tile_x, tile_y, bdx * num_warps)
+    apply_to_gemm(sch, sch.get_sblock("S_gemm"), tile_s, k_major=True)
+    apply_to_gemm(sch, sch.get_sblock("O_gemm"), tile_o, k_major=False)
+    apply_to_so_ewise(sch, sch.get_sblock("S_store"), tile_s)
+    apply_to_so_ewise(sch, sch.get_sblock("O_init"), tile_o)
+    apply_to_so_ewise(sch, sch.get_sblock("O_store"), tile_o)
+    apply_to_qkv_load(sch, sch.get_sblock("Q_load"))
+    if not merged_qk_load:
+        apply_to_qkv_load(sch, sch.get_sblock("K_load"))
+        apply_to_qkv_load(sch, sch.get_sblock("V_load"))
+    else:
+        apply_to_qkv_load(sch, sch.get_sblock("KV_load"))
+    apply_to_md(sch, sch.get_sblock("lse_store"))
+    return sch

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

@@ -0,0 +1,293 @@
+# 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
+# fmt: off
+
+"""TIR kernels that operate on paged KV-cache storage (without doing attention).
+
+This module contains:
+- Append helpers that transpose/write new K/V tokens into the paged layout
+  (``_kv_cache_transpose_append`` and its MLA variant).
+- Debug helpers that extract K/V from the paged layout for inspection
+  (``_kv_cache_debug_get_kv``, ``_kv_cache_debug_get_kv_mla``).
+- Copy helpers used by the cache runtime for forking/sharing pages
+  (``_copy_single_page``, ``_copy_single_page_mla``, ``_copy_single_page_cpu``).
+- Compact helpers that reorganise pages after removals
+  (``_compact_kv_copy``, ``_compact_kv_copy_cpu``).
+"""
+
+# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
+from tvm.script import tirx as T
+from tvm.target import Target
+
+from ._kernel_common import get_max_num_threads_per_block
+
+
+def _kv_cache_transpose_append(num_key_value_heads, head_dim, dtype, page_size: int = 16):
+    """Return the TIR function that appends new k/v data to PagedKVCache."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_transpose_append(
+        var_pages: T.handle,
+        var_k_data: T.handle,
+        var_v_data: T.handle,
+        var_position_map: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        ntoken = T.Var("num_tokens_excluding_cache", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_key_value_heads, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+        k_data = T.match_buffer(var_k_data, (ntoken, num_key_value_heads, head_dim), dtype)
+        v_data = T.match_buffer(var_v_data, (ntoken, num_key_value_heads, head_dim), dtype)
+        position_map = T.match_buffer(var_position_map, (ntoken,), "int32", elem_offset=position_map_elem_offset)
+        for global_pos, h, f in T.grid(ntoken, num_key_value_heads, head_dim):
+            if position_map[global_pos] != T.int32(-1):
+                with T.sblock("k_transpose_append"):
+                    vgpos, vh, vf = T.axis.remap("SSS", [global_pos, h, f])
+                    T.reads(position_map[vgpos], k_data[vgpos, vh, vf])
+                    T.writes(pages[position_map[vgpos] // page_size, 0, vh, position_map[vgpos] % page_size, vf])
+                    position: T.int32 = position_map[vgpos]  # type: ignore
+                    pages[T.floordiv(position, page_size), 0, vh, T.floormod(position, page_size), vf] = k_data[vgpos, vh, vf]
+                with T.sblock("v_transpose_append"):
+                    vgpos, vh, vf = T.axis.remap("SSS", [global_pos, h, f])
+                    T.reads(position_map[vgpos], v_data[vgpos, vh, vf])
+                    T.writes(pages[position_map[vgpos] // page_size, 1, vh, position_map[vgpos] % page_size, vf])
+                    position: T.int32 = position_map[vgpos] # type: ignore[name-defined,no-redef]
+                    pages[T.floordiv(position, page_size), 1, vh, T.floormod(position, page_size), vf] = v_data[vgpos, vh, vf]
+
+    return tir_kv_cache_transpose_append
+
+
+def _kv_cache_transpose_append_mla(d_qk: int, dtype, page_size: int = 16):
+    """Return the TIR function that appends new compressed KV data to PagedKVCache for MLA."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_transpose_append_mla(
+        var_pages: T.handle,
+        var_kv_data: T.handle,
+        var_position_map: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        ntoken = T.Var("num_tokens_excluding_cache", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, page_size, d_qk), dtype, elem_offset=pages_elem_offset)
+        kv_data = T.match_buffer(var_kv_data, (ntoken, d_qk), dtype)
+        position_map = T.match_buffer(var_position_map, (ntoken,), "int32", elem_offset=position_map_elem_offset)
+        for global_pos, f in T.grid(ntoken, d_qk):
+            if position_map[global_pos] != T.int32(-1):
+                with T.sblock("k_transpose_append"):
+                    vgpos, vf = T.axis.remap("SS", [global_pos, f])
+                    T.reads(position_map[vgpos], kv_data[vgpos, vf])
+                    T.writes(pages[position_map[vgpos] // page_size, position_map[vgpos] % page_size, vf])
+                    position: T.int32 = position_map[vgpos]  # type: ignore
+                    pages[T.floordiv(position, page_size), T.floormod(position, page_size), vf] = kv_data[vgpos, vf]
+
+    return tir_kv_cache_transpose_append_mla
+
+
+def _kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, head_dim, dtype):
+    """Return the TIR function that fetches the k/v data on given positions and layer."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_debug_get_kv(
+        var_pages: T.handle,
+        var_position_map: T.handle,
+        var_k_data: T.handle,
+        var_v_data: T.handle,
+        layer_id: T.int64,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        seqlen = T.Var("num_tokens_including_cache", "int64")
+        page_size = T.Var("page_size", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_key_value_heads, page_size, head_dim), dtype,elem_offset=pages_elem_offset)
+        position_map = T.match_buffer(var_position_map, (seqlen,), "int32", elem_offset=position_map_elem_offset)
+        k_data = T.match_buffer(var_k_data, (num_hidden_layers, seqlen, num_key_value_heads, head_dim), dtype)
+        v_data = T.match_buffer(var_v_data, (num_hidden_layers, seqlen, num_key_value_heads, head_dim), dtype)
+        for p, h, d in T.grid(seqlen, num_key_value_heads, head_dim):
+            with T.sblock("copy0"):
+                vp, vh, vd = T.axis.remap("SSS", [p, h, d])
+                T.reads(position_map[vp], pages[position_map[vp] // page_size, 0:2, vh, position_map[vp] % page_size, vd])
+                T.writes(k_data[layer_id, vp, vh, vd], v_data[layer_id, vp, vh, vd])
+                position: T.int32 = position_map[vp] # type: ignore[name-defined]
+                k_data[layer_id, vp, vh, vd] = pages[T.floordiv(position, page_size), 0, vh, T.floormod(position, page_size), vd]
+                v_data[layer_id, vp, vh, vd] = pages[T.floordiv(position, page_size), 1, vh, T.floormod(position, page_size), vd]
+
+    return tir_kv_cache_debug_get_kv
+
+
+def _kv_cache_debug_get_kv_mla(num_hidden_layers, d_qk, dtype):
+    """Return the TIR function that fetches the k/v data on given positions and layer."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_debug_get_kv_mla(
+        var_pages: T.handle,
+        var_position_map: T.handle,
+        var_compressed_kv_with_k_pe_data: T.handle,
+        layer_id: T.int64,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        seqlen = T.Var("num_tokens_including_cache", "int64")
+        page_size = T.Var("page_size", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, page_size, d_qk), dtype, elem_offset=pages_elem_offset)
+        position_map = T.match_buffer(var_position_map, (seqlen,), "int32", elem_offset=position_map_elem_offset)
+        compressed_kv_with_k_pe_data = T.match_buffer(var_compressed_kv_with_k_pe_data, (num_hidden_layers, seqlen, d_qk), dtype)
+        for p, d in T.grid(seqlen, d_qk):
+            with T.sblock("copy0"):
+                vp, vd = T.axis.remap("SS", [p, d])
+                T.reads(position_map[vp], pages[position_map[vp] // page_size, position_map[vp] % page_size, vd])
+                T.writes(compressed_kv_with_k_pe_data[layer_id, vp, vd])
+                position: T.int32 = position_map[vp] # type: ignore[name-defined]
+                compressed_kv_with_k_pe_data[layer_id, vp, vd] = pages[T.floordiv(position, page_size), T.floormod(position, page_size), vd]
+
+    return tir_kv_cache_debug_get_kv_mla
+
+
+def _copy_single_page(num_heads, page_size, head_dim, dtype, target: Target):
+    tx = get_max_num_threads_per_block(target)
+
+    @T.prim_func(s_tir=True)
+    def copy_single_page(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        pages_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+
+        for b in T.thread_binding((copy_length * num_heads * head_dim + tx - 1) // tx, thread="blockIdx.x"):
+            for t in T.thread_binding(tx, thread="threadIdx.x"):
+                with T.sblock("copy"):
+                    T.where(b * tx + t < copy_length * num_heads * head_dim)
+                    vh = T.axis.spatial(num_heads, T.Cast("int32", (b * tx + t) // (copy_length * head_dim)))
+                    vp = T.axis.spatial(copy_length, (b * tx + t) % (copy_length * head_dim) // head_dim)
+                    vd = T.axis.spatial(head_dim, T.Cast("int32", (b * tx + t) % head_dim))
+                    pages[tgt_page_id, 0, vh, vp, vd] = pages[src_page_id, 0, vh, vp, vd]
+                    pages[tgt_page_id, 1, vh, vp, vd] = pages[src_page_id, 1, vh, vp, vd]
+
+    return copy_single_page
+
+
+def _copy_single_page_mla(page_size, head_dim, dtype, target: Target):
+    tx = get_max_num_threads_per_block(target)
+
+    @T.prim_func(s_tir=True)
+    def copy_single_page_mla(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        pages_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+
+        for b in T.thread_binding((copy_length * head_dim + tx - 1) // tx, thread="blockIdx.x"):
+            for t in T.thread_binding(tx, thread="threadIdx.x"):
+                with T.sblock("copy"):
+                    T.where(b * tx + t < copy_length * head_dim)
+                    vp = T.axis.spatial(copy_length, (b * tx + t) // head_dim)
+                    vd = T.axis.spatial(head_dim, T.Cast("int32", (b * tx + t) % head_dim))
+                    pages[tgt_page_id, vp, vd] = pages[src_page_id, vp, vd]
+
+    return copy_single_page_mla
+
+
+def _copy_single_page_cpu(num_heads, page_size, head_dim, dtype):
+    tx = 1
+
+    @T.prim_func(s_tir=True)
+    def copy_single_page_cpu(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype)
+
+        for b in T.serial((copy_length * num_heads * head_dim + tx - 1) // tx):
+            for t in T.serial(tx):
+                with T.sblock("copy"):
+                    T.where(b * tx + t < copy_length * num_heads * head_dim)
+                    vh = T.axis.spatial(num_heads, T.Cast("int32", (b * tx + t) // (copy_length * head_dim)))
+                    vp = T.axis.spatial(copy_length, (b * tx + t) % (copy_length * head_dim) // head_dim)
+                    vd = T.axis.spatial(head_dim, T.Cast("int32", (b * tx + t) % head_dim))
+                    pages[tgt_page_id, 0, vh, vp, vd] = pages[src_page_id, 0, vh, vp, vd]
+                    pages[tgt_page_id, 1, vh, vp, vd] = pages[src_page_id, 1, vh, vp, vd]
+
+    return copy_single_page_cpu
+
+
+def _compact_kv_copy(num_heads, head_dim, dtype, target: Target, page_size: int = 16):
+    tx = get_max_num_threads_per_block(target)
+
+    @T.prim_func(s_tir=True)
+    def compact_kv_copy(var_pages: T.handle, var_copy_length_indptr: T.handle, var_copy_src_dst_pos: T.handle, batch_size: T.int32):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        total_copy_length = T.int32()
+        copy_length_indptr_elem_offset = T.int32()
+        copy_src_dst_pos_elem_offset = T.int32()
+        pages_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+        copy_length_indptr = T.match_buffer(var_copy_length_indptr, (batch_size + 1,), "int32", elem_offset=copy_length_indptr_elem_offset)
+        copy_src_dst_pos = T.match_buffer(var_copy_src_dst_pos, (2, total_copy_length), "int32", elem_offset=copy_src_dst_pos_elem_offset)
+
+        with T.sblock("root"):
+            for bhd_o in T.thread_binding((batch_size * num_heads * head_dim + tx - 1) // tx, thread="blockIdx.x"):
+                for bhd_i in T.thread_binding(tx, thread="threadIdx.x"):
+                    b: T.int32 = (bhd_o * tx + bhd_i) // (num_heads * head_dim)
+                    h: T.int32 = (bhd_o * tx + bhd_i) // head_dim % num_heads
+                    d: T.int32 = (bhd_o * tx + bhd_i) % head_dim
+                    if (bhd_o * tx + bhd_i) < batch_size * num_heads * head_dim:
+                        for i in T.serial(copy_length_indptr[b + 1] - copy_length_indptr[b]):
+                            src_pos: T.int32 = copy_src_dst_pos[0, copy_length_indptr[b] + i]
+                            dst_pos: T.int32 = copy_src_dst_pos[1, copy_length_indptr[b] + i]
+                            pages[dst_pos // page_size, 0, h, dst_pos % page_size, d] = pages[src_pos // page_size, 0, h, src_pos % page_size, d]
+                            pages[dst_pos // page_size, 1, h, dst_pos % page_size, d] = pages[src_pos // page_size, 1, h, src_pos % page_size, d]
+
+    return compact_kv_copy
+
+
+def _compact_kv_copy_cpu(num_heads, head_dim, dtype, page_size: int = 16):
+    tx = 8
+
+    @T.prim_func(s_tir=True)
+    def compact_kv_copy_cpu(var_pages: T.handle, var_copy_length_indptr: T.handle, var_copy_src_dst_pos: T.handle, batch_size: T.int32):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        total_copy_length = T.int32()
+        copy_length_indptr_elem_offset = T.int32()
+        copy_src_dst_pos_elem_offset = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype)
+        copy_length_indptr = T.match_buffer(var_copy_length_indptr, (batch_size + 1,), "int32", elem_offset=copy_length_indptr_elem_offset)
+        copy_src_dst_pos = T.match_buffer(var_copy_src_dst_pos, (2, total_copy_length), "int32", elem_offset=copy_src_dst_pos_elem_offset)
+
+        with T.sblock("root"):
+            for bhd_o in T.serial((batch_size * num_heads * head_dim + tx - 1) // tx):
+                for bhd_i in T.serial(tx):
+                    b: T.int32 = (bhd_o * tx + bhd_i) // (num_heads * head_dim)
+                    h: T.int32 = (bhd_o * tx + bhd_i) // head_dim % num_heads
+                    d: T.int32 = (bhd_o * tx + bhd_i) % head_dim
+                    if (bhd_o * tx + bhd_i) < batch_size * num_heads * head_dim:
+                        for i in T.serial(copy_length_indptr[b + 1] - copy_length_indptr[b]):
+                            src_pos: T.int32 = copy_src_dst_pos[0, copy_length_indptr[b] + i]
+                            dst_pos: T.int32 = copy_src_dst_pos[1, copy_length_indptr[b] + i]
+                            pages[dst_pos // page_size, 0, h, dst_pos % page_size, d] = pages[src_pos // page_size, 0, h, src_pos % page_size, d]
+                            pages[dst_pos // page_size, 1, h, dst_pos % page_size, d] = pages[src_pos // page_size, 1, h, src_pos % page_size, d]
+
+    return compact_kv_copy_cpu

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

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

@@ -0,0 +1,894 @@
+# 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.
+
+"""Operators for positional embeddings, e.g. RoPE."""
+
+import math
+from collections.abc import Callable
+from functools import partial
+from typing import Any
+
+from tvm import tirx
+from tvm.relax.frontend.nn import Tensor, op
+from tvm.script import tirx as T
+
+# pylint: disable=invalid-name
+
+
+def rope_freq_default(s: tirx.Var, d: tirx.Var, d_range: int, theta: float, dtype: str):
+    """Compute the inverse frequency of RoPE and then return the cosine and sine of it.
+
+    Parameters
+    ----------
+    s : tirx.Var
+        The position index.
+
+    d : tirx.Var
+        The dimension index.
+
+    d_range : int
+        The maximum dimension index.
+
+    theta : float
+        The theta value in RoPE, which controls the frequency.
+
+    dtype : str
+        The data type of the output.
+
+    Returns
+    -------
+    cos_freq : Tensor
+        The cosine of the inverse frequency.
+
+    sin_freq : Tensor
+        The sine of the inverse frequency.
+
+    var_map: Dict[tirx.Var, tirx.Expr]
+        The common expression map.
+    """
+    freq = s / tirx.power(theta, d * 2 % d_range / tirx.const(d_range, "float32"))
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = tirx.cos(freq_var).astype(dtype)
+    sin_freq = tirx.sin(freq_var).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def rope_freq_gptj(s: tirx.Var, d: tirx.Var, d_range: int, theta: float, dtype: str):
+    """Compute the inverse frequency of RoPE for gptj RoPE scaling."""
+    freq = s / tirx.power(theta, 2 * (d // 2) % d_range / tirx.const(d_range, "float32"))
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = tirx.cos(freq_var).astype(dtype)
+    sin_freq = tirx.sin(freq_var).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def rope_freq_llama4(  # pylint: disable=too-many-arguments,too-many-locals
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float,
+    dtype: str,
+    factor: float,
+    low_freq_factor: float,
+    high_freq_factor: float,
+    original_max_position_embeddings: float,
+):
+    """Compute the inverse frequency of RoPE for llama4 RoPE scaling."""
+    orig_freq = tirx.const(1, "float32") / tirx.power(
+        theta, 2 * (d // 2) / tirx.const(d_range, "float32")
+    )
+    orig_freq_var = tirx.Var("orig_freq", "float32")
+
+    llama4_inv_scaling_factor = 1.0 / factor
+
+    if high_freq_factor == low_freq_factor:
+        wavelength = tirx.const(2 * math.pi, "float32") / orig_freq_var
+        threshold_wavelen = tirx.const(
+            original_max_position_embeddings / low_freq_factor, "float32"
+        )
+
+        scaled_freq = tirx.if_then_else(
+            wavelength > threshold_wavelen, orig_freq_var / factor, orig_freq_var
+        )
+        smoothed_freq = s * scaled_freq
+
+    else:
+        # Original smooth interpolation logic
+        inv_diff_freq_factor = 1.0 / (high_freq_factor - low_freq_factor)
+
+        llama4_alpha = original_max_position_embeddings / (2 * math.pi) * inv_diff_freq_factor
+        llama4_beta = low_freq_factor * inv_diff_freq_factor
+        smooth = tirx.max(0.0, tirx.min(1.0, llama4_alpha * orig_freq_var - llama4_beta))
+        smoothed_freq = s * (
+            (1.0 - smooth) * orig_freq_var * llama4_inv_scaling_factor + smooth * orig_freq_var
+        )
+
+    smoothed_freq_var = tirx.Var("smoothed_freq", "float32")
+    cos_freq = tirx.cos(smoothed_freq_var).astype(dtype)
+    sin_freq = tirx.sin(smoothed_freq_var).astype(dtype)
+    return (
+        cos_freq,
+        sin_freq,
+        {smoothed_freq_var: smoothed_freq, orig_freq_var: orig_freq},
+    )
+
+
+def rope_freq_llama3(  # pylint: disable=too-many-arguments,too-many-locals
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float,
+    dtype: str,
+    factor: float,
+    low_freq_factor: float,
+    high_freq_factor: float,
+    original_max_position_embeddings: float,
+):
+    """Compute the inverse frequency of RoPE for llama3 RoPE scaling."""
+    orig_freq = tirx.const(1, "float32") / tirx.power(
+        theta, d * 2 % d_range / tirx.const(d_range, "float32")
+    )
+    orig_freq_var = tirx.Var("orig_freq", "float32")
+    inv_diff_freq_factor = 1.0 / (high_freq_factor - low_freq_factor)
+    llama3_inv_scaling_factor = 1.0 / factor
+    llama3_alpha = original_max_position_embeddings / (2 * math.pi) * inv_diff_freq_factor
+    llama3_beta = low_freq_factor * inv_diff_freq_factor
+    smooth = tirx.max(0.0, tirx.min(1.0, llama3_alpha * orig_freq_var - llama3_beta))
+    smoothed_freq = s * (
+        (1.0 - smooth) * orig_freq_var * llama3_inv_scaling_factor + smooth * orig_freq_var
+    )
+    smoothed_freq_var = tirx.Var("smoothed_freq", "float32")
+    cos_freq = tirx.cos(smoothed_freq_var).astype(dtype)
+    sin_freq = tirx.sin(smoothed_freq_var).astype(dtype)
+    return (
+        cos_freq,
+        sin_freq,
+        {smoothed_freq_var: smoothed_freq, orig_freq_var: orig_freq},
+    )
+
+
+def rope_freq_longrope(  # pylint: disable=too-many-arguments
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float,
+    dtype: str,
+    max_position_embeddings: int,
+    original_max_position_embeddings: int,
+    ext_factors: T.Buffer | None = None,
+):
+    """Compute the inverse frequency of RoPE for longrope scaling."""
+    scale = max_position_embeddings / original_max_position_embeddings
+    scaling_factor = (
+        math.sqrt(1 + math.log(scale) / math.log(original_max_position_embeddings))
+        if scale > 1.0
+        else 1.0
+    )
+    divisor = tirx.power(theta, d * 2 % d_range / tirx.const(d_range, "float32"))
+    if ext_factors is not None:
+        divisor = ext_factors[d % (d_range // 2)] * divisor
+    freq = s / divisor
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = (tirx.cos(freq_var) * scaling_factor).astype(dtype)
+    sin_freq = (tirx.sin(freq_var) * scaling_factor).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def yarn_find_correction_dim(
+    num_rotations: int,
+    d: tirx.Var,
+    max_position_embeddings: int,
+    inv_theta_log_scale: float | tirx.Expr | None = None,
+):
+    """Inverse dim formula to find dim based on number of rotations"""
+    return (
+        d * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) * inv_theta_log_scale
+    )
+
+
+def yarn_find_correction_range(
+    low_rot: int,
+    high_rot: int,
+    d: tirx.Var,
+    max_position_embeddings: int,
+    inv_theta_log_scale: float | tirx.Expr | None = None,
+):
+    """Find the correction range based on the number of rotations"""
+    low = yarn_find_correction_dim(
+        low_rot, d, max_position_embeddings, inv_theta_log_scale=inv_theta_log_scale
+    )
+    high = yarn_find_correction_dim(
+        high_rot, d, max_position_embeddings, inv_theta_log_scale=inv_theta_log_scale
+    )
+    return tirx.max(low, 0), tirx.min(high, d - 1)
+
+
+def rope_freq_yarn(
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float | tirx.Expr,
+    dtype: str,
+    original_max_position_embeddings: int,
+    scaling_factor: float,
+    beta_fast: int,
+    beta_slow: int,
+    inv_theta_log_scale: float | tirx.Expr | None = None,
+):  # pylint: disable=too-many-arguments, too-many-locals
+    """Compute the inverse frequency of RoPE for yarn RoPE scaling."""
+
+    exponent = d * 2 % d_range / tirx.const(d_range, "float32")
+    freq_power = tirx.power(theta, exponent)
+    freq_extra = tirx.const(1, "float32") / freq_power
+    freq_inter = tirx.const(1, "float32") / (scaling_factor * freq_power)
+
+    low, high = yarn_find_correction_range(
+        beta_fast,
+        beta_slow,
+        d_range,
+        original_max_position_embeddings,
+        inv_theta_log_scale=inv_theta_log_scale,
+    )
+    high = tirx.if_then_else(low == high, high + 0.001, high)
+    inv_freq_mask = tirx.const(1, "float32") - tirx.max(
+        tirx.min((d - low) / (high - low), 1.0), 0.0
+    ).astype("float32")
+    inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask
+    freq = s * inv_freq
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = tirx.cos(freq_var).astype(dtype)
+    sin_freq = tirx.sin(freq_var).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def switch_rope_freq_func(rope_scaling: dict[str, Any]) -> Callable:
+    """Return the RoPE inverse frequency computation function based
+    on the given RoPE scaling.
+    """
+    if "rope_type" not in rope_scaling:
+        return rope_freq_default
+    if rope_scaling["rope_type"] == "gptj":
+        return rope_freq_gptj
+    if rope_scaling["rope_type"] == "llama3":
+        return partial(
+            rope_freq_llama3,
+            factor=rope_scaling["factor"],
+            low_freq_factor=rope_scaling["low_freq_factor"],
+            high_freq_factor=rope_scaling["high_freq_factor"],
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+        )
+    if rope_scaling["rope_type"] == "llama4":
+        return partial(
+            rope_freq_llama4,
+            factor=rope_scaling["factor"],
+            low_freq_factor=rope_scaling["low_freq_factor"],
+            high_freq_factor=rope_scaling["high_freq_factor"],
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+        )
+    if rope_scaling["rope_type"] == "longrope":
+        return partial(
+            rope_freq_longrope,
+            max_position_embeddings=rope_scaling["max_position_embeddings"],
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+        )
+    if rope_scaling["rope_type"] == "yarn":
+        inv_theta_log_scale = rope_scaling.get("inv_theta_log_scale")
+        assert inv_theta_log_scale is not None, "inv_theta_log_scale must be precomputed for YaRN"
+        return partial(
+            rope_freq_yarn,
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+            scaling_factor=rope_scaling["factor"],
+            beta_fast=rope_scaling["beta_fast"],
+            beta_slow=rope_scaling["beta_slow"],
+            inv_theta_log_scale=inv_theta_log_scale,
+        )
+    raise ValueError(f"Unsupported RoPE scaling type: {rope_scaling['rope_type']}")
+
+
+# mypy: disable-error-code="attr-defined"
+
+
+def llama_rope(  # pylint: disable=too-many-arguments
+    qkv: Tensor,
+    total_seq_len: tirx.Var,
+    theta: float,
+    scale: float,
+    num_q_heads: int,
+    num_kv_heads: int,
+    rope_scaling: dict[str, Any],
+    rotary_dim: int | None = None,
+) -> tuple[Tensor, Tensor, Tensor]:
+    """Llama-style RoPE. Given a fused QKV tensor, it returns three tensors, Q, K, and V, where Q
+    and K are rotated by RoPE while V remains unchanged.
+
+    Parameters
+    ----------
+    qkv : Tensor
+        The fused QKV tensor of shape: [batch_size, seq_len, #q_heads + #kv_heads * 2, head_dim]
+
+    total_seq_len : tirx.Var
+        The total sequence length after being concatenated with KVCache. It is used to compute the
+        offset of RoPE.
+
+    theta : float
+        The theta value, or "base" in RoPE, which controls the frequency.
+
+    scale : float
+        The RoPE scaling factor.
+
+    num_q_heads : int
+        The number of query heads.
+
+    num_kv_heads : int
+        The number of key/value heads. It differs from `num_q_heads` in group-query attention.
+
+    rope_scaling : Dict
+        The configuration of RoPE scaling.
+
+    rotary_dim : Optional[int]
+        The number of dimensions in the embedding that RoPE is applied to. By default, the
+        rotary_dim is the same as head_dim.
+
+    Returns
+    -------
+    q : Tensor
+        The query tensor of shape [batch_size, seq_len, #q_heads, head_dim] w/ RoPE applied
+
+    k : Tensor
+        The key tensor of shape [batch_size, seq_len, #kv_heads, head_dim] w/ RoPE applied
+
+    v : Tensor
+        The value tensor of shape [batch_size, seq_len, #kv_heads, head_dim] w/o RoPE applied
+    """
+    _, _, fused_heads, head_dim = qkv.shape
+    assert fused_heads == num_q_heads + num_kv_heads * 2
+    if rotary_dim is None:
+        rotary_dim = head_dim
+    dtype = qkv.dtype
+    scale = tirx.const(scale, dtype)
+
+    def _rope(  # pylint: disable=too-many-arguments
+        x: T.Buffer,
+        b: tirx.Var,
+        s: tirx.Var,
+        h: tirx.Var,
+        d: tirx.Var,
+        offset: tirx.Var,
+    ):
+        cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+            (s + offset) * scale, d, rotary_dim, theta, dtype
+        )
+        cos = cos_freq * x[b, s, h, d]
+        if rope_scaling["rope_type"] == "gptj":
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[b, s, h, d + 1],
+                x[b, s, h, d - 1],
+            )
+        else:
+            sin = sin_freq * tirx.if_then_else(
+                d < rotary_dim // 2,
+                -x[b, s, h, d + rotary_dim // 2],
+                x[b, s, h, d - rotary_dim // 2],
+            )
+        expr = cos + sin
+        for var, value in var_map.items():
+            expr = tirx.Let(var, value, expr)
+        return expr
+
+    @T.prim_func(private=True, s_tir=True)
+    def fused_rope(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        total_seq_len: T.int64,
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        batch_size = T.int64()
+        seq_len = T.int64()
+        qkv = T.match_buffer(var_qkv, (batch_size, seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (batch_size, seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (batch_size, seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (batch_size, seq_len, num_kv_heads, head_dim), dtype)
+        for iters in T.grid(batch_size, seq_len, fused_heads, head_dim):
+            with T.sblock("llama_fused_rope"):
+                b, s, h, d = T.axis.remap("SSSS", iters)
+                if h < num_q_heads:
+                    q[b, s, h, d] = T.if_then_else(
+                        d < rotary_dim,
+                        _rope(qkv, b, s, h, d, total_seq_len - seq_len),
+                        qkv[b, s, h, d],
+                    )
+                elif h < num_q_heads + num_kv_heads:
+                    k[b, s, h - num_q_heads, d] = T.if_then_else(
+                        d < rotary_dim,
+                        _rope(qkv, b, s, h, d, total_seq_len - seq_len),
+                        qkv[b, s, h, d],
+                    )
+                else:
+                    v[b, s, h - (num_q_heads + num_kv_heads), d] = qkv[b, s, h, d]
+
+    b, s, _, _ = qkv.shape
+    return op.tensor_ir_op(  # pylint: disable=no-member
+        fused_rope,
+        "llama_rope",
+        args=[qkv, total_seq_len],
+        out=(
+            Tensor.placeholder((b, s, num_q_heads, head_dim), dtype),
+            Tensor.placeholder((b, s, num_kv_heads, head_dim), dtype),
+            Tensor.placeholder((b, s, num_kv_heads, head_dim), dtype),
+        ),
+    )
+
+
+def llama_rope_with_position_map(  # pylint: disable=too-many-arguments
+    theta: float,
+    scale: float,
+    head_dim: int,
+    num_q_heads: int,
+    num_kv_heads: int,
+    dtype: str,
+    rope_scaling: dict[str, Any],
+    rotary_dim: int | None = None,
+):
+    """Return the TIR function that computes Llama-style RoPE with q position map.
+
+    Parameters
+    ----------
+    theta : float
+        The theta value, or "base" in RoPE, which controls the frequency.
+
+    scale : float
+        The RoPE scaling factor.
+
+    head_dim : int
+        The number of features on each head.
+
+    num_q_heads : int
+        The number of query heads.
+
+    num_kv_heads : int
+        The number of key/value heads. It differs from `num_q_heads` in group-query attention.
+
+    dtype : str
+        The dtype of qkv data.
+
+    rope_scaling : Dict
+        The configuration of RoPE scaling.
+
+    rotary_dim : int
+        The number of dimensions in the embedding that RoPE is applied to. By default, the
+        rotary_dim is the same as head_dim.
+    """
+    fused_heads = num_q_heads + num_kv_heads * 2
+    if rotary_dim is None:
+        rotary_dim = head_dim
+    scale = tirx.const(scale, "float32")
+    is_longrope_scaling = rope_scaling.get("rope_type") == "longrope"
+    if is_longrope_scaling and "original_max_position_embeddings" in rope_scaling:
+        original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+    else:
+        original_max_position_embeddings = 0
+
+    def _rope(  # pylint: disable=too-many-arguments
+        x: T.Buffer,
+        s: tirx.Var,
+        h: tirx.Var,
+        d: tirx.Var,
+        pos: tirx.Var,
+        ext_factors: T.Buffer | None = None,
+    ):
+        kwargs = {}
+        if ext_factors:
+            kwargs["ext_factors"] = ext_factors
+        cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+            pos * scale, d, rotary_dim, theta, "float32", **kwargs
+        )
+        cos = cos_freq * x[s, h, d].astype("float32")
+        if "rope_type" in rope_scaling and rope_scaling["rope_type"] == "gptj":
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[s, h, d + 1],
+                x[s, h, d - 1],
+            ).astype("float32")
+        else:
+            sin = sin_freq * tirx.if_then_else(
+                d < rotary_dim // 2,
+                -x[s, h, d + rotary_dim // 2],
+                x[s, h, d - rotary_dim // 2],
+            ).astype("float32")
+        expr = (cos + sin).astype(dtype)
+        for var, value in var_map.items():
+            expr = tirx.Let(var, value, expr)
+        return expr
+
+    @T.prim_func(s_tir=True)
+    def fused_rope(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        apply_rope: T.int64,
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int32()
+        position_map_elem_offset = T.int32()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        for iters in T.grid(seq_len, fused_heads, head_dim):
+            with T.sblock("llama_fused_rope"):
+                s, h, d = T.axis.remap("SSS", iters)
+                if h < num_q_heads:
+                    q[s, h, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                elif h < num_q_heads + num_kv_heads:
+                    k[s, h - num_q_heads, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                else:
+                    v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    @T.prim_func(s_tir=True)
+    def fused_rope_longrope_scaling(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        ext_factors: T.Buffer((rotary_dim,), "float32"),  # type: ignore
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int64()
+        position_map_elem_offset = T.int64()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        # long factors is the first half, short factors is the second half
+        long_factors = T.decl_buffer((rotary_dim // 2,), "float32", data=ext_factors.data)
+        short_factors = T.decl_buffer(
+            (rotary_dim // 2,),
+            "float32",
+            data=ext_factors.data,
+            elem_offset=(rotary_dim // 2),
+        )
+
+        if seq_len > original_max_position_embeddings:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+        else:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    if is_longrope_scaling:
+        return fused_rope_longrope_scaling
+    return fused_rope
+
+
+def llama4_rope_with_position_map(  # pylint: disable=too-many-arguments
+    theta: float,
+    scale: float,
+    head_dim: int,
+    num_q_heads: int,
+    num_kv_heads: int,
+    dtype: str,
+    rope_scaling: dict[str, Any],
+    rotary_dim: int | None = None,
+):
+    """Return the TIR function that computes Llama-style RoPE with q position map.
+
+    Parameters
+    ----------
+    theta : float
+        The theta value, or "base" in RoPE, which controls the frequency.
+
+    scale : float
+        The RoPE scaling factor.
+
+    head_dim : int
+        The number of features on each head.
+
+    num_q_heads : int
+        The number of query heads.
+
+    num_kv_heads : int
+        The number of key/value heads. It differs from `num_q_heads` in group-query attention.
+
+    dtype : str
+        The dtype of qkv data.
+
+    rope_scaling : Dict
+        The configuration of RoPE scaling.
+
+    rotary_dim : int
+        The number of dimensions in the embedding that RoPE is applied to. By default, the
+        rotary_dim is the same as head_dim.
+    """
+    fused_heads = num_q_heads + num_kv_heads * 2
+    if rotary_dim is None:
+        rotary_dim = head_dim
+    scale = tirx.const(scale, "float32")
+    is_longrope_scaling = rope_scaling.get("rope_type") == "longrope"
+    if is_longrope_scaling and "original_max_position_embeddings" in rope_scaling:
+        original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+    else:
+        original_max_position_embeddings = 0
+
+    def _rope(  # pylint: disable=too-many-arguments
+        x: T.Buffer,
+        s: tirx.Var,
+        h: tirx.Var,
+        d: tirx.Var,
+        pos: tirx.Var,
+        ext_factors: T.Buffer | None = None,
+    ):
+        kwargs = {}
+        if ext_factors:
+            kwargs["ext_factors"] = ext_factors
+        cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+            pos * scale, d, rotary_dim, theta, "float32", **kwargs
+        )
+        cos = cos_freq * x[s, h, d].astype("float32")
+        if "rope_type" in rope_scaling and rope_scaling["rope_type"] == "gptj":
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[s, h, d + 1],
+                x[s, h, d - 1],
+            ).astype("float32")
+        else:
+            # Data layout is different for llama4 vs llama3
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[s, h, d + 1],
+                x[s, h, d - 1],
+            ).astype("float32")
+        expr = (cos + sin).astype(dtype)
+        for var, value in var_map.items():
+            expr = tirx.Let(var, value, expr)
+        return expr
+
+    @T.prim_func(private=True, s_tir=True)
+    def fused_rope(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        apply_rope: T.int64,
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int32()
+        position_map_elem_offset = T.int32()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        for iters in T.grid(seq_len, fused_heads, head_dim):
+            with T.sblock("llama_fused_rope"):
+                s, h, d = T.axis.remap("SSS", iters)
+                if h < num_q_heads:
+                    q[s, h, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                elif h < num_q_heads + num_kv_heads:
+                    k[s, h - num_q_heads, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                else:
+                    v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    @T.prim_func(s_tir=True)
+    def fused_rope_longrope_scaling(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        ext_factors: T.Buffer((rotary_dim,), "float32"),  # type: ignore
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int64()
+        position_map_elem_offset = T.int64()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        # long factors is the first half, short factors is the second half
+        long_factors = T.decl_buffer((rotary_dim // 2,), "float32", data=ext_factors.data)
+        short_factors = T.decl_buffer(
+            (rotary_dim // 2,),
+            "float32",
+            data=ext_factors.data,
+            elem_offset=(rotary_dim // 2),
+        )
+
+        if seq_len > original_max_position_embeddings:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+        else:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    if is_longrope_scaling:
+        return fused_rope_longrope_scaling
+    return fused_rope