Repository: mlc-ai/mlc-llm
File: python/mlc_chat/nn/kv_cache.py (Python)

Code before:
        )
        return cos + sin

    # fmt: off
    @T.prim_func
    def tir_rotary(
        var_q: T.handle,
        var_k: T.handle,
        var_append_len_indptr: T.handle,
        var_rope_offsets: T.handle,
        _0: T.int32,
        _1: T.int32,
        _2: T.int32,
        _3: T.int32,
        _4: T.int32,
        _5: T.float32,
        _6: T.float32,
    ):
        T.func_attr({"tir.is_scheduled": 1})
        total_len = T.int32()
        batch_size = T.int32()
        q = T.match_buffer(var_q, (total_len, num_q_heads, head_dim), dtype)
        k = T.match_buffer(var_k, (total_len, num_kv_heads, head_dim), dtype)
        rope_offsets = T.match_buffer(var_rope_offsets, (batch_size,), "int32")
        append_len_indptr = T.match_buffer(var_append_len_indptr, (batch_size + 1,), "int32")
        for b_h in T.thread_binding(batch_size * (num_q_heads + num_kv_heads), thread="blockIdx.x"):
            b: T.int32 = b_h // (num_q_heads + num_kv_heads)
            h: T.int32 = b_h % (num_q_heads + num_kv_heads)
            instance_offset: T.int32 = append_len_indptr[b]
            rope_offset: T.int32 = rope_offsets[b]
            append_len: T.int32 = append_len_indptr[b + 1] - append_len_indptr[b]
            for s0 in range(T.ceildiv(append_len, 32)):
                for s1 in T.thread_binding(32, thread="threadIdx.y"):
                    for d0 in T.thread_binding(T.ceildiv(head_dim, 4), thread="threadIdx.x"):
                        for d1 in T.vectorized(4):
                            s: T.int32 = s0 * 32 + s1
                            d: T.int32 = d0 * 4 + d1
                            if s < append_len and d < head_dim:
                                if h < num_q_heads:
                                    q[s + instance_offset, h, d] = _rope(q, s, h, d, rope_offset, instance_offset)
                                else:
                                    k[s + instance_offset, h - num_q_heads, d] = _rope(k, s, h - num_q_heads, d, rope_offset, instance_offset)
    return tir_rotary


# pylint: enable=line-too-long,too-many-arguments,too-many-nested-blocks,invalid-name


def _rope(  # pylint: disable=too-many-arguments
    buffer: T.Buffer,
    offset: tir.Var,

Review comment (suggestion):
The parser does not automatically add a block without `with T.block():`. Have not got a chance to see why.

Code after:
        )
        return cos + sin

    # fmt: off
    @T.prim_func
    def tir_rotary(
        var_q: T.handle,
        var_k: T.handle,
        var_append_len_indptr: T.handle,
        var_rope_offsets: T.handle,
        _0: T.int32,
        _1: T.int32,
        _2: T.int32,
        _3: T.int32,
        _4: T.int32,
        _5: T.float32,
        _6: T.float32,
    ):
        T.func_attr({"tir.is_scheduled": 1})
        total_len = T.int32()
        batch_size = T.int32()
        q = T.match_buffer(var_q, (total_len, num_q_heads, head_dim), dtype)
        k = T.match_buffer(var_k, (total_len, num_kv_heads, head_dim), dtype)
        rope_offsets = T.match_buffer(var_rope_offsets, (batch_size,), "int32")
        append_len_indptr = T.match_buffer(var_append_len_indptr, (batch_size + 1,), "int32")
        with T.block():
            for b_h in T.thread_binding(batch_size * (num_q_heads + num_kv_heads), thread="blockIdx.x"):
                b: T.int32 = b_h // (num_q_heads + num_kv_heads)
                h: T.int32 = b_h % (num_q_heads + num_kv_heads)
                instance_offset: T.int32 = append_len_indptr[b]
                rope_offset: T.int32 = rope_offsets[b]
                append_len: T.int32 = append_len_indptr[b + 1] - append_len_indptr[b]
                for s0 in range(T.ceildiv(append_len, 32)):
                    for s1 in T.thread_binding(32, thread="threadIdx.y"):
                        for d0 in T.thread_binding(T.ceildiv(head_dim, 4), thread="threadIdx.x"):
                            for d1 in T.vectorized(4):
                                s: T.int32 = s0 * 32 + s1
                                d: T.int32 = d0 * 4 + d1
                                if s < append_len and d < head_dim:
                                    if h < num_q_heads:
                                        q[s + instance_offset, h, d] = _rope(q, s, h, d, rope_offset, instance_offset)
                                    else:
                                        k[s + instance_offset, h - num_q_heads, d] = _rope(k, s, h - num_q_heads, d, rope_offset, instance_offset)
    return tir_rotary


# pylint: enable=line-too-long,too-many-arguments,too-many-nested-blocks,invalid-name


def _rope(  # pylint: disable=too-many-arguments
    buffer: T.Buffer,