# Copyright (c) 2024 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_device_place, is_custom_device from test_block_multihead_attention import ( RopeEmbedding, block_cache_to_naive_cache, create_attn_mask, get_cuda_version, get_padding_offset, is_sm_supported, remove_padding, ) 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(2024) np.random.seed(2024) 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: pre_cache_k = pre_cache_k.reshape([batch, kv_head, 1, -1, head_dim]) pre_cache_k = paddle.tile(pre_cache_k, [1, 1, heads // kv_head, 1, 1]) pre_cache_k = pre_cache_k.reshape([batch, heads, -1, head_dim]) 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) dequant_cache_k = dequant_cache_k.reshape( [batch, kv_head, 1, -1, head_dim] ) dequant_cache_k = paddle.tile( dequant_cache_k, [1, 1, heads // kv_head, 1, 1] ) dequant_cache_k = dequant_cache_k.reshape( [batch, heads, -1, head_dim] ) key = paddle.concat([dequant_cache_k, key], axis=2) else: cache_k = cache_k.reshape([batch, kv_head, 1, -1, head_dim]) cache_k = paddle.tile(cache_k, [1, 1, heads // kv_head, 1, 1]) cache_k = cache_k.reshape([batch, heads, -1, head_dim]) 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: pre_cache_v = pre_cache_v.reshape([batch, kv_head, 1, -1, head_dim]) pre_cache_v = paddle.tile(pre_cache_v, [1, 1, heads // kv_head, 1, 1]) pre_cache_v = pre_cache_v.reshape([batch, heads, -1, head_dim]) 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) dequant_cache_v = dequant_cache_v.reshape( [batch, kv_head, 1, -1, head_dim] ) dequant_cache_v = paddle.tile( dequant_cache_v, [1, 1, heads // kv_head, 1, 1] ) dequant_cache_v = dequant_cache_v.reshape( [batch, heads, -1, head_dim] ) value = paddle.concat([dequant_cache_v, value], axis=2) else: cache_v = cache_v.reshape([batch, kv_head, 1, -1, head_dim]) cache_v = paddle.tile(cache_v, [1, 1, heads // kv_head, 1, 1]) cache_v = cache_v.reshape([batch, heads, -1, head_dim]) 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 @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 TestBlockGroupQueryAttnEncDec(unittest.TestCase): def setUp(self): paddle.disable_static() self.name = "TestBlockGroupQueryAttnEncDec" self.place = get_device_place() self.batch_size = 2 self.q_num_head = 8 self.kv_num_head = 2 self.seq_len = 64 self.max_dec_len = 64 self.dim_head = 64 self.q_hid_dim = self.q_num_head * self.dim_head self.kv_hid_dim = self.kv_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.q_shape = ( self.batch_size, self.q_num_head, self.seq_len, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, self.seq_len, self.dim_head, ) self.cache_shape = ( self.max_block_num, self.kv_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.q_num_head, 1, self.seq_len + 1], dtype=self.dtype, ) # self.tgt_mask = None self.scale = 1.0 / np.sqrt(self.q_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.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_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.q_shape = ( self.batch_size, self.q_num_head, 1, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, 1, self.dim_head, ) query = np.random.random(self.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_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=7e-02, atol=7e-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 TestBlockGroupQueryAttnEncDecSkipGetMaxLen(unittest.TestCase): def setUp(self): paddle.disable_static() self.name = "TestBlockGroupQueryAttnEncDecSkipGetMaxLen" self.place = get_device_place() self.batch_size = 2 self.q_num_head = 8 self.kv_num_head = 2 self.seq_len = 64 self.max_dec_len = 64 self.dim_head = 64 self.q_hid_dim = self.q_num_head * self.dim_head self.kv_hid_dim = self.kv_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.q_shape = ( self.batch_size, self.q_num_head, self.seq_len, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, self.seq_len, self.dim_head, ) self.cache_shape = ( self.max_block_num, self.kv_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.q_num_head, 1, self.seq_len + 1], dtype=self.dtype, ) # self.tgt_mask = None self.scale = 1.0 / np.sqrt(self.q_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.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_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.q_shape = ( self.batch_size, self.q_num_head, 1, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, 1, self.dim_head, ) query = np.random.random(self.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_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 TestBlockGroupQueryAttnRoPE(unittest.TestCase): def setUp(self): paddle.disable_static() self.name = "TestBlockGroupQueryAttnRoPE" self.place = get_device_place() self.batch_size = 2 self.q_num_head = 8 self.kv_num_head = 2 self.seq_len = 64 self.max_dec_len = 64 self.dim_head = 64 self.q_hid_dim = self.q_num_head * self.dim_head self.kv_hid_dim = self.kv_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.q_shape = ( self.batch_size, self.q_num_head, self.seq_len, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, self.seq_len, self.dim_head, ) self.cache_shape = ( self.max_block_num, self.kv_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.q_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.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_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.q_shape = ( self.batch_size, self.q_num_head, 1, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, 1, self.dim_head, ) query = np.random.random(self.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_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=7e-02, atol=7e-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 TestBlockGroupQueryAttnEncStatic(unittest.TestCase): def setUp(self): paddle.disable_static() self.name = "TestBlockGroupQueryAttnEncStatic" self.place = get_device_place() self.batch_size = 2 self.q_num_head = 8 self.kv_num_head = 2 self.seq_len = 64 self.max_dec_len = 64 self.dim_head = 64 self.q_hid_dim = self.q_num_head * self.dim_head self.kv_hid_dim = self.kv_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.q_shape = ( self.batch_size, self.q_num_head, self.seq_len, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, self.seq_len, self.dim_head, ) self.cache_shape = ( self.max_block_num, self.kv_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.q_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.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_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.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_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, 2 * self.kv_hid_dim + self.q_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 TestBlockGroupQueryAttnEncDecPTQDequant(unittest.TestCase): def setUp(self): paddle.disable_static() self.name = "TestBlockGroupQueryAttnEncDecPTQDequant" self.place = get_device_place() self.batch_size = 2 self.q_num_head = 8 self.kv_num_head = 2 self.seq_len = 64 self.max_dec_len = 64 self.dim_head = 64 self.q_hid_dim = self.q_num_head * self.dim_head self.kv_hid_dim = self.kv_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.q_shape = ( self.batch_size, self.q_num_head, self.seq_len, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, self.seq_len, self.dim_head, ) self.cache_shape = ( self.max_block_num, self.kv_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.q_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.q_shape, 'int32') q = paddle.to_tensor( query, place=self.place, dtype='int32', stop_gradient=False ) key = np.random.randint(-65535, 65535, self.kv_shape, 'int32') k = paddle.to_tensor( key, place=self.place, dtype='int32', stop_gradient=False ) value = np.random.randint(-65535, 65535, self.kv_shape, 'int32') v = paddle.to_tensor( value, place=self.place, dtype='int32', stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_hid_dim] ), ], axis=1, ).reshape([self.token_num, -1]) q = q.transpose([0, 2, 1, 3]).reshape([self.token_num, self.q_hid_dim]) k = k.transpose([0, 2, 1, 3]).reshape([self.token_num, self.kv_hid_dim]) v = v.transpose([0, 2, 1, 3]).reshape([self.token_num, self.kv_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.q_hid_dim], dtype=self.dtype) k_bias = paddle.ones([self.kv_hid_dim], dtype=self.dtype) v_bias = paddle.ones([self.kv_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.q_num_head, self.dim_head] ).transpose([0, 2, 1, 3]) k = k.reshape( [self.batch_size, self.seq_len, self.kv_num_head, self.dim_head] ).transpose([0, 2, 1, 3]) v = v.reshape( [self.batch_size, self.seq_len, self.kv_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.q_shape = ( self.batch_size, self.q_num_head, 1, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, 1, self.dim_head, ) query = np.random.randint(-65535, 65535, self.q_shape, 'int32') q = paddle.to_tensor( query, place=self.place, dtype='int32', stop_gradient=False ) key = np.random.randint(-65535, 65535, self.kv_shape, 'int32') k = paddle.to_tensor( key, place=self.place, dtype='int32', stop_gradient=False ) value = np.random.randint(-65535, 65535, self.kv_shape, 'int32') v = paddle.to_tensor( value, place=self.place, dtype='int32', stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ), ], axis=1, ).reshape([self.batch_size, -1]) q = q.transpose([0, 2, 1, 3]).reshape([self.batch_size, self.q_hid_dim]) k = k.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ) v = v.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_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.q_hid_dim], dtype=self.dtype) * 0.1 k_bias = paddle.ones([self.kv_hid_dim], dtype=self.dtype) * 0.1 v_bias = paddle.ones([self.kv_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.q_num_head, self.dim_head] ).transpose([0, 2, 1, 3]) k = k.reshape( [self.batch_size, 1, self.kv_num_head, self.dim_head] ).transpose([0, 2, 1, 3]) v = v.reshape( [self.batch_size, 1, self.kv_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 TestBlockGroupQueryAttnEncDecPTQDequantQuantShiftSmooth( unittest.TestCase ): def setUp(self): paddle.disable_static() self.name = "TestBlockGroupQueryAttnEncDecPTQDequantQuantShiftSmooth" self.place = get_device_place() self.batch_size = 2 self.q_num_head = 8 self.kv_num_head = 2 self.seq_len = 64 self.max_dec_len = 64 self.dim_head = 64 self.q_hid_dim = self.q_num_head * self.dim_head self.kv_hid_dim = self.kv_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.q_shape = ( self.batch_size, self.q_num_head, self.seq_len, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, self.seq_len, self.dim_head, ) self.cache_shape = ( self.max_block_num, self.kv_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.q_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.q_shape, 'int32') q = paddle.to_tensor( query, place=self.place, dtype='int32', stop_gradient=False ) key = np.random.randint(-65535, 65535, self.kv_shape, 'int32') k = paddle.to_tensor( key, place=self.place, dtype='int32', stop_gradient=False ) value = np.random.randint(-65535, 65535, self.kv_shape, 'int32') v = paddle.to_tensor( value, place=self.place, dtype='int32', stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_hid_dim] ), ], axis=1, ).reshape([self.token_num, -1]) q = q.transpose([0, 2, 1, 3]).reshape([self.token_num, self.q_hid_dim]) k = k.transpose([0, 2, 1, 3]).reshape([self.token_num, self.kv_hid_dim]) v = v.transpose([0, 2, 1, 3]).reshape([self.token_num, self.kv_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.q_hid_dim], dtype=self.dtype) k_bias = paddle.ones([self.kv_hid_dim], dtype=self.dtype) v_bias = paddle.ones([self.kv_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.q_num_head, self.dim_head] ).transpose([0, 2, 1, 3]) k = k.reshape( [self.batch_size, self.seq_len, self.kv_num_head, self.dim_head] ).transpose([0, 2, 1, 3]) v = v.reshape( [self.batch_size, self.seq_len, self.kv_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.q_num_head * self.dim_head]) shift = paddle.to_tensor(shift, dtype=self.dtype, place=self.place) smooth = np.random.random([self.q_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.q_shape = ( self.batch_size, self.q_num_head, 1, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, 1, self.dim_head, ) query = np.random.randint(-65535, 65535, self.q_shape, 'int32') q = paddle.to_tensor( query, place=self.place, dtype='int32', stop_gradient=False ) key = np.random.randint(-65535, 65535, self.kv_shape, 'int32') k = paddle.to_tensor( key, place=self.place, dtype='int32', stop_gradient=False ) value = np.random.randint(-65535, 65535, self.kv_shape, 'int32') v = paddle.to_tensor( value, place=self.place, dtype='int32', stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ), ], axis=1, ).reshape([self.batch_size, -1]) q = q.transpose([0, 2, 1, 3]).reshape([self.batch_size, self.q_hid_dim]) k = k.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ) v = v.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_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.q_hid_dim], dtype=self.dtype) * 0.1 k_bias = paddle.ones([self.kv_hid_dim], dtype=self.dtype) * 0.1 v_bias = paddle.ones([self.kv_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.q_num_head, self.dim_head] ).transpose([0, 2, 1, 3]) k = k.reshape( [self.batch_size, 1, self.kv_num_head, self.dim_head] ).transpose([0, 2, 1, 3]) v = v.reshape( [self.batch_size, 1, self.kv_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.q_num_head * self.dim_head]) shift = paddle.to_tensor(shift, dtype=self.dtype, place=self.place) smooth = np.random.random([self.q_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 TestBlockGroupQueryAttnEncDecQuant(unittest.TestCase): def setUp(self): paddle.disable_static() self.name = "TestBlockGroupQueryAttnEncDecQuant" self.place = get_device_place() self.batch_size = 2 self.q_num_head = 8 self.kv_num_head = 2 self.seq_len = 64 self.max_dec_len = 64 self.dim_head = 64 self.q_hid_dim = self.q_num_head * self.dim_head self.kv_hid_dim = self.kv_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.q_shape = ( self.batch_size, self.q_num_head, self.seq_len, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, self.seq_len, self.dim_head, ) self.cache_shape = ( self.max_block_num, self.kv_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.q_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 ) # batch_size * seq_len self.token_num = self.padding_offset.shape[0] def test_all(self): paddle.disable_static() # encoder query = np.random.random(self.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_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.q_shape = ( self.batch_size, self.q_num_head, 1, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, 1, self.dim_head, ) query = np.random.random(self.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_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 TestBlockGroupQueryAttnEncDecCacheKVDynamicQuant(unittest.TestCase): def setUp(self): paddle.disable_static() self.name = "TestBlockGroupQueryAttnEncDecCacheKVDynamicQuant" self.place = get_device_place() self.batch_size = 2 self.q_num_head = 8 self.kv_num_head = 2 self.seq_len = 64 self.max_dec_len = 64 self.dim_head = 64 self.q_hid_dim = self.q_num_head * self.dim_head self.kv_hid_dim = self.kv_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.q_shape = ( self.batch_size, self.q_num_head, self.seq_len, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, self.seq_len, self.dim_head, ) self.cache_shape = ( self.max_block_num, self.kv_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.q_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.kv_num_head], dtype='float32' ) self.cache_v_quant_scales = paddle.zeros( shape=[self.batch_size, self.kv_num_head], dtype='float32' ) self.cache_k_dequant_scales = paddle.zeros( shape=[self.batch_size, self.kv_num_head], dtype='float32' ) self.cache_v_dequant_scales = paddle.zeros( shape=[self.batch_size, self.kv_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.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_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.q_shape = ( self.batch_size, self.q_num_head, 1, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, 1, self.dim_head, ) query = np.random.random(self.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_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 TestBlockGroupQueryAttnEncDecCacheKVStaticQuant(unittest.TestCase): def setUp(self): paddle.disable_static() self.name = "TestBlockGroupQueryAttnEncDecCacheKVStaticQuant" self.place = get_device_place() self.batch_size = 2 self.q_num_head = 8 self.kv_num_head = 2 self.seq_len = 64 self.max_dec_len = 64 self.dim_head = 64 self.q_hid_dim = self.q_num_head * self.dim_head self.kv_hid_dim = self.kv_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.q_shape = ( self.batch_size, self.q_num_head, self.seq_len, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, self.seq_len, self.dim_head, ) self.cache_shape = ( self.max_block_num, self.kv_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.q_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.kv_num_head], dtype='float32' ) self.cache_v_quant_scales = paddle.zeros( shape=[self.kv_num_head], dtype='float32' ) self.cache_k_dequant_scales = paddle.zeros( shape=[self.kv_num_head], dtype='float32' ) self.cache_v_dequant_scales = paddle.zeros( shape=[self.kv_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.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.token_num, self.kv_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.q_shape = ( self.batch_size, self.q_num_head, 1, self.dim_head, ) self.kv_shape = ( self.batch_size, self.kv_num_head, 1, self.dim_head, ) query = np.random.random(self.q_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) key = np.random.random(self.kv_shape) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) value = np.random.random(self.kv_shape) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) qkv = paddle.concat( [ q.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.q_hid_dim] ), k.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_hid_dim] ), v.transpose([0, 2, 1, 3]).reshape( [self.batch_size, self.kv_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()