# 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.
# pylint: disable=invalid-name
# ruff: noqa: E501, F841

"""Operators for tree attention."""

import math
from typing import Any

from tvm import s_tir, tirx
from tvm.runtime import DataType
from tvm.script import tirx as T
from tvm.target import Target

# The helpers below are shared with the main KV-cache kernels. They live in
# ``_kernel_common`` so that ``kv_cache.py``, this file, and the split kernel
# modules can all pull from a single source of truth.
from ._kernel_common import (
    _alloc_mha_qkvo_buffers,
    _alloc_softmax_state_buffers,
    _alloc_tile_walk_state,
    _declare_length_info,
    _get_kv_chunk_len,
    _get_seq_offset,
    _rope,
    check_thread_limits,
)

# mypy: disable-error-code="attr-defined,valid-type,no-redef"
# pylint: disable=too-many-statements,too-many-locals,too-many-arguments


def _check_tree_order(tree_order_indptr, tree_order, batch, row, col, kv_len, qo_len):
    tree_order_len = tree_order_indptr[batch + 1] - tree_order_indptr[batch]

    tree_start = kv_len - tree_order_len
    child_idx_in_tree = row + tree_order_len - qo_len
    parent_idx_in_tree = col - tree_start
    return tirx.all(
        col < kv_len,
        tirx.any(
            col < tree_start,
            tirx.all(
                tree_order[tree_order_indptr[batch] + child_idx_in_tree, 0]
                >= tree_order[tree_order_indptr[batch] + parent_idx_in_tree, 0],
                tree_order[tree_order_indptr[batch] + child_idx_in_tree, 0]
                < tree_order[tree_order_indptr[batch] + parent_idx_in_tree, 1],
            ),
        ),
    )


