wehub-resource-sync
1260 lines
65 KiB
Python
1260 lines
65 KiB
Python
# 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)
|