wehub-resource-sync
3197 lines
105 KiB
Python
3197 lines
105 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed 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.
|
|
import unittest
|
|
|
|
import numpy as np
|
|
from op_test import get_cuda_version, get_device_place, is_custom_device
|
|
|
|
import paddle
|
|
from paddle import base
|
|
from paddle.framework import core
|
|
from paddle.incubate.nn.functional import block_multihead_attention
|
|
from paddle.static import Program, program_guard
|
|
|
|
paddle.seed(2023)
|
|
np.random.seed(2023)
|
|
|
|
|
|
is_sm8x = (
|
|
(core.is_compiled_with_cuda() or is_custom_device())
|
|
and paddle.device.cuda.get_device_capability()[0] == 8
|
|
and paddle.device.cuda.get_device_capability()[1] >= 0
|
|
)
|
|
|
|
is_sm9x = (
|
|
(core.is_compiled_with_cuda() or is_custom_device())
|
|
and paddle.device.cuda.get_device_capability()[0] == 9
|
|
and paddle.device.cuda.get_device_capability()[1] >= 0
|
|
)
|
|
|
|
is_sm7x = (
|
|
(core.is_compiled_with_cuda() or is_custom_device())
|
|
and paddle.device.cuda.get_device_capability()[0] == 7
|
|
and paddle.device.cuda.get_device_capability()[1] >= 0
|
|
)
|
|
|
|
is_sm_supported = is_sm8x or is_sm9x or is_sm7x
|
|
|
|
|
|
def create_attn_mask(
|
|
mask_type,
|
|
batch_size,
|
|
seq_lens,
|
|
pre_cache_length=0,
|
|
):
|
|
max_seq_len = max(seq_lens)
|
|
mask = paddle.zeros(
|
|
[batch_size, 1, max_seq_len, max_seq_len + pre_cache_length],
|
|
dtype=mask_type,
|
|
)
|
|
mask[:, :, :, :pre_cache_length] = 1
|
|
for i in range(batch_size):
|
|
seq_len = seq_lens[i]
|
|
mask[i, 0, :seq_len, :seq_len] = (
|
|
paddle.tril(paddle.ones(shape=(seq_len, seq_len), dtype=mask_type))
|
|
- 1
|
|
) * 1e4
|
|
return mask
|
|
|
|
|
|
def naive_attention_impl(
|
|
query,
|
|
key,
|
|
value,
|
|
cache_k=None,
|
|
cache_v=None,
|
|
pre_cache_k=None,
|
|
pre_cache_v=None,
|
|
mask=None,
|
|
scale=1.0,
|
|
cache_k_dequant_scales=None,
|
|
cache_v_dequant_scales=None,
|
|
use_cachekv_int8="None",
|
|
):
|
|
batch = query.shape[0]
|
|
heads = query.shape[1]
|
|
seq_len = query.shape[2]
|
|
head_dim = query.shape[3]
|
|
kv_head = key.shape[1]
|
|
|
|
key = key.reshape([batch, kv_head, 1, seq_len, head_dim])
|
|
key = paddle.tile(key, [1, 1, heads // kv_head, 1, 1])
|
|
key = key.reshape([batch, heads, seq_len, head_dim])
|
|
|
|
if use_cachekv_int8 == "dynamic":
|
|
unsqueeze_shape = [2, 3]
|
|
elif use_cachekv_int8 == "static":
|
|
unsqueeze_shape = [0, 2, 3]
|
|
if pre_cache_k is not None:
|
|
key = paddle.concat([pre_cache_k, key], axis=2)
|
|
if cache_k is not None:
|
|
if cache_k_dequant_scales is not None:
|
|
dequant_cache_k = (
|
|
(cache_k.astype('float32') - 128.0)
|
|
* cache_k_dequant_scales.unsqueeze(unsqueeze_shape)
|
|
).astype(key.dtype)
|
|
key = paddle.concat([dequant_cache_k, key], axis=2)
|
|
else:
|
|
key = paddle.concat([cache_k, key], axis=2)
|
|
|
|
value = value.reshape([batch, kv_head, 1, seq_len, head_dim])
|
|
value = paddle.tile(value, [1, 1, heads // kv_head, 1, 1])
|
|
value = value.reshape([batch, heads, seq_len, head_dim])
|
|
if pre_cache_v is not None:
|
|
value = paddle.concat([pre_cache_v, value], axis=2)
|
|
if cache_v is not None:
|
|
if cache_v_dequant_scales is not None:
|
|
dequant_cache_v = (
|
|
(cache_v.astype('float32') - 128.0)
|
|
* cache_v_dequant_scales.unsqueeze(unsqueeze_shape)
|
|
).astype(value.dtype)
|
|
value = paddle.concat([dequant_cache_v, value], axis=2)
|
|
else:
|
|
value = paddle.concat([cache_v, value], axis=2)
|
|
|
|
qk_res = paddle.matmul(query, key, transpose_y=True)
|
|
attention = qk_res * scale
|
|
if mask is not None:
|
|
attention = attention + mask
|
|
softmax_result = paddle.nn.functional.softmax(attention, -1)
|
|
result = paddle.matmul(softmax_result, value)
|
|
return result
|
|
|
|
|
|
def get_padding_offset(bsz, max_seq_len, seq_lens_this_time):
|
|
cum_offsets_now = paddle.cumsum(max_seq_len - seq_lens_this_time)
|
|
cum_offsets = paddle.zeros(shape=(bsz + 1), dtype="int32")
|
|
cum_offsets[1:] = cum_offsets_now
|
|
token_num = paddle.sum(seq_lens_this_time)
|
|
padding_offsets = paddle.zeros(shape=(token_num), dtype="int32")
|
|
cu_seqlens_q = paddle.zeros(shape=(bsz + 1), dtype="int32")
|
|
cu_seqlens_k = paddle.zeros(shape=(bsz + 1), dtype="int32")
|
|
for i in range(bsz):
|
|
seq_len_now = seq_lens_this_time[i]
|
|
cum_offset = cum_offsets[i]
|
|
for j in range(seq_len_now):
|
|
padding_offsets[i * max_seq_len - cum_offset + j] = cum_offset
|
|
cum_seq_len = (i + 1) * max_seq_len - cum_offsets[i + 1]
|
|
cu_seqlens_q[i + 1] = cum_seq_len
|
|
cu_seqlens_k[i + 1] = cum_seq_len
|
|
return padding_offsets, cum_offsets[:-1], cu_seqlens_q, cu_seqlens_k
|
|
|
|
|
|
class RopeEmbedding:
|
|
def _rotary_position_embedding(self, seq_len, head_dim, dtype):
|
|
pos_seq = paddle.arange(0, seq_len, 1, dtype=dtype)
|
|
indices = paddle.arange(0, head_dim, 2, dtype=dtype)
|
|
indices = 1 / 10000 ** (indices / head_dim)
|
|
|
|
sinusoid_inp = pos_seq.unsqueeze(1) * indices.unsqueeze(0)
|
|
pos_emb = paddle.concat(
|
|
[paddle.sin(sinusoid_inp), paddle.cos(sinusoid_inp)], axis=-1
|
|
)
|
|
pos_emb = paddle.reshape(pos_emb, (1, 1, seq_len, head_dim))
|
|
pos_emb.stop_gradient = True
|
|
return pos_emb
|
|
|
|
def _apply_rope(self, rp, q, k, v=None):
|
|
# sin [sequence_length, embed_size_per_head//2]
|
|
# cos [sequence_length, embed_size_per_head//2]
|
|
sin, cos = paddle.chunk(rp, 2, axis=-1)
|
|
# sin [θ0,θ1,θ2......θd/2-1] -> sin_pos [θ0,θ0,θ1,θ1,θ2,θ2......θd/2-1,θd/2-1]
|
|
sin_pos = paddle.reshape(paddle.stack([sin, sin], axis=-1), rp.shape)
|
|
# cos [θ0,θ1,θ2......θd/2-1] -> cos_pos [θ0,θ0,θ1,θ1,θ2,θ2......θd/2-1,θd/2-1]
|
|
cos_pos = paddle.reshape(paddle.stack([cos, cos], axis=-1), rp.shape)
|
|
# rotate_half_query_layer [-q1,q0,-q3,q2......,-qd-1,qd-2]
|
|
rotate_half_q = paddle.reshape(
|
|
paddle.stack([-q[:, :, :, 1::2], q[:, :, :, 0::2]], axis=-1),
|
|
paddle.shape(q),
|
|
)
|
|
query = paddle.add(
|
|
paddle.multiply(q, cos_pos), paddle.multiply(rotate_half_q, sin_pos)
|
|
)
|
|
# rotate_half_key_layer [-k1,k0,-k3,k2......,-kd-1,kd-2]
|
|
rotate_half_k = paddle.reshape(
|
|
paddle.stack([-k[:, :, :, 1::2], k[:, :, :, 0::2]], axis=-1),
|
|
paddle.shape(k),
|
|
)
|
|
key = paddle.add(
|
|
paddle.multiply(k, cos_pos), paddle.multiply(rotate_half_k, sin_pos)
|
|
)
|
|
if v is not None:
|
|
# rotate_half_value_layer [-v1,v0,-v3,v2......,-vd-1,vd-2]
|
|
rotate_half_v = paddle.reshape(
|
|
paddle.stack([-v[:, :, :, 1::2], v[:, :, :, 0::2]], axis=-1),
|
|
paddle.shape(v),
|
|
)
|
|
value = paddle.add(
|
|
paddle.multiply(v, cos_pos),
|
|
paddle.multiply(rotate_half_v, sin_pos),
|
|
)
|
|
return query, key, value
|
|
return query, key
|
|
|
|
|
|
def remove_padding(seq_lens, cu_seq_lens, inputs, token_num):
|
|
bsz, num_head, seq_len, dim_head = inputs.shape
|
|
output = paddle.zeros(
|
|
shape=[token_num, num_head * dim_head], dtype=inputs.dtype
|
|
)
|
|
inputs = inputs.transpose([0, 2, 1, 3]).reshape([bsz, seq_len, -1])
|
|
for i in range(bsz):
|
|
seq_len_now = seq_lens[i]
|
|
start_idx = cu_seq_lens[i]
|
|
end_idx = cu_seq_lens[i + 1]
|
|
output[start_idx:end_idx, :] = inputs[i, :seq_len_now, :]
|
|
return output
|
|
|
|
|
|
def block_cache_to_naive_cache(
|
|
cache_k, cache_v, bsz, block_tables, cache_seq_len
|
|
):
|
|
_, num_head, blocksize, dim_head = cache_k.shape
|
|
out_cache_k = paddle.zeros(
|
|
shape=[bsz, num_head, cache_seq_len, dim_head], dtype=cache_k.dtype
|
|
)
|
|
out_cache_v = paddle.zeros(
|
|
shape=[bsz, num_head, cache_seq_len, dim_head], dtype=cache_v.dtype
|
|
)
|
|
for i in range(bsz):
|
|
for j in range(cache_seq_len):
|
|
out_cache_k[i, :, j, :] = cache_k[
|
|
block_tables[i, j // blocksize], :, j % blocksize, :
|
|
]
|
|
out_cache_v[i, :, j, :] = cache_v[
|
|
block_tables[i, j // blocksize], :, j % blocksize, :
|
|
]
|
|
return out_cache_k, out_cache_v
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11040
|
|
or not is_sm_supported,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.4"
|
|
"and device's compute capability must be 8.x or 90",
|
|
)
|
|
class TestBlockMultiHeadAttnEncDec(unittest.TestCase):
|
|
def setUp(self):
|
|
paddle.disable_static()
|
|
self.name = "TestBlockMultiHeadAttnEncDec"
|
|
self.place = get_device_place()
|
|
self.batch_size = 2
|
|
self.num_head = 8
|
|
self.seq_len = 64
|
|
self.max_dec_len = 64
|
|
self.dim_head = 64
|
|
self.hid_dim = self.num_head * self.dim_head
|
|
self.blocksize = 64
|
|
self.block_num_per_seq = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
self.max_block_num = self.block_num_per_seq * self.batch_size
|
|
self.free_list = list(range(self.max_block_num - 1, -1, -1))
|
|
self.seq_lens_encoder = paddle.to_tensor(
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_decoder = paddle.to_tensor(
|
|
[
|
|
0,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_this_time = self.seq_lens_encoder
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
self.seq_len,
|
|
self.dim_head,
|
|
)
|
|
self.cache_shape = (
|
|
self.max_block_num,
|
|
self.num_head,
|
|
self.blocksize,
|
|
self.dim_head,
|
|
)
|
|
self.dtype = 'float16'
|
|
self.attention_mask = create_attn_mask(
|
|
self.dtype,
|
|
self.batch_size,
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
)
|
|
|
|
self.tgt_mask = paddle.randn(
|
|
[self.batch_size, self.num_head, 1, self.seq_len + 1],
|
|
dtype=self.dtype,
|
|
)
|
|
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.cache_k = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.cache_v = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.block_tables = paddle.zeros(
|
|
shape=(self.batch_size, self.block_num_per_seq), dtype="int32"
|
|
)
|
|
for i in range(self.batch_size):
|
|
need_block_num = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
for j in range(need_block_num):
|
|
self.block_tables[i, j] = self.free_list.pop()
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(
|
|
self.batch_size, self.seq_len, self.seq_lens_this_time
|
|
)
|
|
self.token_num = self.padding_offset.shape[0]
|
|
|
|
def test_all(self):
|
|
paddle.disable_static()
|
|
# encoder
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.token_num, -1])
|
|
out_ = naive_attention_impl(
|
|
q, k, v, None, None, None, None, self.attention_mask, self.scale
|
|
)
|
|
out_ = remove_padding(
|
|
self.seq_lens_this_time, self.cu_seqlens_q, out_, self.token_num
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
self.seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style,
|
|
)[0]
|
|
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=5e-03,
|
|
atol=1e-03,
|
|
)
|
|
|
|
# decoder
|
|
naive_cache_k, naive_cache_v = block_cache_to_naive_cache(
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.batch_size,
|
|
self.block_tables,
|
|
self.seq_len,
|
|
)
|
|
|
|
self.seq_lens_decoder[:] = self.seq_lens_encoder
|
|
self.seq_lens_encoder[:] = 0
|
|
self.seq_lens_this_time[:] = 1
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
1,
|
|
self.dim_head,
|
|
)
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.batch_size, -1])
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(self.batch_size, 1, self.seq_lens_this_time)
|
|
|
|
out_ = (
|
|
naive_attention_impl(
|
|
q,
|
|
k,
|
|
v,
|
|
naive_cache_k,
|
|
naive_cache_v,
|
|
None,
|
|
None,
|
|
self.tgt_mask,
|
|
self.scale,
|
|
)
|
|
.transpose([0, 2, 1, 3])
|
|
.reshape([self.batch_size, -1])
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
self.tgt_mask, # tgt_mask
|
|
1, # seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
)[0]
|
|
# NOTE: The diff of decoder is a little big
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=5e-02,
|
|
atol=5e-02,
|
|
)
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11040
|
|
or not is_sm_supported,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.4"
|
|
"and device's compute capability must be 8.x or 90",
|
|
)
|
|
class TestBlockMultiHeadAttnEncDecSkipGetMaxLen(unittest.TestCase):
|
|
def setUp(self):
|
|
paddle.disable_static()
|
|
self.name = "TestBlockMultiHeadAttnEncDecSkipGetMaxLen"
|
|
self.place = get_device_place()
|
|
self.batch_size = 2
|
|
self.num_head = 8
|
|
self.seq_len = 64
|
|
self.max_dec_len = 64
|
|
self.dim_head = 64
|
|
self.hid_dim = self.num_head * self.dim_head
|
|
self.blocksize = 64
|
|
self.block_num_per_seq = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
self.max_block_num = self.block_num_per_seq * self.batch_size
|
|
self.free_list = list(range(self.max_block_num - 1, -1, -1))
|
|
self.seq_lens_encoder = paddle.to_tensor(
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_decoder = paddle.to_tensor(
|
|
[
|
|
0,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_this_time = self.seq_lens_encoder
|
|
self.max_enc_len_this_time = paddle.to_tensor(
|
|
[self.seq_len], "int32"
|
|
).cpu()
|
|
self.max_dec_len_this_time = paddle.to_tensor([0], "int32").cpu()
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
self.seq_len,
|
|
self.dim_head,
|
|
)
|
|
self.cache_shape = (
|
|
self.max_block_num,
|
|
self.num_head,
|
|
self.blocksize,
|
|
self.dim_head,
|
|
)
|
|
self.dtype = 'float16'
|
|
self.attention_mask = create_attn_mask(
|
|
self.dtype,
|
|
self.batch_size,
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
)
|
|
|
|
self.tgt_mask = paddle.randn(
|
|
[self.batch_size, self.num_head, 1, self.seq_len + 1],
|
|
dtype=self.dtype,
|
|
)
|
|
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.cache_k = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.cache_v = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.block_tables = paddle.zeros(
|
|
shape=(self.batch_size, self.block_num_per_seq), dtype="int32"
|
|
)
|
|
for i in range(self.batch_size):
|
|
need_block_num = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
for j in range(need_block_num):
|
|
self.block_tables[i, j] = self.free_list.pop()
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(
|
|
self.batch_size, self.seq_len, self.seq_lens_this_time
|
|
)
|
|
self.token_num = self.padding_offset.shape[0]
|
|
|
|
def test_all(self):
|
|
paddle.disable_static()
|
|
# encoder
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.token_num, -1])
|
|
out_ = naive_attention_impl(
|
|
q, k, v, None, None, None, None, self.attention_mask, self.scale
|
|
)
|
|
out_ = remove_padding(
|
|
self.seq_lens_this_time, self.cu_seqlens_q, out_, self.token_num
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
self.max_enc_len_this_time, # max_enc_len_this_time
|
|
self.max_dec_len_this_time, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
self.seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style,
|
|
)[0]
|
|
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=5e-03,
|
|
atol=1e-03,
|
|
)
|
|
|
|
# decoder
|
|
naive_cache_k, naive_cache_v = block_cache_to_naive_cache(
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.batch_size,
|
|
self.block_tables,
|
|
self.seq_len,
|
|
)
|
|
|
|
self.seq_lens_decoder[:] = self.seq_lens_encoder
|
|
self.seq_lens_encoder[:] = 0
|
|
self.seq_lens_this_time[:] = 1
|
|
self.max_enc_len_this_time = paddle.to_tensor([0], "int32").cpu()
|
|
self.max_dec_len_this_time = paddle.to_tensor(
|
|
[self.seq_len], "int32"
|
|
).cpu()
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
1,
|
|
self.dim_head,
|
|
)
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.batch_size, -1])
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(self.batch_size, 1, self.seq_lens_this_time)
|
|
|
|
out_ = (
|
|
naive_attention_impl(
|
|
q,
|
|
k,
|
|
v,
|
|
naive_cache_k,
|
|
naive_cache_v,
|
|
None,
|
|
None,
|
|
self.tgt_mask,
|
|
self.scale,
|
|
)
|
|
.transpose([0, 2, 1, 3])
|
|
.reshape([self.batch_size, -1])
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
self.max_enc_len_this_time, # max_enc_len_this_time
|
|
self.max_dec_len_this_time, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
self.tgt_mask, # tgt_mask
|
|
1, # seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
)[0]
|
|
# NOTE: The diff of decoder is a little big
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=5e-02,
|
|
atol=5e-02,
|
|
)
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11040
|
|
or not is_sm_supported,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.4"
|
|
"and device's compute capability must be 8.x or 90",
|
|
)
|
|
class TestBlockMultiHeadAttnRoPE(unittest.TestCase):
|
|
def setUp(self):
|
|
paddle.disable_static()
|
|
self.name = "TestBlockMultiHeadAttnRoPE"
|
|
self.place = get_device_place()
|
|
self.batch_size = 2
|
|
self.num_head = 8
|
|
self.seq_len = 64
|
|
self.max_dec_len = 64
|
|
self.dim_head = 64
|
|
self.hid_dim = self.num_head * self.dim_head
|
|
self.blocksize = 64
|
|
self.block_num_per_seq = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
self.rope = RopeEmbedding()
|
|
self.max_block_num = self.block_num_per_seq * self.batch_size
|
|
self.free_list = list(range(self.max_block_num - 1, -1, -1))
|
|
self.seq_lens_encoder = paddle.to_tensor(
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_decoder = paddle.to_tensor(
|
|
[
|
|
0,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_this_time = self.seq_lens_encoder
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
self.seq_len,
|
|
self.dim_head,
|
|
)
|
|
self.cache_shape = (
|
|
self.max_block_num,
|
|
self.num_head,
|
|
self.blocksize,
|
|
self.dim_head,
|
|
)
|
|
self.dtype = 'float16'
|
|
self.attention_mask = create_attn_mask(
|
|
self.dtype,
|
|
self.batch_size,
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
)
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.cache_k = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.cache_v = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.block_tables = paddle.zeros(
|
|
shape=(self.batch_size, self.block_num_per_seq), dtype="int32"
|
|
)
|
|
for i in range(self.batch_size):
|
|
need_block_num = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
for j in range(need_block_num):
|
|
self.block_tables[i, j] = self.free_list.pop()
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(
|
|
self.batch_size, self.seq_len, self.seq_lens_this_time
|
|
)
|
|
self.token_num = self.padding_offset.shape[0]
|
|
|
|
def get_rotary_position_embedding(self, position_ids, head_dim):
|
|
bsz, max_seq_len = position_ids.shape[:2]
|
|
rot_emb = paddle.zeros(
|
|
(2, bsz, max_seq_len, 1, head_dim // 2), dtype="float32"
|
|
)
|
|
inv_freq = 10000 ** (
|
|
-paddle.arange(0, head_dim, 2, dtype="float32") / head_dim
|
|
)
|
|
|
|
# shape: [B, S, D/2]
|
|
freqs = paddle.einsum(
|
|
"ij,k->ijk", position_ids.cast("float32"), inv_freq
|
|
)
|
|
# shape: [B, S, D]
|
|
# emb = paddle.stack([freqs, freqs], axis=-1).reshape((bsz, max_seq_len, head_dim))
|
|
emb = paddle.stack([freqs], axis=-1).reshape(
|
|
(bsz, max_seq_len, head_dim // 2)
|
|
)
|
|
# shape: [B, S, 1, D]
|
|
emb = paddle.unsqueeze(emb, 2)
|
|
|
|
rot_emb[0] = paddle.cos(emb)
|
|
rot_emb[1] = paddle.sin(emb)
|
|
return rot_emb
|
|
|
|
def test_all(self):
|
|
paddle.disable_static()
|
|
tmp_position_ids = paddle.arange(
|
|
self.seq_len + self.max_dec_len
|
|
).reshape((1, -1))
|
|
self.rope_emb = self.get_rotary_position_embedding(
|
|
tmp_position_ids, self.dim_head
|
|
)
|
|
# encoder
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.token_num, -1])
|
|
|
|
sinusoidal_pos = self.rope._rotary_position_embedding(
|
|
self.seq_len, self.dim_head, self.dtype
|
|
)
|
|
q, k = self.rope._apply_rope(sinusoidal_pos, q, k)
|
|
|
|
out_ = naive_attention_impl(
|
|
q, k, v, None, None, None, None, self.attention_mask, self.scale
|
|
)
|
|
out_ = remove_padding(
|
|
self.seq_lens_this_time, self.cu_seqlens_q, out_, self.token_num
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
self.rope_emb, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
self.seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
)[0]
|
|
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=5e-03,
|
|
atol=1e-03,
|
|
)
|
|
|
|
# decoder
|
|
naive_cache_k, naive_cache_v = block_cache_to_naive_cache(
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.batch_size,
|
|
self.block_tables,
|
|
self.seq_len,
|
|
)
|
|
|
|
self.seq_lens_decoder[:] = self.seq_lens_encoder
|
|
self.seq_lens_encoder[:] = 0
|
|
self.seq_lens_this_time[:] = 1
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
1,
|
|
self.dim_head,
|
|
)
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.batch_size, -1])
|
|
|
|
sinusoidal_pos = self.rope._rotary_position_embedding(
|
|
self.seq_len + 1, self.dim_head, self.dtype
|
|
)[:, :, -1:, :]
|
|
q, k = self.rope._apply_rope(sinusoidal_pos, q, k)
|
|
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(self.batch_size, 1, self.seq_lens_this_time)
|
|
|
|
out_ = (
|
|
naive_attention_impl(
|
|
q,
|
|
k,
|
|
v,
|
|
naive_cache_k,
|
|
naive_cache_v,
|
|
None,
|
|
None,
|
|
None,
|
|
self.scale,
|
|
)
|
|
.transpose([0, 2, 1, 3])
|
|
.reshape([self.batch_size, -1])
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
self.rope_emb, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
1, # seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
)[0]
|
|
# NOTE: The diff of decoder is a little big
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=5e-02,
|
|
atol=5e-02,
|
|
)
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11040
|
|
or not is_sm_supported,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.4"
|
|
"and device's compute capability must be 8.x or 90",
|
|
)
|
|
class TestBlockMultiHeadAttnPreCache(unittest.TestCase):
|
|
def setUp(self):
|
|
paddle.disable_static()
|
|
self.name = "TestBlockMultiHeadAttnPreCacbe"
|
|
self.place = get_device_place()
|
|
self.batch_size = 2
|
|
self.num_head = 8
|
|
self.seq_len = 64
|
|
self.max_dec_len = 64
|
|
self.dim_head = 64
|
|
self.hid_dim = self.num_head * self.dim_head
|
|
self.blocksize = 64
|
|
self.pre_cache_length = 64
|
|
self.max_seq_len = self.seq_len + self.pre_cache_length
|
|
self.block_num_per_seq = (
|
|
self.max_seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
self.max_block_num = self.block_num_per_seq * self.batch_size
|
|
self.free_list = list(range(self.max_block_num - 1, -1, -1))
|
|
self.seq_lens_encoder = paddle.to_tensor(
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_decoder = paddle.to_tensor(
|
|
[
|
|
0,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_this_time = self.seq_lens_encoder
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
self.seq_len,
|
|
self.dim_head,
|
|
)
|
|
self.cache_shape = (
|
|
self.max_block_num,
|
|
self.num_head,
|
|
self.blocksize,
|
|
self.dim_head,
|
|
)
|
|
self.pre_cache_shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
self.pre_cache_length,
|
|
self.dim_head,
|
|
)
|
|
self.dtype = 'float16'
|
|
self.attention_mask = create_attn_mask(
|
|
self.dtype,
|
|
self.batch_size,
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
self.pre_cache_length,
|
|
)
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.cache_k = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.cache_v = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.pre_cache_k = paddle.randn(
|
|
shape=self.pre_cache_shape, dtype=self.dtype
|
|
)
|
|
self.pre_cache_v = paddle.randn(
|
|
shape=self.pre_cache_shape, dtype=self.dtype
|
|
)
|
|
self.block_tables = paddle.zeros(
|
|
shape=(self.batch_size, self.block_num_per_seq), dtype="int32"
|
|
)
|
|
for i in range(self.batch_size):
|
|
need_block_num = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
for j in range(need_block_num):
|
|
self.block_tables[i, j] = self.free_list.pop()
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(
|
|
self.batch_size, self.seq_len, self.seq_lens_this_time
|
|
)
|
|
self.token_num = self.padding_offset.shape[0]
|
|
|
|
def test_all(self):
|
|
paddle.disable_static()
|
|
# encoder
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.token_num, -1])
|
|
out_ = naive_attention_impl(
|
|
q,
|
|
k,
|
|
v,
|
|
None,
|
|
None,
|
|
self.pre_cache_k,
|
|
self.pre_cache_v,
|
|
self.attention_mask,
|
|
self.scale,
|
|
)
|
|
out_ = remove_padding(
|
|
self.seq_lens_this_time, self.cu_seqlens_q, out_, self.token_num
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
self.pre_cache_k, # pre_key_cache
|
|
self.pre_cache_v, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
self.attention_mask, # attn_mask
|
|
None, # tgt_mask
|
|
self.seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
)[0]
|
|
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=5e-03,
|
|
atol=1e-03,
|
|
)
|
|
|
|
# decoder
|
|
naive_cache_k, naive_cache_v = block_cache_to_naive_cache(
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.batch_size,
|
|
self.block_tables,
|
|
self.seq_len + self.pre_cache_length,
|
|
)
|
|
|
|
self.seq_lens_decoder[:] = self.seq_lens_encoder
|
|
self.seq_lens_encoder[:] = 0
|
|
self.seq_lens_this_time[:] = 1
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
1,
|
|
self.dim_head,
|
|
)
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.batch_size, -1])
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(self.batch_size, 1, self.seq_lens_this_time)
|
|
|
|
out_ = (
|
|
naive_attention_impl(
|
|
q,
|
|
k,
|
|
v,
|
|
naive_cache_k,
|
|
naive_cache_v,
|
|
None,
|
|
None,
|
|
None,
|
|
self.scale,
|
|
)
|
|
.transpose([0, 2, 1, 3])
|
|
.reshape([self.batch_size, -1])
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
self.pre_cache_k, # pre_key_cache
|
|
self.pre_cache_v, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
self.attention_mask, # attn_mask
|
|
None, # tgt_mask
|
|
1, # seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
)[0]
|
|
# NOTE: The diff of decoder is a little big
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=5e-02,
|
|
atol=5e-02,
|
|
)
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11040
|
|
or not is_sm_supported,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.4"
|
|
"and device's compute capability must be 8.x or 90",
|
|
)
|
|
class TestBlockMultiHeadAttnEncStatic(unittest.TestCase):
|
|
def setUp(self):
|
|
paddle.disable_static()
|
|
self.name = "TestBlockMultiHeadAttnEncStatic"
|
|
self.place = get_device_place()
|
|
self.batch_size = 2
|
|
self.num_head = 8
|
|
self.seq_len = 64
|
|
self.max_dec_len = 64
|
|
self.dim_head = 64
|
|
self.hid_dim = self.num_head * self.dim_head
|
|
self.blocksize = 64
|
|
self.block_num_per_seq = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
self.max_block_num = self.block_num_per_seq * self.batch_size
|
|
self.free_list = list(range(self.max_block_num - 1, -1, -1))
|
|
self.seq_lens_encoder = paddle.to_tensor(
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_decoder = paddle.to_tensor(
|
|
[
|
|
0,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_this_time = self.seq_lens_encoder
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
self.seq_len,
|
|
self.dim_head,
|
|
)
|
|
self.cache_shape = (
|
|
self.max_block_num,
|
|
self.num_head,
|
|
self.blocksize,
|
|
self.dim_head,
|
|
)
|
|
self.dtype = 'float16'
|
|
self.attention_mask = create_attn_mask(
|
|
self.dtype,
|
|
self.batch_size,
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
)
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.cache_k = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.cache_v = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.block_tables = paddle.zeros(
|
|
shape=(self.batch_size, self.block_num_per_seq), dtype="int32"
|
|
)
|
|
for i in range(self.batch_size):
|
|
need_block_num = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
for j in range(need_block_num):
|
|
self.block_tables[i, j] = self.free_list.pop()
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(
|
|
self.batch_size, self.seq_len, self.seq_lens_this_time
|
|
)
|
|
self.token_num = self.padding_offset.shape[0]
|
|
|
|
def test_all(self):
|
|
paddle.disable_static()
|
|
# encoder
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
out_ = naive_attention_impl(
|
|
q, k, v, None, None, None, None, self.attention_mask, self.scale
|
|
)
|
|
out_ = remove_padding(
|
|
self.seq_lens_this_time, self.cu_seqlens_q, out_, self.token_num
|
|
)
|
|
|
|
qkv_numpy = (
|
|
paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
)
|
|
.reshape([self.token_num, -1])
|
|
.numpy()
|
|
)
|
|
paddle.enable_static()
|
|
with program_guard(Program(), Program()):
|
|
qkv = paddle.static.data(
|
|
name="qkv",
|
|
shape=(self.token_num, 3 * self.hid_dim),
|
|
dtype=self.dtype,
|
|
)
|
|
cache_k = paddle.static.data(
|
|
name="cache_k", shape=self.cache_shape, dtype=self.dtype
|
|
)
|
|
cache_v = paddle.static.data(
|
|
name="cache_v", shape=self.cache_shape, dtype=self.dtype
|
|
)
|
|
seq_lens_encoder = paddle.static.data(
|
|
name="seq_lens_encoder", shape=(self.batch_size,), dtype='int32'
|
|
)
|
|
seq_lens_decoder = paddle.static.data(
|
|
name="seq_lens_decoder", shape=(self.batch_size,), dtype='int32'
|
|
)
|
|
seq_lens_this_time = paddle.static.data(
|
|
name="seq_lens_this_time",
|
|
shape=(self.batch_size,),
|
|
dtype='int32',
|
|
)
|
|
cu_seqlens_q = paddle.static.data(
|
|
name="cu_seqlens_q", shape=(self.batch_size + 1,), dtype='int32'
|
|
)
|
|
cu_seqlens_k = paddle.static.data(
|
|
name="cu_seqlens_k", shape=(self.batch_size + 1,), dtype='int32'
|
|
)
|
|
padding_offsets = paddle.static.data(
|
|
name="padding_offset", shape=(self.token_num,), dtype='int32'
|
|
)
|
|
cum_offsets = paddle.static.data(
|
|
name="cum_offset", shape=(self.batch_size,), dtype='int32'
|
|
)
|
|
block_tables = paddle.static.data(
|
|
name="block_tables",
|
|
shape=(self.batch_size, self.block_num_per_seq),
|
|
dtype='int32',
|
|
)
|
|
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
cache_k,
|
|
cache_v,
|
|
seq_lens_encoder,
|
|
seq_lens_decoder,
|
|
seq_lens_this_time,
|
|
padding_offsets,
|
|
cum_offsets,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
self.seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
)
|
|
exe = base.Executor()
|
|
res = exe.run(
|
|
feed={
|
|
"qkv": qkv_numpy,
|
|
"cache_k": self.cache_k.numpy(),
|
|
"cache_v": self.cache_v.numpy(),
|
|
"seq_lens_encoder": self.seq_lens_encoder.numpy(),
|
|
"seq_lens_decoder": self.seq_lens_decoder.numpy(),
|
|
"seq_lens_this_time": self.seq_lens_this_time.numpy(),
|
|
"cu_seqlens_q": self.cu_seqlens_q.numpy(),
|
|
"cu_seqlens_k": self.cu_seqlens_k.numpy(),
|
|
"padding_offset": self.padding_offset.numpy(),
|
|
"cum_offset": self.cum_offset.numpy(),
|
|
"block_tables": self.block_tables.numpy(),
|
|
},
|
|
fetch_list=[out],
|
|
)
|
|
paddle.disable_static()
|
|
np.testing.assert_allclose(
|
|
res[0],
|
|
out_.numpy(),
|
|
rtol=5e-03,
|
|
atol=1e-03,
|
|
)
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11040
|
|
or not is_sm_supported,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.4"
|
|
"and device's compute capability must be 8.x or 90",
|
|
)
|
|
class TestBlockMultiHeadAttnEncDecPTQDequant(unittest.TestCase):
|
|
def setUp(self):
|
|
paddle.disable_static()
|
|
self.name = "TestBlockMultiHeadAttnEncDec"
|
|
self.place = get_device_place()
|
|
self.batch_size = 2
|
|
self.num_head = 8
|
|
self.seq_len = 64
|
|
self.max_dec_len = 64
|
|
self.dim_head = 64
|
|
self.hid_dim = self.num_head * self.dim_head
|
|
self.blocksize = 64
|
|
self.block_num_per_seq = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
self.max_block_num = self.block_num_per_seq * self.batch_size
|
|
self.free_list = list(range(self.max_block_num - 1, -1, -1))
|
|
self.seq_lens_encoder = paddle.to_tensor(
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_decoder = paddle.to_tensor(
|
|
[
|
|
0,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_this_time = self.seq_lens_encoder
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
self.seq_len,
|
|
self.dim_head,
|
|
)
|
|
self.cache_shape = (
|
|
self.max_block_num,
|
|
self.num_head,
|
|
self.blocksize,
|
|
self.dim_head,
|
|
)
|
|
self.dtype = 'float16'
|
|
self.attention_mask = create_attn_mask(
|
|
self.dtype,
|
|
self.batch_size,
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
)
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.cache_k = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.cache_v = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.block_tables = paddle.zeros(
|
|
shape=(self.batch_size, self.block_num_per_seq), dtype="int32"
|
|
)
|
|
for i in range(self.batch_size):
|
|
need_block_num = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
for j in range(need_block_num):
|
|
self.block_tables[i, j] = self.free_list.pop()
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(
|
|
self.batch_size, self.seq_len, self.seq_lens_this_time
|
|
)
|
|
self.token_num = self.padding_offset.shape[0]
|
|
|
|
def test_all(self):
|
|
paddle.disable_static()
|
|
# encoder
|
|
query = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
key = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
value = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.token_num, -1])
|
|
|
|
q = q.transpose([0, 2, 1, 3]).reshape([self.token_num, self.hid_dim])
|
|
k = k.transpose([0, 2, 1, 3]).reshape([self.token_num, self.hid_dim])
|
|
v = v.transpose([0, 2, 1, 3]).reshape([self.token_num, self.hid_dim])
|
|
|
|
q_out_scale = 10.0 / paddle.max(q, axis=0).astype('float32')
|
|
k_out_scale = 10.0 / paddle.max(k, axis=0).astype('float32')
|
|
v_out_scale = 10.0 / paddle.max(v, axis=0).astype('float32')
|
|
|
|
qkv_out_scale = paddle.concat(
|
|
[q_out_scale, k_out_scale, v_out_scale], axis=0
|
|
)
|
|
|
|
q_bias = paddle.ones([self.hid_dim], dtype=self.dtype)
|
|
k_bias = paddle.ones([self.hid_dim], dtype=self.dtype)
|
|
v_bias = paddle.ones([self.hid_dim], dtype=self.dtype)
|
|
|
|
qkv_bias = paddle.concat([q_bias, k_bias, v_bias], axis=-1)
|
|
|
|
# dequant
|
|
q = (q.astype('float32') * q_out_scale).astype(self.dtype)
|
|
k = (k.astype('float32') * k_out_scale).astype(self.dtype)
|
|
v = (v.astype('float32') * v_out_scale).astype(self.dtype)
|
|
|
|
# add bias
|
|
q = q + q_bias
|
|
k = k + k_bias
|
|
v = v + v_bias
|
|
|
|
# transpose to origin
|
|
q = q.reshape(
|
|
[self.batch_size, self.seq_len, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
k = k.reshape(
|
|
[self.batch_size, self.seq_len, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
v = v.reshape(
|
|
[self.batch_size, self.seq_len, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
|
|
out_ = naive_attention_impl(
|
|
q, k, v, None, None, None, None, self.attention_mask, self.scale
|
|
)
|
|
out_ = remove_padding(
|
|
self.seq_lens_this_time, self.cu_seqlens_q, out_, self.token_num
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
qkv_out_scale, # qkv_out_scale
|
|
qkv_bias, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
self.seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style,
|
|
compute_dtype="fp16",
|
|
)[0]
|
|
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=100,
|
|
atol=1,
|
|
)
|
|
|
|
# decoder
|
|
naive_cache_k, naive_cache_v = block_cache_to_naive_cache(
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.batch_size,
|
|
self.block_tables,
|
|
self.seq_len,
|
|
)
|
|
|
|
self.seq_lens_decoder[:] = self.seq_lens_encoder
|
|
self.seq_lens_encoder[:] = 0
|
|
self.seq_lens_this_time[:] = 1
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
1,
|
|
self.dim_head,
|
|
)
|
|
query = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
key = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
value = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.batch_size, -1])
|
|
|
|
q = q.transpose([0, 2, 1, 3]).reshape([self.batch_size, self.hid_dim])
|
|
k = k.transpose([0, 2, 1, 3]).reshape([self.batch_size, self.hid_dim])
|
|
v = v.transpose([0, 2, 1, 3]).reshape([self.batch_size, self.hid_dim])
|
|
|
|
q_out_scale = 1.0 / paddle.max(q, axis=0).astype('float32')
|
|
k_out_scale = 1.0 / paddle.max(k, axis=0).astype('float32')
|
|
v_out_scale = 1.0 / paddle.max(v, axis=0).astype('float32')
|
|
|
|
qkv_out_scale = paddle.concat(
|
|
[q_out_scale, k_out_scale, v_out_scale], axis=0
|
|
)
|
|
|
|
q_bias = paddle.ones([self.hid_dim], dtype=self.dtype) * 0.1
|
|
k_bias = paddle.ones([self.hid_dim], dtype=self.dtype) * 0.1
|
|
v_bias = paddle.ones([self.hid_dim], dtype=self.dtype) * 0.1
|
|
|
|
qkv_bias = paddle.concat([q_bias, k_bias, v_bias], axis=-1)
|
|
|
|
# dequant
|
|
q = (q.astype('float32') * q_out_scale).astype(self.dtype)
|
|
k = (k.astype('float32') * k_out_scale).astype(self.dtype)
|
|
v = (v.astype('float32') * v_out_scale).astype(self.dtype)
|
|
|
|
# add bias
|
|
q = q + q_bias
|
|
k = k + k_bias
|
|
v = v + v_bias
|
|
|
|
# transpose to origin
|
|
q = q.reshape(
|
|
[self.batch_size, 1, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
k = k.reshape(
|
|
[self.batch_size, 1, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
v = v.reshape(
|
|
[self.batch_size, 1, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(self.batch_size, 1, self.seq_lens_this_time)
|
|
|
|
out_ = (
|
|
naive_attention_impl(
|
|
q,
|
|
k,
|
|
v,
|
|
naive_cache_k,
|
|
naive_cache_v,
|
|
None,
|
|
None,
|
|
None,
|
|
self.scale,
|
|
)
|
|
.transpose([0, 2, 1, 3])
|
|
.reshape([self.batch_size, -1])
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
qkv_out_scale, # qkv_out_scale
|
|
qkv_bias, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
1, # seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
compute_dtype="fp16",
|
|
)[0]
|
|
# NOTE: The diff of decoder is a little big
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=100,
|
|
atol=1,
|
|
)
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11040
|
|
or not is_sm_supported,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.4"
|
|
" and device's compute capability must be 7.x, 8.x or 9.x",
|
|
)
|
|
class TestBlockMultiHeadAttnEncDecPTQDequantQuantShiftSmooth(unittest.TestCase):
|
|
def setUp(self):
|
|
paddle.disable_static()
|
|
self.name = "TestBlockMultiHeadAttnEncDec"
|
|
self.place = get_device_place()
|
|
self.batch_size = 2
|
|
self.num_head = 8
|
|
self.seq_len = 64
|
|
self.max_dec_len = 64
|
|
self.dim_head = 64
|
|
self.hid_dim = self.num_head * self.dim_head
|
|
self.blocksize = 64
|
|
self.block_num_per_seq = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
self.max_block_num = self.block_num_per_seq * self.batch_size
|
|
self.free_list = list(range(self.max_block_num - 1, -1, -1))
|
|
self.seq_lens_encoder = paddle.to_tensor(
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_decoder = paddle.to_tensor(
|
|
[
|
|
0,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_this_time = self.seq_lens_encoder
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
self.seq_len,
|
|
self.dim_head,
|
|
)
|
|
self.cache_shape = (
|
|
self.max_block_num,
|
|
self.num_head,
|
|
self.blocksize,
|
|
self.dim_head,
|
|
)
|
|
self.dtype = 'float16'
|
|
self.attention_mask = create_attn_mask(
|
|
self.dtype,
|
|
self.batch_size,
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
)
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.cache_k = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.cache_v = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.block_tables = paddle.zeros(
|
|
shape=(self.batch_size, self.block_num_per_seq), dtype="int32"
|
|
)
|
|
for i in range(self.batch_size):
|
|
need_block_num = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
for j in range(need_block_num):
|
|
self.block_tables[i, j] = self.free_list.pop()
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(
|
|
self.batch_size, self.seq_len, self.seq_lens_this_time
|
|
)
|
|
self.token_num = self.padding_offset.shape[0]
|
|
|
|
def test_all(self):
|
|
paddle.disable_static()
|
|
# encoder
|
|
query = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
key = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
value = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.token_num, -1])
|
|
|
|
q = q.transpose([0, 2, 1, 3]).reshape([self.token_num, self.hid_dim])
|
|
k = k.transpose([0, 2, 1, 3]).reshape([self.token_num, self.hid_dim])
|
|
v = v.transpose([0, 2, 1, 3]).reshape([self.token_num, self.hid_dim])
|
|
|
|
q_out_scale = 1.0 / paddle.max(q, axis=0).astype('float32')
|
|
k_out_scale = 1.0 / paddle.max(k, axis=0).astype('float32')
|
|
v_out_scale = 1.0 / paddle.max(v, axis=0).astype('float32')
|
|
|
|
qkv_out_scale = paddle.concat(
|
|
[q_out_scale, k_out_scale, v_out_scale], axis=0
|
|
)
|
|
|
|
q_bias = paddle.ones([self.hid_dim], dtype=self.dtype)
|
|
k_bias = paddle.ones([self.hid_dim], dtype=self.dtype)
|
|
v_bias = paddle.ones([self.hid_dim], dtype=self.dtype)
|
|
|
|
qkv_bias = paddle.concat([q_bias, k_bias, v_bias], axis=-1)
|
|
|
|
# dequant
|
|
q = (q.astype('float32') * q_out_scale).astype(self.dtype)
|
|
k = (k.astype('float32') * k_out_scale).astype(self.dtype)
|
|
v = (v.astype('float32') * v_out_scale).astype(self.dtype)
|
|
|
|
# add bias
|
|
q = q + q_bias
|
|
k = k + k_bias
|
|
v = v + v_bias
|
|
|
|
# transpose to origin
|
|
q = q.reshape(
|
|
[self.batch_size, self.seq_len, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
k = k.reshape(
|
|
[self.batch_size, self.seq_len, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
v = v.reshape(
|
|
[self.batch_size, self.seq_len, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
|
|
out_ = naive_attention_impl(
|
|
q, k, v, None, None, None, None, self.attention_mask, self.scale
|
|
)
|
|
|
|
out_ = remove_padding(
|
|
self.seq_lens_this_time, self.cu_seqlens_q, out_, self.token_num
|
|
)
|
|
|
|
# shift smooth
|
|
shift = np.random.random([self.num_head * self.dim_head])
|
|
shift = paddle.to_tensor(shift, dtype=self.dtype, place=self.place)
|
|
|
|
smooth = np.random.random([self.num_head * self.dim_head])
|
|
smooth = paddle.to_tensor(smooth, dtype=self.dtype, place=self.place)
|
|
|
|
out_ = (out_ + shift) * smooth
|
|
|
|
# quant
|
|
out_ *= 127.0
|
|
|
|
out_ = paddle.where(out_ <= -127, paddle.full_like(out_, -127), out_)
|
|
out_ = paddle.where(out_ >= 127, paddle.full_like(out_, 127), out_)
|
|
out_ = paddle.round(out_).astype('int8')
|
|
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
qkv_out_scale, # qkv_out_scale
|
|
qkv_bias, # qkv_bias
|
|
shift, # out_shift
|
|
smooth, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
self.seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style,
|
|
compute_dtype="fp16",
|
|
out_scale=1.0,
|
|
)[0]
|
|
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=1,
|
|
atol=1,
|
|
)
|
|
|
|
# decoder
|
|
naive_cache_k, naive_cache_v = block_cache_to_naive_cache(
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.batch_size,
|
|
self.block_tables,
|
|
self.seq_len,
|
|
)
|
|
|
|
self.seq_lens_decoder[:] = self.seq_lens_encoder
|
|
self.seq_lens_encoder[:] = 0
|
|
self.seq_lens_this_time[:] = 1
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
1,
|
|
self.dim_head,
|
|
)
|
|
query = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
key = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
value = np.random.randint(-65535, 65535, self.shape, 'int32')
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype='int32', stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.batch_size, -1])
|
|
|
|
q = q.transpose([0, 2, 1, 3]).reshape([self.batch_size, self.hid_dim])
|
|
k = k.transpose([0, 2, 1, 3]).reshape([self.batch_size, self.hid_dim])
|
|
v = v.transpose([0, 2, 1, 3]).reshape([self.batch_size, self.hid_dim])
|
|
|
|
q_out_scale = 1.0 / paddle.max(q, axis=0).astype('float32')
|
|
k_out_scale = 1.0 / paddle.max(k, axis=0).astype('float32')
|
|
v_out_scale = 1.0 / paddle.max(v, axis=0).astype('float32')
|
|
|
|
qkv_out_scale = paddle.concat(
|
|
[q_out_scale, k_out_scale, v_out_scale], axis=0
|
|
)
|
|
|
|
q_bias = paddle.ones([self.hid_dim], dtype=self.dtype) * 0.1
|
|
k_bias = paddle.ones([self.hid_dim], dtype=self.dtype) * 0.1
|
|
v_bias = paddle.ones([self.hid_dim], dtype=self.dtype) * 0.1
|
|
|
|
qkv_bias = paddle.concat([q_bias, k_bias, v_bias], axis=-1)
|
|
|
|
# dequant
|
|
q = (q.astype('float32') * q_out_scale).astype(self.dtype)
|
|
k = (k.astype('float32') * k_out_scale).astype(self.dtype)
|
|
v = (v.astype('float32') * v_out_scale).astype(self.dtype)
|
|
|
|
# add bias
|
|
q = q + q_bias
|
|
k = k + k_bias
|
|
v = v + v_bias
|
|
|
|
# transpose to origin
|
|
q = q.reshape(
|
|
[self.batch_size, 1, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
k = k.reshape(
|
|
[self.batch_size, 1, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
v = v.reshape(
|
|
[self.batch_size, 1, self.num_head, self.dim_head]
|
|
).transpose([0, 2, 1, 3])
|
|
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(self.batch_size, 1, self.seq_lens_this_time)
|
|
|
|
out_ = (
|
|
naive_attention_impl(
|
|
q,
|
|
k,
|
|
v,
|
|
naive_cache_k,
|
|
naive_cache_v,
|
|
None,
|
|
None,
|
|
None,
|
|
self.scale,
|
|
)
|
|
.transpose([0, 2, 1, 3])
|
|
.reshape([self.batch_size, -1])
|
|
)
|
|
|
|
# shift smooth
|
|
shift = np.random.random([self.num_head * self.dim_head])
|
|
shift = paddle.to_tensor(shift, dtype=self.dtype, place=self.place)
|
|
|
|
smooth = np.random.random([self.num_head * self.dim_head])
|
|
smooth = paddle.to_tensor(smooth, dtype=self.dtype, place=self.place)
|
|
|
|
out_ = (out_ + shift) * smooth
|
|
|
|
# quant
|
|
out_ *= 127.0
|
|
|
|
out_ = paddle.where(out_ <= -127, paddle.full_like(out_, -127), out_)
|
|
out_ = paddle.where(out_ >= 127, paddle.full_like(out_, 127), out_)
|
|
out_ = paddle.round(out_).astype('int8')
|
|
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
qkv_out_scale, # qkv_out_scale
|
|
qkv_bias, # qkv_bias
|
|
shift, # out_shift
|
|
smooth, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
1, # seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
compute_dtype="fp16",
|
|
out_scale=1.0,
|
|
)[0]
|
|
# NOTE: The diff of decoder is a little big
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=20,
|
|
atol=57,
|
|
)
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11040
|
|
or not is_sm_supported,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.4"
|
|
" and device's compute capability must be 7.x, 8.x or 9.x",
|
|
)
|
|
class TestBlockMultiHeadAttnEncDecQuant(unittest.TestCase):
|
|
def setUp(self):
|
|
paddle.disable_static()
|
|
self.name = "TestBlockMultiHeadAttnEncDec"
|
|
self.place = get_device_place()
|
|
self.batch_size = 2
|
|
self.num_head = 8
|
|
self.seq_len = 64
|
|
self.max_dec_len = 64
|
|
self.dim_head = 64
|
|
self.hid_dim = self.num_head * self.dim_head
|
|
self.blocksize = 64
|
|
self.block_num_per_seq = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
self.max_block_num = self.block_num_per_seq * self.batch_size
|
|
self.free_list = list(range(self.max_block_num - 1, -1, -1))
|
|
self.seq_lens_encoder = paddle.to_tensor(
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_decoder = paddle.to_tensor(
|
|
[
|
|
0,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_this_time = self.seq_lens_encoder
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
self.seq_len,
|
|
self.dim_head,
|
|
)
|
|
self.cache_shape = (
|
|
self.max_block_num,
|
|
self.num_head,
|
|
self.blocksize,
|
|
self.dim_head,
|
|
)
|
|
self.dtype = 'float16'
|
|
self.attention_mask = create_attn_mask(
|
|
self.dtype,
|
|
self.batch_size,
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
)
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.cache_k = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.cache_v = paddle.zeros(shape=self.cache_shape, dtype=self.dtype)
|
|
self.block_tables = paddle.zeros(
|
|
shape=(self.batch_size, self.block_num_per_seq), dtype="int32"
|
|
)
|
|
for i in range(self.batch_size):
|
|
need_block_num = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
for j in range(need_block_num):
|
|
self.block_tables[i, j] = self.free_list.pop()
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(
|
|
self.batch_size, self.seq_len, self.seq_lens_this_time
|
|
)
|
|
self.token_num = self.padding_offset.shape[0]
|
|
|
|
def test_all(self):
|
|
paddle.disable_static()
|
|
# encoder
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.token_num, -1])
|
|
out_ = naive_attention_impl(
|
|
q, k, v, None, None, None, None, self.attention_mask, self.scale
|
|
)
|
|
# quant
|
|
out_ *= 127.0
|
|
|
|
out_ = paddle.where(out_ <= -127, paddle.full_like(out_, -127), out_)
|
|
out_ = paddle.where(out_ >= 127, paddle.full_like(out_, 127), out_)
|
|
out_ = paddle.round(out_).astype('int8')
|
|
|
|
out_ = remove_padding(
|
|
self.seq_lens_this_time, self.cu_seqlens_q, out_, self.token_num
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
self.seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style,
|
|
out_scale=1.0,
|
|
)[0]
|
|
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=0.1,
|
|
atol=1,
|
|
)
|
|
|
|
# decoder
|
|
naive_cache_k, naive_cache_v = block_cache_to_naive_cache(
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.batch_size,
|
|
self.block_tables,
|
|
self.seq_len,
|
|
)
|
|
|
|
self.seq_lens_decoder[:] = self.seq_lens_encoder
|
|
self.seq_lens_encoder[:] = 0
|
|
self.seq_lens_this_time[:] = 1
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
1,
|
|
self.dim_head,
|
|
)
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.batch_size, -1])
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(self.batch_size, 1, self.seq_lens_this_time)
|
|
|
|
out_ = (
|
|
naive_attention_impl(
|
|
q,
|
|
k,
|
|
v,
|
|
naive_cache_k,
|
|
naive_cache_v,
|
|
None,
|
|
None,
|
|
None,
|
|
self.scale,
|
|
)
|
|
.transpose([0, 2, 1, 3])
|
|
.reshape([self.batch_size, -1])
|
|
)
|
|
# quant
|
|
out_ *= 127.0
|
|
|
|
out_ = paddle.where(out_ <= -127, paddle.full_like(out_, -127), out_)
|
|
out_ = paddle.where(out_ >= 127, paddle.full_like(out_, 127), out_)
|
|
out_ = paddle.round(out_).astype('int8')
|
|
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
None, # cache_k_quant_scales
|
|
None, # cache_v_quant_scales
|
|
None, # cache_k_dequant_scales
|
|
None, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
1, # seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
out_scale=1.0,
|
|
)[0]
|
|
# NOTE: The diff of decoder is a little big
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=0.1,
|
|
atol=1,
|
|
)
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11040
|
|
or not is_sm_supported,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.4"
|
|
"and device's compute capability must be 8.x or 90",
|
|
)
|
|
class TestBlockMultiHeadAttnEncDecCacheKVDynamicQuant(unittest.TestCase):
|
|
def setUp(self):
|
|
paddle.disable_static()
|
|
self.name = "TestBlockMultiHeadAttnEncDec"
|
|
self.place = get_device_place()
|
|
self.batch_size = 2
|
|
self.num_head = 8
|
|
self.seq_len = 64
|
|
self.max_dec_len = 64
|
|
self.dim_head = 64
|
|
self.hid_dim = self.num_head * self.dim_head
|
|
self.blocksize = 64
|
|
self.block_num_per_seq = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
self.max_block_num = self.block_num_per_seq * self.batch_size
|
|
self.free_list = list(range(self.max_block_num - 1, -1, -1))
|
|
self.seq_lens_encoder = paddle.to_tensor(
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_decoder = paddle.to_tensor(
|
|
[
|
|
0,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_this_time = self.seq_lens_encoder
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
self.seq_len,
|
|
self.dim_head,
|
|
)
|
|
self.cache_shape = (
|
|
self.max_block_num,
|
|
self.num_head,
|
|
self.blocksize,
|
|
self.dim_head,
|
|
)
|
|
self.dtype = 'float16'
|
|
self.attention_mask = create_attn_mask(
|
|
self.dtype,
|
|
self.batch_size,
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
)
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.cache_k = paddle.zeros(shape=self.cache_shape, dtype='uint8')
|
|
self.cache_v = paddle.zeros(shape=self.cache_shape, dtype='uint8')
|
|
self.cache_k_quant_scales = paddle.zeros(
|
|
shape=[self.batch_size, self.num_head], dtype='float32'
|
|
)
|
|
self.cache_v_quant_scales = paddle.zeros(
|
|
shape=[self.batch_size, self.num_head], dtype='float32'
|
|
)
|
|
self.cache_k_dequant_scales = paddle.zeros(
|
|
shape=[self.batch_size, self.num_head], dtype='float32'
|
|
)
|
|
self.cache_v_dequant_scales = paddle.zeros(
|
|
shape=[self.batch_size, self.num_head], dtype='float32'
|
|
)
|
|
|
|
self.block_tables = paddle.zeros(
|
|
shape=(self.batch_size, self.block_num_per_seq), dtype="int32"
|
|
)
|
|
for i in range(self.batch_size):
|
|
need_block_num = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
for j in range(need_block_num):
|
|
self.block_tables[i, j] = self.free_list.pop()
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(
|
|
self.batch_size, self.seq_len, self.seq_lens_this_time
|
|
)
|
|
self.token_num = self.padding_offset.shape[0]
|
|
|
|
def test_all(self):
|
|
paddle.disable_static()
|
|
# encoder
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.token_num, -1])
|
|
out_ = naive_attention_impl(
|
|
q, k, v, None, None, None, None, self.attention_mask, self.scale
|
|
)
|
|
|
|
out_ = remove_padding(
|
|
self.seq_lens_this_time, self.cu_seqlens_q, out_, self.token_num
|
|
)
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
self.cache_k_quant_scales, # cache_k_quant_scales
|
|
self.cache_v_quant_scales, # cache_v_quant_scales
|
|
self.cache_k_dequant_scales, # cache_k_dequant_scales
|
|
self.cache_v_dequant_scales, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
self.seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style,
|
|
use_dynamic_cachekv_quant=True,
|
|
)[0]
|
|
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=0.1,
|
|
atol=1,
|
|
)
|
|
|
|
# decoder
|
|
naive_cache_k, naive_cache_v = block_cache_to_naive_cache(
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.batch_size,
|
|
self.block_tables,
|
|
self.seq_len,
|
|
)
|
|
|
|
self.seq_lens_decoder[:] = self.seq_lens_encoder
|
|
self.seq_lens_encoder[:] = 0
|
|
self.seq_lens_this_time[:] = 1
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
1,
|
|
self.dim_head,
|
|
)
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.batch_size, -1])
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(self.batch_size, 1, self.seq_lens_this_time)
|
|
|
|
out_ = (
|
|
naive_attention_impl(
|
|
q,
|
|
k,
|
|
v,
|
|
naive_cache_k,
|
|
naive_cache_v,
|
|
None,
|
|
None,
|
|
None,
|
|
self.scale,
|
|
cache_k_dequant_scales=self.cache_k_dequant_scales,
|
|
cache_v_dequant_scales=self.cache_v_dequant_scales,
|
|
use_cachekv_int8="dynamic",
|
|
)
|
|
.transpose([0, 2, 1, 3])
|
|
.reshape([self.batch_size, -1])
|
|
)
|
|
# quant
|
|
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
self.cache_k_quant_scales, # cache_k_quant_scales
|
|
self.cache_v_quant_scales, # cache_v_quant_scales
|
|
self.cache_k_dequant_scales, # cache_k_dequant_scales
|
|
self.cache_v_dequant_scales, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
1, # seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
use_dynamic_cachekv_quant=True,
|
|
)[0]
|
|
# NOTE: The diff of decoder is a little big
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=0.1,
|
|
atol=1,
|
|
)
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11040
|
|
or not is_sm_supported,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.4"
|
|
"and device's compute capability must be 8.x or 90",
|
|
)
|
|
class TestBlockMultiHeadAttnEncDecCacheKVStaticQuant(unittest.TestCase):
|
|
def setUp(self):
|
|
paddle.disable_static()
|
|
self.name = "TestBlockMultiHeadAttnEncDec"
|
|
self.place = get_device_place()
|
|
self.batch_size = 2
|
|
self.num_head = 8
|
|
self.seq_len = 64
|
|
self.max_dec_len = 64
|
|
self.dim_head = 64
|
|
self.hid_dim = self.num_head * self.dim_head
|
|
self.blocksize = 64
|
|
self.block_num_per_seq = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
self.max_block_num = self.block_num_per_seq * self.batch_size
|
|
self.free_list = list(range(self.max_block_num - 1, -1, -1))
|
|
self.seq_lens_encoder = paddle.to_tensor(
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_decoder = paddle.to_tensor(
|
|
[
|
|
0,
|
|
]
|
|
* self.batch_size,
|
|
"int32",
|
|
)
|
|
self.seq_lens_this_time = self.seq_lens_encoder
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
self.seq_len,
|
|
self.dim_head,
|
|
)
|
|
self.cache_shape = (
|
|
self.max_block_num,
|
|
self.num_head,
|
|
self.blocksize,
|
|
self.dim_head,
|
|
)
|
|
self.dtype = 'float16'
|
|
self.attention_mask = create_attn_mask(
|
|
self.dtype,
|
|
self.batch_size,
|
|
[
|
|
self.seq_len,
|
|
]
|
|
* self.batch_size,
|
|
)
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.cache_k = paddle.zeros(shape=self.cache_shape, dtype='uint8')
|
|
self.cache_v = paddle.zeros(shape=self.cache_shape, dtype='uint8')
|
|
self.cache_k_quant_scales = paddle.zeros(
|
|
shape=[self.num_head], dtype='float32'
|
|
)
|
|
self.cache_v_quant_scales = paddle.zeros(
|
|
shape=[self.num_head], dtype='float32'
|
|
)
|
|
self.cache_k_dequant_scales = paddle.zeros(
|
|
shape=[self.num_head], dtype='float32'
|
|
)
|
|
self.cache_v_dequant_scales = paddle.zeros(
|
|
shape=[self.num_head], dtype='float32'
|
|
)
|
|
|
|
self.block_tables = paddle.zeros(
|
|
shape=(self.batch_size, self.block_num_per_seq), dtype="int32"
|
|
)
|
|
for i in range(self.batch_size):
|
|
need_block_num = (
|
|
self.seq_len + self.max_dec_len + self.blocksize - 1
|
|
) // self.blocksize
|
|
for j in range(need_block_num):
|
|
self.block_tables[i, j] = self.free_list.pop()
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(
|
|
self.batch_size, self.seq_len, self.seq_lens_this_time
|
|
)
|
|
self.token_num = self.padding_offset.shape[0]
|
|
|
|
def test_all(self):
|
|
paddle.disable_static()
|
|
# encoder
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.token_num, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.token_num, -1])
|
|
|
|
out_ = naive_attention_impl(
|
|
q, k, v, None, None, None, None, self.attention_mask, self.scale
|
|
)
|
|
|
|
out_ = remove_padding(
|
|
self.seq_lens_this_time, self.cu_seqlens_q, out_, self.token_num
|
|
)
|
|
|
|
self.cache_k_quant_scales = (
|
|
127.0 / paddle.max(k, axis=[0, 2, 3])
|
|
).astype("float32")
|
|
self.cache_v_quant_scales = (
|
|
127.0 / paddle.max(k, axis=[0, 2, 3])
|
|
).astype("float32")
|
|
|
|
self.cache_k_dequant_scales = 1.0 / self.cache_k_quant_scales
|
|
self.cache_v_dequant_scales = 1.0 / self.cache_v_quant_scales
|
|
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
self.cache_k_quant_scales, # cache_k_quant_scales
|
|
self.cache_v_quant_scales, # cache_v_quant_scales
|
|
self.cache_k_dequant_scales, # cache_k_dequant_scales
|
|
self.cache_v_dequant_scales, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
self.seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style,
|
|
use_dynamic_cachekv_quant=False,
|
|
)[0]
|
|
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=0.1,
|
|
atol=1,
|
|
)
|
|
|
|
# decoder
|
|
naive_cache_k, naive_cache_v = block_cache_to_naive_cache(
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.batch_size,
|
|
self.block_tables,
|
|
self.seq_len,
|
|
)
|
|
|
|
self.seq_lens_decoder[:] = self.seq_lens_encoder
|
|
self.seq_lens_encoder[:] = 0
|
|
self.seq_lens_this_time[:] = 1
|
|
self.shape = (
|
|
self.batch_size,
|
|
self.num_head,
|
|
1,
|
|
self.dim_head,
|
|
)
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
qkv = paddle.stack(
|
|
[
|
|
q.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
k.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
v.transpose([0, 2, 1, 3]).reshape(
|
|
[self.batch_size, self.hid_dim]
|
|
),
|
|
],
|
|
axis=1,
|
|
).reshape([self.batch_size, -1])
|
|
(
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
) = get_padding_offset(self.batch_size, 1, self.seq_lens_this_time)
|
|
|
|
out_ = (
|
|
naive_attention_impl(
|
|
q,
|
|
k,
|
|
v,
|
|
naive_cache_k,
|
|
naive_cache_v,
|
|
None,
|
|
None,
|
|
None,
|
|
self.scale,
|
|
cache_k_dequant_scales=self.cache_k_dequant_scales,
|
|
cache_v_dequant_scales=self.cache_v_dequant_scales,
|
|
use_cachekv_int8="static",
|
|
)
|
|
.transpose([0, 2, 1, 3])
|
|
.reshape([self.batch_size, -1])
|
|
)
|
|
|
|
out = block_multihead_attention(
|
|
qkv,
|
|
self.cache_k,
|
|
self.cache_v,
|
|
self.seq_lens_encoder,
|
|
self.seq_lens_decoder,
|
|
self.seq_lens_this_time,
|
|
self.padding_offset,
|
|
self.cum_offset,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.block_tables,
|
|
None, # pre_key_cache
|
|
None, # pre_value_cache
|
|
self.cache_k_quant_scales, # cache_k_quant_scales
|
|
self.cache_v_quant_scales, # cache_v_quant_scales
|
|
self.cache_k_dequant_scales, # cache_k_dequant_scales
|
|
self.cache_v_dequant_scales, # cache_v_dequant_scales
|
|
None, # qkv_out_scale
|
|
None, # qkv_bias
|
|
None, # out_shift
|
|
None, # out_smooth
|
|
None, # max_enc_len_this_time
|
|
None, # max_dec_len_this_time
|
|
None, # rotary_embs
|
|
None, # attn_mask
|
|
None, # tgt_mask
|
|
1, # seq_len,
|
|
self.blocksize,
|
|
False, # use_neox_rotary_style
|
|
use_dynamic_cachekv_quant=False,
|
|
)[0]
|
|
# NOTE: The diff of decoder is a little big
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
out_.numpy(),
|
|
rtol=0.1,
|
|
atol=1,
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|