def tree_attn_cpu(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any]):
    """Generate tree attention kernel for batched tree attention.

    Parameters
    ----------
    h_kv : int
        Number of heads for key and value.
    h_q : int
        Number of heads for query.
    d : int
        Hidden dimension.
    dtype : str
        Data type.
    target : Target
        The target device.

    Returns
    -------
    mod : tvm.IRModule
        The generated IR module.
    """
    group_size = h_q // h_kv

    # fmt: off
    @T.prim_func(s_tir=True)
    def batch_tree_attn(  # pylint: disable=too-many-branches,line-too-long
        var_q: T.handle,  # [total_len, h_q, d]
        var_q_indptr: T.handle,  # [batch_size + 1]
        var_k: T.handle,  # [total_len, h_kv, d]
        var_v: T.handle,  # [total_len, h_kv, d]
        var_kv_indptr: T.handle,  # [batch_size + 1], kv_indptr should be the same as q_indptr in this case
        var_q_rope_position: T.handle,  # [total_q_len]
        var_mn_indptr: T.handle,  # [batch_size + 1]
        var_mask: T.handle,  # [mn_indptr[batch_size]]
        var_output: T.handle,  # [total_len, h_q, d]
        var_lse: T.handle,  # [total_len, h_q]
        rotary_mode: T.int32,
        rope_scale: T.float32,
        rope_theta: T.float32,
        sm_scale: T.float32,
    ):
        qo_len = T.int32()
        kv_len = T.int32()
        q_indptr_elem_offset = T.int32()
        kv_indptr_elem_offset = T.int32()
        q_rope_position_elem_offset = T.int32()
        mn_indptr_elem_offset = T.int32()
        mask_elem_offset = T.int32()
        tree_size = T.int32()
        batch_size_plus_1 = T.int32()

        q = T.match_buffer(var_q, (qo_len, h_q, d), dtype)
        q_indptr = T.match_buffer(
            var_q_indptr, (batch_size_plus_1,), "int32", elem_offset=q_indptr_elem_offset
        )
        k = T.match_buffer(var_k, (kv_len, h_kv, d), dtype)
        v = T.match_buffer(var_v, (kv_len, h_kv, d), dtype)
        kv_indptr = T.match_buffer(
            var_kv_indptr, (batch_size_plus_1,), "int32", elem_offset=kv_indptr_elem_offset
        )
        q_rope_position = T.match_buffer(
            var_q_rope_position, (qo_len,), "int32", elem_offset=q_rope_position_elem_offset
        )
        mn_indptr = T.match_buffer(
            var_mn_indptr, (batch_size_plus_1,), "int32", elem_offset=mn_indptr_elem_offset
        )
        mask = T.match_buffer(var_mask, (tree_size, 2), "int32", elem_offset=mask_elem_offset)
        output = T.match_buffer(var_output, (qo_len, h_q, d), dtype)
        lse = T.match_buffer(var_lse, (qo_len, h_q), "float32")  # pylint: disable=unused-variable

        for b in T.serial(batch_size_plus_1 - 1):
            with T.sblock("attn"):

                softmax_sum = T.sblock_alloc_buffer([h_q], "float32")
                m_prev = T.sblock_alloc_buffer([h_q], "float32")
                m_new = T.sblock_alloc_buffer([h_q], "float32")
                d_prev = T.sblock_alloc_buffer([h_q], "float32")
                d_new = T.sblock_alloc_buffer([h_q], "float32")
                p_sum = T.sblock_alloc_buffer([d], "float32")

                max_score = T.sblock_alloc_buffer([h_q], "float32")
                attention_scores = T.sblock_alloc_buffer([kv_len, h_q], "float32")
                exp_scores = T.sblock_alloc_buffer([kv_len, h_q], "float32")
                attention_score = T.sblock_alloc_buffer(
                    [
                        1,
                    ],
                    "float32",
                )
                query_val = T.sblock_alloc_buffer(
                    [
                        1,
                    ],
                    "float32",
                )
                key_val = T.sblock_alloc_buffer(
                    [
                        1,
                    ],
                    "float32",
                )
                result = T.sblock_alloc_buffer(
                    [
                        1,
                    ],
                    "float32",
                )

                for q_idx in T.serial(q_indptr[b + 1] - q_indptr[b]):
                    for i in T.serial(h_q):
                        max_score[i] = -5e4
                        m_prev[i] = -5e4
                        d_prev[i] = 1.0

                    for k_idx in T.serial(kv_indptr[b + 1] - kv_indptr[b]):
                        for h in T.serial(h_q):
                            h_kv_idx: T.let[T.int32] = h // group_size

                            if _check_tree_order(
                                row=q_idx,
                                col=k_idx,
                                batch=b,
                                tree_order=mask,
                                tree_order_indptr=mn_indptr,
                                kv_len=kv_indptr[b + 1] - kv_indptr[b],
                                qo_len=q_indptr[b + 1] - q_indptr[b],
                            ):
                                result[0] = 0.0
                                for d_idx in T.serial(d):
                                    query_val[0] = T.if_then_else(
                                        rotary_mode == 1,
                                        _rope(
                                            q,
                                            q_rope_position[q_indptr[b] + q_idx],
                                            d,
                                            rope_theta,
                                            rope_scale,
                                            (q_indptr[b] + q_idx, h, d_idx),
                                            dtype,
                                            rope_scaling,
                                        ),
                                        q[q_indptr[b] + q_idx, h, d_idx],
                                    )

                                    key_val[0] = T.if_then_else(
                                        rotary_mode == 1,
                                        _rope(
                                            k,
                                            q_rope_position[kv_indptr[b] + k_idx],
                                            d,
                                            rope_theta,
                                            rope_scale,
                                            (kv_indptr[b] + k_idx, h_kv_idx, d_idx),
                                            dtype,
                                            rope_scaling,
                                        ),
                                        k[kv_indptr[b] + k_idx, h_kv_idx, d_idx],
                                    )

                                    result[0] += query_val[0] * key_val[0]
                                attention_score[0] = (
                                    result[0] * math.log2(math.exp(1)) * sm_scale
                                )
                            else:
                                attention_score[0] = -5e4 * math.log2(math.exp(1)) * sm_scale
                            attention_scores[k_idx, h] = attention_score[0]
                            max_score[h] = T.max(max_score[h], attention_score[0])
                            m_new[h] = T.max(m_prev[h], max_score[h])

                    for h in T.serial(h_q):
                        d_new[h] = d_prev[h] * T.exp2(m_prev[h] - m_new[h])

                    for h in T.serial(h_q):
                        softmax_sum[h] = 0.0
                        for k_idx in T.serial(kv_indptr[b + 1] - kv_indptr[b]):
                            exp_scores[k_idx, h] = T.exp2(attention_scores[k_idx, h] - m_new[h])
                            softmax_sum[h] += exp_scores[k_idx, h]
                        d_new[h] += softmax_sum[h]

                    for h in T.serial(h_q):
                        h_kv_idx: T.let[T.int32] = h // group_size
                        for i in T.serial(d):
                            p_sum[i] = 0.0
                        for v_idx in T.serial(kv_indptr[b + 1] - kv_indptr[b]):
                            weight: T.let[T.float32] = exp_scores[v_idx, h] / d_new[h]
                            for i in T.serial(d):
                                p_sum[i] += v[kv_indptr[b] + v_idx, h_kv_idx, i] * weight
                        for i in T.serial(d):
                            output[q_indptr[b] + q_idx, h, i] = p_sum[i]
                        lse[q_indptr[b] + q_idx, h] = m_new[h] + T.log2(d_new[h])

    # fmt: on
    # pylint: enable=line-too-long,too-many-branches
    return batch_tree_attn


