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

Code before:
        and ((d + 127) // 128) * ((DataType(dtype).bits + 15) // 16) >= 4
    ):
        tile_z = 8
        num_warps = 2

    # fmt: off
    @T.prim_func
    def batch_sequence_prefill_kv(  # pylint: disable=too-many-branches
        var_q: T.handle, # [total_len, h_q, d]
        var_k: T.handle, # [total_len, h_kv, d]
        var_v: T.handle, # [total_len, h_kv, d]
        var_output: T.handle, # [total_len, h_q, d]
        var_lse: T.handle # [total_len, h_q]
    ):
        qo_len = T.int32(is_size_var=True)
        kv_len = T.int32(is_size_var=True)
        q = T.match_buffer(var_q, (batch_size, qo_len, h_q, d), dtype)
        k = T.match_buffer(var_k, (batch_size, kv_len, h_kv, d), dtype)
        v = T.match_buffer(var_v, (batch_size, kv_len, h_kv, d), dtype)
        output = T.match_buffer(var_output, (batch_size, qo_len, h_q, d), dtype)
        lse = T.match_buffer(var_lse, (batch_size, qo_len, h_q), dtype)  # pylint: disable=unused-variable

        batch_tiles: T.int32 = T.ceildiv(qo_len * group_size, tile_x)

        # kernel code
        for lbx in T.thread_binding(batch_size * batch_tiles, 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.block("attn"):
                            vbx, by, ty, tx = T.axis.remap("SSSS", [lbx, lby, lty, ltx])
                            T.reads()
                            T.writes()

                            Q_smem = T.alloc_buffer((tile_x, d), dtype, scope="shared")
                            K_smem = T.alloc_buffer((tile_z, d), dtype, scope="shared")
                            V_smem = T.alloc_buffer((tile_z, d), dtype, scope="shared")
                            S_smem = T.alloc_buffer((tile_x, tile_z), "float32", scope="shared")

                            S_local = T.alloc_buffer((tile_x, tile_z), "float32", scope="local")
                            O_local = T.alloc_buffer((tile_x, d), "float32", scope="local")

                            m_smem = T.alloc_buffer((tile_x,), "float32", scope="shared")
                            m_prev_smem = T.alloc_buffer((tile_x,), "float32", scope="shared")
                            d_smem = T.alloc_buffer((tile_x,), "float32", scope="shared")

                            m_new = T.alloc_buffer(
                                (math.ceil(tile_x / (bdx * num_warps)),), "float32", scope="local"
                            )
                            m_prev = T.alloc_buffer(
                                (math.ceil(tile_x / (bdx * num_warps)),), "float32", scope="local"

Review comment (performance):
```suggestion
        for lbx in T.thread_binding(T.cast(batch_size, "int32") * batch_tiles, thread="blockIdx.x"):
```

Code after:
        and ((d + 127) // 128) * ((DataType(dtype).bits + 15) // 16) >= 4
    ):
        tile_z = 8
        num_warps = 2

    # fmt: off
    @T.prim_func
    def batch_sequence_prefill_kv(  # pylint: disable=too-many-branches
        var_q: T.handle, # [total_len, h_q, d]
        var_k: T.handle, # [total_len, h_kv, d]
        var_v: T.handle, # [total_len, h_kv, d]
        var_output: T.handle, # [total_len, h_q, d]
        var_lse: T.handle # [total_len, h_q]
    ):
        qo_len = T.int32(is_size_var=True)
        kv_len = T.int32(is_size_var=True)
        q = T.match_buffer(var_q, (batch_size, qo_len, h_q, d), dtype)
        k = T.match_buffer(var_k, (batch_size, kv_len, h_kv, d), dtype)
        v = T.match_buffer(var_v, (batch_size, kv_len, h_kv, d), dtype)
        output = T.match_buffer(var_output, (batch_size, qo_len, h_q, d), dtype)
        lse = T.match_buffer(var_lse, (batch_size, qo_len, h_q), dtype)  # pylint: disable=unused-variable

        batch_tiles: T.int32 = T.ceildiv(qo_len * group_size, tile_x)

        # kernel code
        for lbx in T.thread_binding(T.cast(batch_size, "int32") * batch_tiles, 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.block("attn"):
                            vbx, by, ty, tx = T.axis.remap("SSSS", [lbx, lby, lty, ltx])
                            T.reads()
                            T.writes()

                            Q_smem = T.alloc_buffer((tile_x, d), dtype, scope="shared")
                            K_smem = T.alloc_buffer((tile_z, d), dtype, scope="shared")
                            V_smem = T.alloc_buffer((tile_z, d), dtype, scope="shared")
                            S_smem = T.alloc_buffer((tile_x, tile_z), "float32", scope="shared")

                            S_local = T.alloc_buffer((tile_x, tile_z), "float32", scope="local")
                            O_local = T.alloc_buffer((tile_x, d), "float32", scope="local")

                            m_smem = T.alloc_buffer((tile_x,), "float32", scope="shared")
                            m_prev_smem = T.alloc_buffer((tile_x,), "float32", scope="shared")
                            d_smem = T.alloc_buffer((tile_x,), "float32", scope="shared")

                            m_new = T.alloc_buffer(
                                (math.ceil(tile_x / (bdx * num_warps)),), "float32", scope="local"
                            )
                            m_prev = T.alloc_buffer(
                                (math.ceil(tile_x / (bdx * num_warps)),), "float32", scope="local"