wehub-resource-sync
244 lines
8.1 KiB
Python
244 lines
8.1 KiB
Python
import math
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.relax.frontend.nn.llm import tree_attn
|
|
|
|
# test category "op_correctness"
|
|
pytestmark = [pytest.mark.op_correctness]
|
|
|
|
|
|
@pytest.mark.parametrize("nbatch", [1, 4, 32])
|
|
@pytest.mark.parametrize("h_q", [8, 16])
|
|
@pytest.mark.parametrize("h_kv", [4, 8])
|
|
@pytest.mark.parametrize("d", [128])
|
|
@pytest.mark.parametrize("rotary_mode", [0, 1])
|
|
def test_tree_attn(nbatch, h_q, h_kv, d, rotary_mode):
|
|
np.random.seed(0)
|
|
np.set_printoptions(linewidth=10000)
|
|
|
|
def gen_chain(num_nodes):
|
|
mask = np.tril(np.ones((num_nodes, num_nodes)))
|
|
return num_nodes, list(mask.flatten()), np.arange(num_nodes)
|
|
|
|
def gen_full_binary_tree(height):
|
|
mask = list()
|
|
pos = list()
|
|
num_nodes = 2**height - 1
|
|
for i in range(num_nodes):
|
|
if i == 0:
|
|
mask_0 = [0] * num_nodes
|
|
mask_0[0] = 1
|
|
mask.append(mask_0)
|
|
pos.append(0)
|
|
else:
|
|
mask_i = mask[(i + 1) // 2 - 1].copy()
|
|
mask_i[i] = 1
|
|
mask.append(mask_i)
|
|
pos.append(pos[(i + 1) // 2 - 1] + 1)
|
|
return num_nodes, list(np.array(mask).flatten()), pos
|
|
|
|
### Inputs
|
|
num_nodes = 0
|
|
m_list = list()
|
|
mn_list = list()
|
|
mask_list = list()
|
|
q_pos_list = list()
|
|
|
|
mn_list.append(0)
|
|
|
|
for _ in range(nbatch):
|
|
choice = np.random.choice(2, 1, p=[1, 0])
|
|
if choice == 0:
|
|
nodes_batch = np.random.randint(3, 32)
|
|
res = gen_chain(nodes_batch)
|
|
num_nodes += nodes_batch
|
|
else:
|
|
height = np.random.randint(2, 6)
|
|
res = gen_full_binary_tree(height)
|
|
num_nodes += 2**height - 1
|
|
m_list.append(res[0])
|
|
mn_list.append(res[0] ** 2)
|
|
mask_list.extend(res[1])
|
|
q_pos_list.extend(res[2])
|
|
|
|
qkv_indptr = np.array(np.cumsum([0, *m_list])).astype(np.int32)
|
|
m_list = np.array(m_list).astype(np.int32)
|
|
mn_list = np.array(mn_list).astype(np.int32)
|
|
mn_list = np.cumsum(mn_list).astype(np.int32)
|
|
mask_list = np.array(mask_list).astype(np.int32)
|
|
q_pos_list = np.array(q_pos_list).astype(np.int32)
|
|
|
|
# print("qkv_indptr:", qkv_indptr)
|
|
# print("m_list:", m_list)
|
|
# print("mn_list:", mn_list)
|
|
# for num_nodes, base in zip(m_list, mn_list):
|
|
# print("num_nodes:", num_nodes)
|
|
# print("indptr:", base)
|
|
# print(
|
|
# "mask:",
|
|
# mask_list[base : base + num_nodes * num_nodes].reshape(num_nodes, num_nodes),
|
|
# )
|
|
# print("q_pos:", q_pos_list[base : base + num_nodes])
|
|
|
|
q = np.random.rand(num_nodes, h_q, d).astype(np.float16)
|
|
q_indptr = qkv_indptr
|
|
k = np.random.rand(num_nodes, h_kv, d).astype(np.float16)
|
|
v = np.random.rand(num_nodes, h_kv, d).astype(np.float16)
|
|
kv_indptr = qkv_indptr
|
|
q_rope_position = q_pos_list
|
|
m_arr = m_list
|
|
mn_indptr = mn_list
|
|
mask = mask_list
|
|
output = np.zeros((num_nodes, h_q, d), dtype=np.float16)
|
|
lse = np.zeros((num_nodes, h_q), dtype=np.float32)
|
|
rotary_scale = 1.0
|
|
rotary_theta = 10000.0
|
|
attn_score_scaling_factor = 1.0
|
|
|
|
### TVM Inputs
|
|
dev = tvm.cuda(0)
|
|
q_tvm = tvm.runtime.tensor(q, dev)
|
|
q_indptr_tvm = tvm.runtime.tensor(q_indptr, dev)
|
|
k_tvm = tvm.runtime.tensor(k, dev)
|
|
v_tvm = tvm.runtime.tensor(v, dev)
|
|
kv_indptr_tvm = tvm.runtime.tensor(kv_indptr, dev)
|
|
q_rope_position_tvm = tvm.runtime.tensor(q_rope_position, dev)
|
|
# m_arr_tvm = tvm.runtime.tensor(m_arr, dev)
|
|
mn_indptr_tvm = tvm.runtime.tensor(mn_indptr, dev)
|
|
mask_tvm = tvm.runtime.tensor(mask, dev)
|
|
output_tvm = tvm.runtime.tensor(output, dev)
|
|
lse_tvm = tvm.runtime.tensor(lse, dev)
|
|
|
|
target = tvm.target.Target("cuda")
|
|
kernel = tree_attn(h_kv=h_kv, h_q=h_q, d=d, dtype="float16", rope_scaling={}, target=target)
|
|
mod = tvm.build(kernel, target=target)
|
|
mod(
|
|
q_tvm,
|
|
q_indptr_tvm,
|
|
k_tvm,
|
|
v_tvm,
|
|
kv_indptr_tvm,
|
|
q_rope_position_tvm,
|
|
# m_arr_tvm,
|
|
mn_indptr_tvm,
|
|
mask_tvm,
|
|
output_tvm,
|
|
lse_tvm,
|
|
rotary_mode,
|
|
rotary_scale,
|
|
rotary_theta,
|
|
attn_score_scaling_factor,
|
|
nbatch,
|
|
)
|
|
|
|
### Numpy reference
|
|
def numpy_reference(
|
|
q,
|
|
q_indptr,
|
|
k,
|
|
v,
|
|
kv_indptr,
|
|
q_rope_position,
|
|
m_arr,
|
|
mn_indptr,
|
|
mask,
|
|
rotary_mode,
|
|
rotary_scale,
|
|
rotary_theta,
|
|
attn_score_scaling_factor,
|
|
output_tvm,
|
|
):
|
|
def rope_freq(s, d, d_range, theta, dtype):
|
|
freq = s / math.pow(theta, (d * 2 % d_range) / float(d_range))
|
|
cos_freq = np.cos(freq).astype(dtype)
|
|
sin_freq = np.sin(freq).astype(dtype)
|
|
return cos_freq, sin_freq
|
|
|
|
def rope(buffer, offset, rotary_dim, theta, scale, dtype):
|
|
result = buffer.copy()
|
|
for pos, h, d in np.ndindex(buffer.shape):
|
|
cos_freq, sin_freq = rope_freq(offset[pos] * scale, d, rotary_dim, theta, dtype)
|
|
cos = cos_freq * buffer[pos, h, d]
|
|
sin = sin_freq * (
|
|
-buffer[pos, h, d + rotary_dim // 2]
|
|
if d < rotary_dim // 2
|
|
else buffer[pos, h, d - rotary_dim // 2]
|
|
)
|
|
result[pos, h, d] = cos + sin
|
|
return result
|
|
|
|
for i in range(len(m_arr)):
|
|
num_nodes = m_arr[i]
|
|
base = mn_indptr[i]
|
|
q_base = q_indptr[i]
|
|
kv_base = kv_indptr[i]
|
|
q_pos = q_rope_position[q_base : q_base + num_nodes] # (num_nodes,)
|
|
q_i = q[q_base : q_base + num_nodes] # (num_nodes, h_q, d)
|
|
k_i = k[kv_base : kv_base + num_nodes] # (num_nodes, h_kv, d)
|
|
v_i = v[kv_base : kv_base + num_nodes] # (num_nodes, h_kv, d)
|
|
mask_i = mask[base : base + num_nodes * num_nodes].reshape(num_nodes, num_nodes)
|
|
|
|
if rotary_mode == 1:
|
|
q_i = rope(q_i, q_pos, d, rotary_theta, rotary_scale, q_i.dtype)
|
|
k_i = rope(k_i, q_pos, d, rotary_theta, rotary_scale, k_i.dtype)
|
|
|
|
# group attention
|
|
# q: (num_nodes, h_q, d)
|
|
# k: (num_nodes, h_kv, d)
|
|
# v: (num_nodes, h_kv, d)
|
|
group_size = h_q // h_kv
|
|
q_reshape = q_i.transpose(1, 0, 2) # (h_q, num_nodes, d)
|
|
k_reshape = k_i.transpose(1, 2, 0) # (h_kv, d, num_nodes)
|
|
v_reshape = v_i.transpose(1, 0, 2) # (h_kv, num_nodes, d)
|
|
# expand k_reshape
|
|
k_reshape = k_reshape.reshape(h_kv, 1, d, num_nodes)
|
|
k_reshape = np.repeat(k_reshape, group_size, axis=1)
|
|
k_reshape = k_reshape.reshape(h_q, d, num_nodes)
|
|
# expand v_reshape
|
|
v_reshape = v_reshape.reshape(h_kv, 1, num_nodes, d)
|
|
v_reshape = np.repeat(v_reshape, group_size, axis=1)
|
|
v_reshape = v_reshape.reshape(h_q, num_nodes, d)
|
|
# print("q_reshape:", q_reshape.shape)
|
|
# print("k_reshape:", k_reshape.shape)
|
|
# print("v_reshape:", v_reshape.shape)
|
|
|
|
# qk: (h_q, num_nodes, num_nodes)
|
|
qk = np.matmul(q_reshape, k_reshape) * attn_score_scaling_factor / math.sqrt(float(d))
|
|
# softmax(qk, axis=-1), numerical stability
|
|
qk[:, mask_i == 0] = -np.inf
|
|
qk_max = np.max(qk, axis=-1, keepdims=True)
|
|
qk = np.exp(qk - qk_max)
|
|
qk = qk / np.sum(qk, axis=-1, keepdims=True)
|
|
|
|
# attention
|
|
output_i = np.matmul(qk, v_reshape).transpose(1, 0, 2) # (num_nodes, h_q, d)
|
|
# print(output_i)
|
|
|
|
tvm.testing.assert_allclose(
|
|
output_i, output_tvm[q_base : q_base + num_nodes], rtol=1e-3, atol=1e-3
|
|
)
|
|
|
|
numpy_reference(
|
|
q,
|
|
q_indptr,
|
|
k,
|
|
v,
|
|
kv_indptr,
|
|
q_rope_position,
|
|
m_arr,
|
|
mn_indptr,
|
|
mask,
|
|
rotary_mode,
|
|
rotary_scale,
|
|
rotary_theta,
|
|
attn_score_scaling_factor,
|
|
output_tvm.numpy(),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|