def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target):  # pylint: disable=unused-argument
    """Generate tree attention kernel for batched tree attention.

    Parameters
    ----------
    h_kv : int
        Number of heads for key and value.
    h_q : int
        Number of heads for query.
    d : int
        Hidden dimension.
    dtype : str
        Data type.
    target : Target
        The target device.

    Returns
    -------
    mod : tvm.IRModule
        The generated IR module.
    """
    # pylint: disable=invalid-name,line-too-long
    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

    # fmt: off
    @T.prim_func(s_tir=True)
    def batch_tree_attn(  # pylint: disable=too-many-branches
        var_q: T.handle, # [total_len, h_q, d]
        var_q_indptr: T.handle, # [batch_size + 1]
        var_k: T.handle, # [total_len, h_kv, d]
        var_v: T.handle, # [total_len, h_kv, d]
        var_kv_indptr: T.handle, # [batch_size + 1], kv_indptr should be the same as q_indptr in this case
        var_q_rope_position: T.handle, # [total_q_len]
        var_mn_indptr: T.handle, # [batch_size + 1]
        var_mask: T.handle, # [mn_indptr[batch_size]]
        var_output: T.handle, # [total_len, h_q, d]
        var_lse: T.handle, # [total_len, h_q]
        rotary_mode: T.int32,
        rope_scale: T.float32,
        rope_theta: T.float32,
        sm_scale: T.float32,
    ):
        qo_len = T.int32()
        kv_len = T.int32()
        q_indptr_elem_offset = T.int32()
        kv_indptr_elem_offset = T.int32()
        q_rope_position_elem_offset = T.int32()
        mn_indptr_elem_offset = T.int32()
        mask_elem_offset = T.int32()
        tree_size = T.int32()
        batch_size_plus_1 = T.int32()

        q = T.match_buffer(var_q, (qo_len, h_q, d), dtype)
        q_indptr = T.match_buffer(var_q_indptr, (batch_size_plus_1,), "int32", elem_offset=q_indptr_elem_offset)
        k = T.match_buffer(var_k, (kv_len, h_kv, d), dtype)
        v = T.match_buffer(var_v, (kv_len, h_kv, d), dtype)
        kv_indptr = T.match_buffer(var_kv_indptr, (batch_size_plus_1,), "int32", elem_offset=kv_indptr_elem_offset)
        q_rope_position = T.match_buffer(var_q_rope_position, (qo_len,), "int32", elem_offset=q_rope_position_elem_offset)
        mn_indptr = T.match_buffer(var_mn_indptr, (batch_size_plus_1,), "int32", elem_offset=mn_indptr_elem_offset)
        mask = T.match_buffer(var_mask, (tree_size, 2), "int32", elem_offset=mask_elem_offset)
        output = T.match_buffer(var_output, (qo_len, h_q, d), dtype)
        lse = T.match_buffer(var_lse, (qo_len, h_q), "float32")  # pylint: disable=unused-variable

        # kernel code
        for lbx in T.thread_binding(NUM_BLKS, thread="blockIdx.x"):
            for lby in T.thread_binding(h_kv, thread="blockIdx.y"):
                for lty in T.thread_binding(num_warps, thread="threadIdx.y"):
                    for ltx in T.thread_binding(bdx, thread="threadIdx.x"):
                        with T.sblock("attn"):
                            bx, by, ty, tx = T.axis.remap("SSSS", [lbx, lby, lty, ltx])
                            T.reads()
                            T.writes()
                            tile_id, batch_idx, batch_tiles, batch_rows, iterator, kv_chunk_len = _alloc_tile_walk_state()
                            Q_smem, K_smem, V_smem, O_local = _alloc_mha_qkvo_buffers(tile_x, tile_z, d, d, dtype)
                            S_smem, S_local, m_smem, m_prev_smem, d_smem, m_new, m_prev, d_new = _alloc_softmax_state_buffers(tile_x, tile_z, bdx, num_warps)

                            tile_id[0] = bx
                            batch_idx[0] = 0
                            batch_rows[0] = (q_indptr[1] - q_indptr[0]) * group_size
                            batch_tiles[0] = T.ceildiv(batch_rows[0], tile_x)
                            while T.tvm_thread_invariant(batch_idx[0] < batch_size_plus_1 - 1):
                                # advance to next tile
                                while tile_id[0] >= batch_tiles[0] and batch_idx[0] < batch_size_plus_1 - 1:
                                    tile_id[0] -= batch_tiles[0]
                                    batch_idx[0] += 1
                                    if batch_idx[0] < batch_size_plus_1 - 1:
                                        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)

                                if T.tvm_thread_invariant(batch_idx[0] < batch_size_plus_1 - 1):
                                    b_idx: T.let[T.int32()] = batch_idx[0]
                                    LH_start: T.let[T.int32()] = tile_id[0] * tile_x
                                    q_indptr_val: T.let[T.int32] = q_indptr[b_idx]

                                    kv_chunk_len[0] = kv_indptr[b_idx + 1] - kv_indptr[b_idx]
                                    T.tvm_storage_sync("shared")

                                    # init states
                                    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_y):
                                        with T.sblock("O_init"):
                                            i, j = T.axis.remap("SS", [li, lj])
                                            O_local[i, j] = 0.0
                                    T.tvm_storage_sync("shared")

                                    # Load Q from gmem to smem
                                    for li, lj in T.grid(tile_x, tile_y):
                                        with T.sblock("Q_load"):
                                            i, j = T.axis.remap("SS", [li, lj])
                                            T.reads()
                                            T.writes()
                                            cur_L: T.let[T.int32] = q_indptr_val + (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]:
                                                Q_smem[i, j] = T.if_then_else(
                                                    rotary_mode == 1,
                                                    _rope(q, q_rope_position[cur_L], d, rope_theta, rope_scale, (cur_L, cur_H_qo, j), dtype, rope_scaling),
                                                    q[cur_L, cur_H_qo, j]
                                                )
                                            else:
                                                Q_smem[i, j] = 0.0
                                    T.tvm_storage_sync("shared")

                                    for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_z)):
                                        L_kv_start: T.let[T.int32] = iterator * tile_z
                                        L_kv_base: T.let[T.int32] = kv_indptr[b_idx]
                                        for lz, ly in T.grid(tile_z, tile_y):
                                            with T.sblock("KV_load"):
                                                i, j = T.axis.remap("SS", [lz, ly])
                                                T.reads()
                                                T.writes()
                                                cur_L: T.let[T.int32] = L_kv_base + L_kv_start + i
                                                if L_kv_start + i < kv_chunk_len[0]:
                                                    K_smem[i, j] = T.if_then_else(
                                                        rotary_mode == 1,
                                                        _rope(k, q_rope_position[cur_L], d, rope_theta, rope_scale, (cur_L, by, j), dtype, rope_scaling),
                                                        k[cur_L, by, j]
                                                    )
                                                    V_smem[i, j] = v[cur_L, by, j]
                                                else:
                                                    K_smem[i, j] = 0.0
                                                    V_smem[i, j] = 0.0
                                        T.tvm_storage_sync("shared")

                                        # Compute S
                                        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")

                                        # Update S, m, d
                                        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]
                                                    # mask out of kv_chunk_len S
                                                    row_: T.let[T.int32] = (LH_start + row) // group_size
                                                    for j in T.serial(tile_z):
                                                        if _check_tree_order(
                                                            row=row_,
                                                            col=L_kv_start + j,
                                                            batch=b_idx,
                                                            tree_order=mask,
                                                            tree_order_indptr=mn_indptr,
                                                            qo_len=q_indptr[b_idx + 1] - q_indptr[b_idx],
                                                            kv_len=kv_chunk_len[0]):
                                                            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):
                                                    # this is to avoid sync inside condition branch
                                                    if row < tile_x:
                                                        row_: T.let[T.int32] = (LH_start + row) // group_size
                                                        if _check_tree_order(
                                                            row=row_,
                                                            col=L_kv_start + j,
                                                            batch=b_idx,
                                                            tree_order=mask,
                                                            tree_order_indptr=mn_indptr,
                                                            qo_len=q_indptr[b_idx + 1] - q_indptr[b_idx],
                                                            kv_len=kv_chunk_len[0]):
                                                            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")

                                        # Update O
                                        with T.sblock():
                                            for li, lj, lk in T.grid(tile_x, tile_y, 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")

                                    # Store O from smem to gmem
                                    for li, lj in T.grid(tile_x, tile_y):
                                        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]

                                    # Store LSE to gmem
                                    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])

                                    # move to next tile
                                    tile_id[0] += NUM_BLKS
    # fmt: on
    # pylint: enable=line-too-long,too-many-branches
    sch = s_tir.Schedule(batch_tree_attn)

    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) 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:]
        loop = sch.fuse(loop_x, loop_y)
        _, ty, tx, vec = sch.split(
            loop, factors=[None, num_warps, bdx, LOAD_VEC], preserve_unit_iters=True
        )
        sch.bind(ty, "threadIdx.y")
        sch.bind(tx, "threadIdx.x")
        sch.vectorize(vec)

    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)
        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(  # pylint: disable=unused-argument
        sch: s_tir.Schedule, block, tile, read_0, read_1, r_len=8, 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)
        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")

    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, 0, 1, k_major=True)
    apply_to_gemm(sch, sch.get_sblock("O_gemm"), tile_o, 2, 3, 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"))
    apply_to_qkv_load(sch, sch.get_sblock("KV_load"))

    apply_to_md(sch, sch.get_sblock("lse_store"))
    return sch.mod["main"].with_attr("tirx.is_scheduled", True)


def tree_attn_with_paged_kv_cache_cpu(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any]):
    """Generate tree attention kernel for batched tree attention with paged key-value cache.

    Parameters
    ----------
    h_kv : int
        Number of heads for key and value.
    h_q : int
        Number of heads for query.
    d : int
        Hidden dimension.
    dtype : str
        Data type.
    target : Target
        The target device.

    Returns
    -------
    mod : tvm.IRModule
        The generated IR module.
    """
    global_symbol = "tree_attn_paged_kv_cpu"
    sliding_window = False
    group_size = h_q // h_kv

    # pylint: disable=line-too-long,too-many-branches
    # fmt: off
    @T.prim_func(s_tir=True)
    def tree_attn_paged_kv_cpu(
        var_q: T.handle, # [total_len, h_q, d]
        var_q_indptr: T.handle, # [batch_size + 1]
        var_pages: T.handle, # [max_num_pages, 2, h_kv, page_size, d]
        var_page_indptr: T.handle, # [batch_size + 1]
        var_page_values: T.handle, # [nnz_pages]
        var_length_info: T.handle, # [b] when sliding window = False, or otherwise [3, b]
        var_k_rope_pos_offset: T.handle, # [b]
        var_q_rope_position: T.handle, # [total_len]
        var_output: T.handle, # [total_len, h_q, d]
        var_lse: T.handle, # [total_len, h_q]
        rotary_mode: T.int32,
        rope_scale: T.float32,
        rope_theta: T.float32,
        sm_scale: T.float32,
        tree_order_indptr_handle: T.handle,  # [batch_size + 1]
        tree_order_handle: T.handle,  # [total_len, 2]
    ):
        T.func_attr({"global_symbol": global_symbol})
        batch_size = T.int32()
        total_len = T.int32()
        nnz_pages = T.int32()
        max_num_pages = T.int32()
        q_indptr_elem_offset = 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()
        tree_order_elem_offset = T.int32()
        tree_order_indptr_elem_offset = T.int32()

        q = T.match_buffer(var_q, (total_len, h_q, d), dtype)
        q_indptr = T.match_buffer(var_q_indptr, (batch_size + 1,), "int32", elem_offset=q_indptr_elem_offset)
        pages = T.match_buffer(var_pages, (max_num_pages, 2, h_kv, 16, d), dtype)
        page_indptr = T.match_buffer(var_page_indptr, (batch_size + 1,), "int32", elem_offset=page_indptr_elem_offset)
        page_values = T.match_buffer(var_page_values, (nnz_pages,), "int32", elem_offset=page_values_elem_offset)
        k_rope_pos_offset = T.match_buffer(var_k_rope_pos_offset, (batch_size,), "int32", elem_offset=k_rope_pos_offset_elem_offset)
        q_rope_position = T.match_buffer(var_q_rope_position, (total_len,), "int32", elem_offset=q_rope_position_elem_offset)
        output = T.match_buffer(var_output, (total_len, h_q, d), dtype)
        lse = T.match_buffer(var_lse, (total_len, h_q), "float32")  # pylint: disable=unused-variable
        tree_order_indptr = T.match_buffer(
            tree_order_indptr_handle,
            (batch_size + 1,),
            "int32",
            elem_offset=tree_order_indptr_elem_offset,
        )
        total_tree_order_len = T.int32()
        tree_order = T.match_buffer(
            tree_order_handle,
            (total_tree_order_len, 2),
            "int32",
            elem_offset=tree_order_elem_offset,
        )
        # 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, batch_size, sliding_window, length_info_elem_offset)


        T.Assert(
            rotary_mode == T.int32(0), "Inline rotary mode is not supported in tree attention."
        )

        for h_qo in T.serial(h_q):
            for b_idx in T.serial(batch_size):
                with T.sblock("attn"):
                    T.reads()
                    T.writes()
                    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_indptr[b_idx]
                    cur_page_indptr_end: T.let[T.int32] = page_indptr[b_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, 16, b_idx, length_info, sliding_window),
                        0
                    )

                    for q_idx in T.serial(q_indptr[b_idx + 1] - q_indptr[b_idx]):
                        #init m, d, O
                        m_val[0] = -5e4
                        d_val[0] = 1.0
                        for d_idx in T.serial(d):
                            O_local[d_idx] = 0.0
                        curl_q: T.let[T.int32] = q_indptr[b_idx] + q_idx

                        for d_idx in T.serial(d):
                            Q_local[d_idx] = T.if_then_else(
                                rotary_mode == 1,
                                _rope(q, q_rope_position[curl_q], d, rope_theta, rope_scale, (curl_q, h_qo, d_idx), dtype, rope_scaling),
                                q[curl_q, h_qo, d_idx]
                            )
                        for row_idx in T.serial(max_num_pages * 16):
                            if row_idx < kv_chunk_len[0]:
                                page_no: T.let[T.int32()] = page_values[cur_page_indptr_begin + (_get_seq_offset(row_idx, b_idx, length_info, sliding_window) // 16)]
                                page_offset: T.let[T.int32()] = _get_seq_offset(row_idx, b_idx, length_info, sliding_window) % 16

                                # Load KV
                                for d_idx in T.serial(d):
                                    K_local[d_idx] = T.if_then_else(
                                        rotary_mode == 1,
                                        _rope(pages, k_rope_pos_offset[b_idx] + row_idx, d, rope_theta, rope_scale, (page_no, 0, h_qo // group_size, page_offset, d_idx), dtype, rope_scaling),
                                        pages[page_no, 0, h_qo // group_size, page_offset, d_idx]
                                    )
                                    V_local[d_idx] = pages[page_no, 1, h_qo // group_size, page_offset, d_idx]

                                # Compute S
                                S_val[0] = 0.0
                                for d_idx in T.serial(d):
                                    S_val[0] += Q_local[d_idx] * K_local[d_idx]
                                S_val[0] *= sm_scale * math.log2(math.exp(1))

                                # update m_val, d_val , O_local
                                if _check_tree_order(
                                    tree_order_indptr=tree_order_indptr,
                                    tree_order=tree_order,
                                    batch=b_idx,
                                    row=q_idx,
                                    col=row_idx,
                                    kv_len=kv_chunk_len[0],
                                    qo_len=q_indptr[b_idx + 1] - q_indptr[b_idx],
                                ):
                                    new_m[0] = T.max(m_val[0], S_val[0])
                                else:
                                    S_val[0] = -5e4
                                # update d_val
                                d_val[0] *= T.exp2(m_val[0] - new_m[0])
                                d_val[0] += T.exp2(S_val[0] - new_m[0])

                                # restore O_local then update O_local
                                scale_O[0] = T.exp2(m_val[0] - new_m[0])
                                m_val[0] = new_m[0]
                                factor[0] = T.exp2(S_val[0] - m_val[0])
                                for d_idx in T.serial(d):
                                    O_local[d_idx] = O_local[d_idx] * scale_O[d_idx]


                                for d_idx in T.serial(d):
                                    O_local[d_idx] += V_local[d_idx] * factor[0]
                        # Store Output
                        for d_idx in T.serial(d):
                            O_local[d_idx] = O_local[d_idx] /d_val[0]
                            output[curl_q, h_qo, d_idx] = O_local[d_idx]
                        lse[curl_q, h_qo] = m_val[0] + T.log2(d_val[0])
    return tree_attn_paged_kv_cpu


def tree_attn_with_paged_kv_cache(
    h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target
):
    """Generate tree attention kernel for batched tree attention with paged key-value cache.

    Parameters
    ----------
    h_kv : int
        Number of heads for key and value.
    h_q : int
        Number of heads for query.
    d : int
        Hidden dimension.
    dtype : str
        Data type.
    target : Target
        The target device.

    Returns
    -------
    mod : tvm.IRModule
        The generated IR module.
    """
    # pylint: disable=invalid-name, line-too-long
    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
    check_thread_limits(target, bdx=bdx, bdy=num_warps, bdz=1, gdz=1)

    global_symbol = "tree_attn_paged_kv"
    sliding_window = False  # Sliding window is not supported in this kernel.

    # fmt: off
    @T.prim_func(s_tir=True)
    def tree_attn_paged_kv(
        var_q: T.handle,  # [total_len, h_q, d]
        var_q_indptr: T.handle,  # [batch_size + 1]
        var_pages: T.handle,  # [max_num_pages, 2, h_kv, page_size, d]
        var_page_indptr: T.handle,  # [batch_size + 1]
        var_page_values: T.handle,  # [nnz_pages]
        var_length_info: T.handle,  # [b] when sliding window = False, or otherwise [3, b]
        var_k_rope_pos_offset: T.handle,  # [b]
        var_q_rope_position: T.handle,  # [total_len]
        var_output: T.handle,  # [total_len, h_q, d]
        var_lse: T.handle,  # [total_len, h_q]
        rotary_mode: T.int32,
        rope_scale: T.float32,
        rope_theta: T.float32,
        sm_scale: T.float32,
        tree_order_indptr_handle: T.handle,  # [batch_size + 1]
        tree_order_handle: T.handle,  # [total_len, 2]
    ):
        # pylint: disable=unused-variable, too-many-branches
        T.func_attr({"global_symbol": global_symbol})
        batch_size = T.int32()
        total_len = T.int32()
        nnz_pages = T.int32()
        max_num_pages = T.int32()
        q_indptr_elem_offset = T.int32()
        k_rope_pos_offset_elem_offset = T.int32()
        q_rope_position_elem_offset = T.int32()
        page_indptr_elem_offset = T.int32()
        page_values_elem_offset = T.int32()
        length_info_elem_offset = T.int32()
        tree_order_elem_offset = T.int32()
        tree_order_indptr_elem_offset = T.int32()

        q = T.match_buffer(var_q, (total_len, h_q, d), dtype)
        q_indptr = T.match_buffer(
            var_q_indptr, (batch_size + 1,), "int32", elem_offset=q_indptr_elem_offset
        )
        pages = T.match_buffer(var_pages, (max_num_pages, 2, h_kv, 16, d), dtype)
        page_indptr = T.match_buffer(
            var_page_indptr, (batch_size + 1,), "int32", elem_offset=page_indptr_elem_offset
        )
        page_values = T.match_buffer(
            var_page_values, (nnz_pages,), "int32", elem_offset=page_values_elem_offset
        )
        k_rope_pos_offset = T.match_buffer(
            var_k_rope_pos_offset, (batch_size,), "int32", elem_offset=k_rope_pos_offset_elem_offset
        )
        q_rope_position = T.match_buffer(
            var_q_rope_position, (total_len,), "int32", elem_offset=q_rope_position_elem_offset
        )
        output = T.match_buffer(var_output, (total_len, h_q, d), dtype)
        lse = T.match_buffer(
            var_lse, (total_len, h_q), "float32"
        )  # pylint: disable=unused-variable
        tree_order_indptr = T.match_buffer(
            tree_order_indptr_handle,
            (batch_size + 1,),
            "int32",
            elem_offset=tree_order_indptr_elem_offset,
        )
        total_tree_order_len = T.int32()
        tree_order = T.match_buffer(
            tree_order_handle,
            (total_tree_order_len, 2),
            "int32",
            elem_offset=tree_order_elem_offset,
        )
        # 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, batch_size, sliding_window, length_info_elem_offset
        )

        T.Assert(
            rotary_mode == T.int32(0), "Inline rotary mode is not supported in tree attention."
        )

        # kernel code
        for lbx in T.thread_binding(NUM_BLKS, thread="blockIdx.x"):
            for lby in T.thread_binding(h_kv, thread="blockIdx.y"):
                for lty in T.thread_binding(num_warps, thread="threadIdx.y"):
                    for ltx in T.thread_binding(bdx, thread="threadIdx.x"):
                        with T.sblock("attn"):
                            bx, by, ty, tx = T.axis.remap("SSSS", [lbx, lby, lty, ltx])
                            T.reads()
                            T.writes()
                            tile_id, batch_idx, batch_tiles, batch_rows, iterator, kv_chunk_len = _alloc_tile_walk_state()
                            Q_smem, K_smem, V_smem, O_local = _alloc_mha_qkvo_buffers(tile_x, tile_z, d, d, dtype)
                            S_smem, S_local, m_smem, m_prev_smem, d_smem, m_new, m_prev, d_new = _alloc_softmax_state_buffers(tile_x, tile_z, bdx, num_warps)

                            tile_id[0] = bx
                            batch_idx[0] = 0
                            batch_rows[0] = (q_indptr[1] - q_indptr[0]) * group_size
                            batch_tiles[0] = T.ceildiv(batch_rows[0], tile_x)
                            while T.tvm_thread_invariant(batch_idx[0] < batch_size):
                                # advance to next tile
                                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)

                                if T.tvm_thread_invariant(batch_idx[0] < batch_size):
                                    b_idx: T.let[T.int32()] = batch_idx[0]
                                    LH_start: T.let[T.int32()] = tile_id[0] * tile_x
                                    q_indptr_val: T.let[T.int32] = q_indptr[b_idx]

                                    cur_page_indptr_begin: T.let[T.int32] = page_indptr[b_idx]
                                    cur_page_indptr_end: T.let[T.int32] = page_indptr[b_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,
                                            16,
                                            b_idx,
                                            length_info,
                                            sliding_window,
                                        ),
                                        0,
                                    )
                                    T.tvm_storage_sync("shared")

                                    # init states
                                    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_y):
                                        with T.sblock("O_init"):
                                            i, j = T.axis.remap("SS", [li, lj])
                                            O_local[i, j] = 0.0
                                    T.tvm_storage_sync("shared")

                                    # Load Q from gmem to smem
                                    for li, lj in T.grid(tile_x, tile_y):
                                        with T.sblock("Q_load"):
                                            i, j = T.axis.remap("SS", [li, lj])
                                            T.reads()
                                            T.writes()
                                            cur_L: T.let[T.int32] = q_indptr_val + (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]:
                                                Q_smem[i, j] = T.if_then_else(
                                                    rotary_mode == 1,
                                                    _rope(
                                                        q,
                                                        q_rope_position[cur_L],
                                                        d,
                                                        rope_theta,
                                                        rope_scale,
                                                        (cur_L, cur_H_qo, j),
                                                        dtype,
                                                        rope_scaling,
                                                    ),
                                                    q[cur_L, cur_H_qo, j],
                                                )
                                            else:
                                                Q_smem[i, j] = 0.0
                                    T.tvm_storage_sync("shared")

                                    for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_z)):
                                        L_kv_start: T.let[T.int32] = iterator * tile_z
                                        for lz, ly in T.grid(tile_z, tile_y):
                                            with T.sblock("K_load"):
                                                i, j = T.axis.remap("SS", [lz, ly])
                                                T.reads()
                                                T.writes()
                                                cur_L: T.let[T.int32] = L_kv_start + i
                                                if cur_L < kv_chunk_len[0]:
                                                    seq_offset: T.let[T.int32()] = _get_seq_offset(cur_L, b_idx, length_info, sliding_window)  # type: ignore
                                                    page_no: T.let[T.int32()] = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, 16)]  # type: ignore
                                                    page_offset: T.let[T.int32()] = T.floormod(seq_offset, 16)  # type: ignore
                                                    K_smem[i, j] = pages[
                                                        page_no, 0, by, page_offset, j
                                                    ]
                                                else:
                                                    K_smem[i, j] = 0.0

                                        T.tvm_storage_sync("shared")
                                        for lz, ly in T.grid(tile_z, tile_y):
                                            with T.sblock("V_load"):
                                                i, j = T.axis.remap("SS", [lz, ly])
                                                T.reads()
                                                T.writes()
                                                cur_L: T.let[T.int32] = L_kv_start + i
                                                if cur_L < kv_chunk_len[0]:
                                                    seq_offset: T.let[T.int32()] = _get_seq_offset(cur_L, b_idx, length_info, sliding_window)  # type: ignore
                                                    page_no: T.let[T.int32()] = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, 16)]  # type: ignore
                                                    page_offset: T.let[T.int32()] = T.floormod(seq_offset, 16)  # type: ignore
                                                    V_smem[i, j] = pages[
                                                        page_no, 1, by, page_offset, j
                                                    ]
                                                else:
                                                    V_smem[i, j] = 0.0
                                        T.tvm_storage_sync("shared")

                                        # Compute S
                                        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")

                                        # Update S, m, d
                                        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]
                                                    # mask out of kv_chunk_len S
                                                    row_: T.let[T.int32] = (LH_start + row) // group_size
                                                    for j in T.serial(tile_z):
                                                        if _check_tree_order(
                                                            tree_order_indptr=tree_order_indptr,
                                                            tree_order=tree_order,
                                                            batch=b_idx,
                                                            row=row_,
                                                            col=L_kv_start + j,
                                                            kv_len=kv_chunk_len[0],
                                                            qo_len=q_indptr[b_idx + 1]
                                                            - q_indptr[b_idx],
                                                        ):
                                                            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):
                                                    # this is to avoid sync inside condition branch
                                                    if row < tile_x:
                                                        row_: T.let[T.int32] = (
                                                            LH_start + row
                                                        ) // group_size
                                                        if _check_tree_order(
                                                            tree_order_indptr=tree_order_indptr,
                                                            tree_order=tree_order,
                                                            batch=b_idx,
                                                            row=row_,
                                                            col=L_kv_start + j,
                                                            kv_len=kv_chunk_len[0],
                                                            qo_len=q_indptr[b_idx + 1]
                                                            - q_indptr[b_idx],
                                                        ):
                                                            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")

                                        # Update O
                                        with T.sblock():
                                            for li, lj, lk in T.grid(tile_x, tile_y, 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"
                                                    )

                                    # Store O from smem to gmem
                                    for li, lj in T.grid(tile_x, tile_y):
                                        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]
                                                )

                                    # Store LSE to gmem
                                    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])

                                    # move to next tile
                                    tile_id[0] += NUM_BLKS

    # fmt: on
    # pylint: enable=line-too-long,too-many-branches
    sch = s_tir.Schedule(tree_attn_paged_kv)

    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) 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:]
        loop = sch.fuse(loop_x, loop_y)
        _, ty, tx, vec = sch.split(
            loop, factors=[None, num_warps, bdx, LOAD_VEC], preserve_unit_iters=True
        )
        sch.bind(ty, "threadIdx.y")
        sch.bind(tx, "threadIdx.x")
        sch.vectorize(vec)

    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)
        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(  # pylint: disable=unused-argument
        sch: s_tir.Schedule, block, tile, read_0, read_1, r_len=8, 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)
        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")

    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, 0, 1, k_major=True)
    apply_to_gemm(sch, sch.get_sblock("O_gemm"), tile_o, 2, 3, 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"))
    apply_to_qkv_load(sch, sch.get_sblock("K_load"))
    apply_to_qkv_load(sch, sch.get_sblock("V_load"))
    apply_to_md(sch, sch.get_sblock("lse_store"))
    return sch.mod["main"].with_attr("tirx.is_scheduled", True)
