# 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 logging import unittest import numpy as np from op_test import get_cuda_version, get_device_place, is_custom_device import paddle import paddle.nn.functional as F from paddle import base from paddle.base import core from paddle.nn.functional import ( scaled_dot_product_attention, sdp_kernel, ) from paddle.nn.functional.flash_attention import ( calc_reduced_attention_scores, flash_attention, flash_attention_v3_varlen, flash_attn_qkvpacked, flash_attn_unpadded, flash_attn_varlen_qkvpacked, flashmask_attention, ) logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s") def attention_naive(q, k, v, causal=False): qt = paddle.transpose(q, [0, 2, 1, 3]) kt = paddle.transpose(k, [0, 2, 1, 3]) vt = paddle.transpose(v, [0, 2, 1, 3]) scale = 1.0 / np.sqrt(q.shape[-1]) s = paddle.matmul(qt * scale, paddle.transpose(kt, [0, 1, 3, 2])) p = ( paddle.incubate.softmax_mask_fuse_upper_triangle(s) if causal else F.softmax(s) ) o = paddle.matmul(p, vt) return paddle.transpose(o, [0, 2, 1, 3]) def attention_naive_with_mask(q, k, v, attn_bias): qt = paddle.transpose(q, [0, 2, 1, 3]) kt = paddle.transpose(k, [0, 2, 1, 3]) vt = paddle.transpose(v, [0, 2, 1, 3]) scale = 1.0 / np.sqrt(q.shape[-1]) s = paddle.matmul(qt, paddle.transpose(kt, [0, 1, 3, 2])) s = paddle.scale(s, scale) p = F.softmax(s + attn_bias) o = paddle.matmul(p, vt) return paddle.transpose(o, [0, 2, 1, 3]) def attention_naive_with_mask_and_scale(q, k, v, attn_bias, scale): """ Naive attention implementation that accepts a custom scale factor. """ q = q.float() k = k.float() v = v.float() attn_bias = attn_bias.float() if attn_bias is not None else None qt = paddle.transpose(q, [0, 2, 1, 3]) kt = paddle.transpose(k, [0, 2, 1, 3]) vt = paddle.transpose(v, [0, 2, 1, 3]) scale_factor = scale if scale is not None else (1.0 / np.sqrt(q.shape[-1])) s = paddle.matmul(qt * scale_factor, paddle.transpose(kt, [0, 1, 3, 2])) if attn_bias is not None: s = s + attn_bias p = F.softmax(s) o = paddle.matmul(p, vt) return paddle.transpose(o, [0, 2, 1, 3]) is_sm80 = ( (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_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_sm90 = ( (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_sm_supported = is_sm8x or is_sm90 def is_flashattn_supported(): if ( not (core.is_compiled_with_cuda() or is_custom_device()) or get_cuda_version() < 11040 or not is_sm_supported ): return False return True @unittest.skipIf( not is_flashattn_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 TestFlashAttentionAPI(unittest.TestCase): def setUp(self): self.place = get_device_place() self.shape = (2, 128, 8, 16) self.dtype = 'float16' self.dropout = 0.0 self.causal = False self.return_softmax = False self.use_sdp_kernel = False self.use_sdp_api = False def test_unpadded(self): print( f"Test unpadded case shape {self.shape} dtype {self.dtype} causal {self.causal}" ) paddle.disable_static() query = np.random.random(self.shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) q_ = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) out_ = attention_naive(q_, q_, q_, self.causal) scale = 1.0 / np.sqrt(q.shape[-1]) bs = self.shape[0] ms = self.shape[1] nh = self.shape[2] hd = self.shape[3] cu_q = paddle.arange(0, (bs + 1) * ms, ms, dtype='int32') qq = paddle.reshape(q, [bs * ms, nh, hd]) if ( is_sm90 and paddle.base.framework.get_flags(["FLAGS_flash_attn_version"]) == 3 ): assert self.dropout == 0.0, ( "flash_attention_v3_varlen not support dropout" ) out, _ = flash_attention_v3_varlen( query=qq, key=qq, value=qq, cu_seqlens_q=cu_q, cu_seqlens_k=cu_q, max_seqlen_q=ms, max_seqlen_k=ms, causal=self.causal, ) else: out, _ = flash_attn_unpadded( qq, qq, qq, cu_q, cu_q, ms, ms, scale, self.dropout, self.causal, self.return_softmax, ) out_ = paddle.reshape(out_, [bs * ms, nh, hd]) np.testing.assert_allclose(out.numpy(), out_, rtol=5e-03, atol=1e-03) out.backward() out_.backward() np.testing.assert_allclose( q.grad.numpy(), q_.grad.numpy(), rtol=5e-03, atol=1e-03 ) # test static paddle.enable_static() main_program = paddle.static.Program() with paddle.static.program_guard(main_program): qs = paddle.static.data( name="q", shape=self.shape, dtype=self.dtype ) cu_q = paddle.arange(0, (bs + 1) * ms, ms, dtype='int32') qs = paddle.reshape(qs, [bs * ms, nh, hd]) outs, softmax = flash_attn_unpadded( qs, qs, qs, cu_q, cu_q, ms, ms, scale, self.dropout, self.causal, self.return_softmax, ) shape_analysis = ( paddle.base.libpaddle.pir.get_shape_constraint_ir_analysis( main_program ) ) first_out_shape_or_data = shape_analysis.get_shape_or_data_for_var( outs[0] ) exe = base.Executor(self.place) fetches_result = exe.run( main_program, feed={ "q": query.astype('float16'), "k": query.astype('float16'), "v": query.astype('float16'), }, fetch_list=[outs], ) self.assertTrue( first_out_shape_or_data.is_equal(list(outs[0].shape)) ) np.testing.assert_allclose( fetches_result[0], out_, rtol=5e-03, atol=1e-03 ) paddle.disable_static() def test_all(self): print( f"Test case shape {self.shape} dtype {self.dtype} causal {self.causal}" ) # test dynamic paddle.disable_static() query = np.random.random(self.shape) key = np.random.random(self.shape) value = np.random.random(self.shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) q_ = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) k_ = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) v_ = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) if self.use_sdp_kernel: with paddle.nn.functional.sdp_kernel( enable_math=self.enable_math, enable_flash=self.enable_flash, enable_mem_efficient=self.enable_mem_efficient, ): if self.use_sdp_api: out = scaled_dot_product_attention( q, k, v, None, self.dropout, self.causal ) else: out, _ = flash_attention( q, k, v, self.dropout, self.causal, self.return_softmax ) else: out, _ = flash_attention( q, k, v, self.dropout, self.causal, self.return_softmax ) out_ = attention_naive(q_, k_, v_, self.causal) out.backward() out_.backward() np.testing.assert_allclose(out.numpy(), out_, rtol=5e-03, atol=1e-03) self.assertEqual(q.grad.shape, q.shape) self.assertEqual(q_.grad.shape, q.shape) np.testing.assert_allclose( q.grad.numpy(), q_.grad.numpy(), rtol=5e-03, atol=2e-03 ) # test static paddle.enable_static() with paddle.static.program_guard(paddle.static.Program()): qs = paddle.static.data( name="q", shape=self.shape, dtype=self.dtype ) ks = paddle.static.data( name="k", shape=self.shape, dtype=self.dtype ) vs = paddle.static.data( name="v", shape=self.shape, dtype=self.dtype ) if self.use_sdp_kernel: with paddle.nn.functional.sdp_kernel( enable_math=self.enable_math, enable_flash=self.enable_flash, enable_mem_efficient=self.enable_mem_efficient, ): if self.use_sdp_api: outs = scaled_dot_product_attention( qs, ks, vs, None, self.dropout, self.causal ) else: outs, softmax = flash_attention( qs, ks, vs, self.dropout, self.causal, self.return_softmax, ) else: outs, softmax = flash_attention( qs, ks, vs, self.dropout, self.causal, self.return_softmax ) exe = base.Executor(self.place) fetches_result = exe.run( feed={ "q": query.astype('float16'), "k": key.astype('float16'), "v": value.astype('float16'), }, fetch_list=[outs], ) np.testing.assert_allclose( fetches_result[0], out_, rtol=5e-03, atol=1e-03 ) paddle.disable_static() @unittest.skipIf( not is_flashattn_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.5 or 8.x", ) class TestFlashAttentionWithMaskAPI(unittest.TestCase): def setUp(self): self.place = get_device_place() self.shape = (2, 128, 8, 32) self.dtype = 'float16' self.dropout = 0.0 self.causal = False def test_dot_scale_product(self): # test dynamic paddle.disable_static() query = np.random.random(self.shape) key = np.random.random(self.shape) value = np.random.random(self.shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) q_ = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) k_ = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) v_ = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) mask_shape = (self.shape[0], 1, self.shape[1], self.shape[1]) mask = np.random.random(mask_shape) m = paddle.to_tensor( mask, place=self.place, dtype=self.dtype, stop_gradient=False ) out = scaled_dot_product_attention( q, k, v, m, self.dropout, self.causal ) out_ = attention_naive_with_mask(q_, k_, v_, m) out.backward() out_.backward() np.testing.assert_allclose(out.numpy(), out_, rtol=5e-03, atol=1e-03) class TestFlashAttentionAPITest1(TestFlashAttentionAPI): def setUp(self): self.place = get_device_place() self.shape = (2, 128, 8, 16) self.dtype = 'float16' self.dropout = 0.0 self.causal = False self.return_softmax = False self.use_sdp_kernel = False class TestFlashAttentionAPITest2(TestFlashAttentionAPI): def setUp(self): self.place = get_device_place() self.shape = (2, 256, 8, 16) self.dtype = 'float16' self.dropout = 0.0 self.causal = False self.return_softmax = False self.use_sdp_kernel = False class TestFlashAttentionAPITest3(TestFlashAttentionAPI): def setUp(self): self.place = get_device_place() self.shape = (2, 512, 8, 16) self.dtype = 'float16' self.dropout = 0.0 self.causal = True self.return_softmax = False self.use_sdp_kernel = False class TestFlashAttentionAPITest4(TestFlashAttentionAPI): def setUp(self): self.place = get_device_place() self.shape = (8, 1024, 16, 128) self.dtype = 'float16' self.dropout = 0.0 self.causal = False self.return_softmax = False self.use_sdp_kernel = False class TestFlashAttentionAPITest5(TestFlashAttentionAPI): def setUp(self): self.place = get_device_place() self.shape = ( (8, 1024, 16, 256) if (is_sm80 or is_sm90) else (8, 1024, 16, 192) ) self.dtype = 'float16' self.dropout = 0.0 self.causal = False self.return_softmax = False self.use_sdp_kernel = False class TestFlashAttentionAPITest6(TestFlashAttentionAPI): def setUp(self): self.place = get_device_place() self.shape = (0, 256, 8, 16) self.dtype = 'float16' self.dropout = 0.0 self.causal = True self.return_softmax = False self.use_sdp_kernel = False def test_unpadded(self): pass class TestMathAttentionAPITest(TestFlashAttentionAPI): def setUp(self): self.place = get_device_place() self.shape = (8, 1024, 16, 128) self.dtype = 'float16' self.dropout = 0.0 self.causal = False self.return_softmax = False self.use_sdp_kernel = True self.use_sdp_api = False self.enable_math = True self.enable_flash = False self.enable_mem_efficient = False class TestSDPAttentionAPITest(TestFlashAttentionAPI): def setUp(self): self.place = get_device_place() self.shape = (8, 1024, 16, 128) self.dtype = 'float16' self.dropout = 0.0 self.causal = False self.return_softmax = False self.use_sdp_kernel = True self.use_sdp_api = True self.enable_math = True self.enable_flash = False self.enable_mem_efficient = False class TestFlashAttentionWithMaskAPITest(TestFlashAttentionWithMaskAPI): def setUp(self): self.place = get_device_place() self.shape = (8, 1024, 16, 128) self.dtype = 'float16' self.dropout = 0.0 self.causal = False # cpu case class TestSDPAttentionWithMaskAPITest(TestFlashAttentionWithMaskAPI): def setUp(self): self.place = paddle.CPUPlace() self.shape = (8, 1024, 16, 128) self.dtype = 'float32' self.dropout = 0.0 self.causal = False # fp32 case class TestSDPAttentionWithMaskAPITest2(TestFlashAttentionWithMaskAPI): def setUp(self): self.place = get_device_place() self.shape = (8, 1024, 16, 128) self.dtype = 'float32' self.dropout = 0.0 self.causal = False # low sm case @unittest.skipIf( 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.5 or 8.x", ) class TestSDPAttentionWithMaskAPITest3(TestFlashAttentionWithMaskAPI): def setUp(self): self.place = get_device_place() self.shape = (8, 1024, 16, 128) self.dtype = 'float16' self.dropout = 0.0 self.causal = False @unittest.skipIf( 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.5 or 8.x", ) class TestSDPAttentionWithMaskAPITest4(TestFlashAttentionWithMaskAPI): def setUp(self): self.place = get_device_place() self.shape = (0, 1024, 16, 128) self.dtype = 'float32' self.dropout = 0.0 self.causal = True @unittest.skipIf( not is_flashattn_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.5 or 8.x", ) class TestFlashAttentionNoKVGrad(unittest.TestCase): def setUp(self): self.place = get_device_place() self.shape = (2, 128, 8, 16) self.dtype = 'float16' self.dropout = 0.0 self.causal = True self.return_softmax = False self.enable_math = False self.enable_flash = True self.enable_mem_efficient = False def _init_tensor_from_numpy(self, array, stop_gradient): t = paddle.to_tensor( array, place=self.place, dtype=self.dtype, stop_gradient=stop_gradient, ) return t def test_all(self): logging.info( f"Test case shape {self.shape} dtype {self.dtype} causal {self.causal}" ) # test dynamic paddle.disable_static() query = np.random.random(self.shape) key = np.random.random(self.shape) value = np.random.random(self.shape) q = self._init_tensor_from_numpy(query, stop_gradient=False) k = self._init_tensor_from_numpy(key, stop_gradient=True) v = self._init_tensor_from_numpy(value, stop_gradient=True) q_ = self._init_tensor_from_numpy(query, stop_gradient=False) k_ = self._init_tensor_from_numpy(key, stop_gradient=True) v_ = self._init_tensor_from_numpy(value, stop_gradient=True) with paddle.nn.functional.sdp_kernel( enable_math=self.enable_math, enable_flash=self.enable_flash, enable_mem_efficient=self.enable_mem_efficient, ): out = scaled_dot_product_attention( q, k, v, None, self.dropout, self.causal ) out_ = attention_naive(q_, k_, v_, self.causal) np.testing.assert_allclose(out.numpy(), out_, rtol=5e-03, atol=1e-03) out.backward() out_.backward() self.assertEqual(q.grad.shape, q.shape) self.assertEqual(q_.grad.shape, q.shape) np.testing.assert_allclose( q.grad.numpy(), q_.grad.numpy(), rtol=5e-03, atol=1e-03 ) @unittest.skipIf( not is_flashattn_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.5 or 8.x", ) class TestFlashAttentionGQA(unittest.TestCase): def setUp(self): self.batch_size = 2 self.num_head = 8 self.seq_len = 8192 self.head_dim = 128 self.num_group = 2 self.dtype = 'bfloat16' def gen_unpadded_data(self, dtype): seq_len_q = np.random.randint( low=1, high=self.seq_len, size=[self.batch_size] ) seq_len_k = np.random.randint( low=1, high=self.seq_len, size=[self.batch_size] ) cu_seqlen_q = paddle.to_tensor( [0, *np.cumsum(seq_len_q).tolist()], dtype=paddle.int32 ) cu_seqlen_k = paddle.to_tensor( [0, *np.cumsum(seq_len_k).tolist()], dtype=paddle.int32 ) qs, ks, vs = [], [], [] for i in range(self.batch_size): tmp_q = ( paddle.randn( [seq_len_q[i] * self.num_head * self.head_dim], dtype=dtype ) / 1e2 ) tmp_k = ( paddle.randn( [ seq_len_k[i] * self.num_head * self.head_dim // self.num_group ], dtype=dtype, ) / 1e2 ) tmp_v = ( paddle.randn( [ seq_len_k[i] * self.num_head * self.head_dim // self.num_group ], dtype=dtype, ) / 1e2 ) qs.append(tmp_q) ks.append(tmp_k) vs.append(tmp_v) q = paddle.concat(qs, axis=0).reshape( [-1, self.num_head, self.head_dim] ) k = paddle.concat(ks, axis=0).reshape( [-1, self.num_head // self.num_group, self.head_dim] ) v = paddle.concat(vs, axis=0).reshape( [-1, self.num_head // self.num_group, self.head_dim] ) return q, k, v, cu_seqlen_q, cu_seqlen_k def gen_test_data(self, dtype, use_unpadded): assert self.num_head % self.num_group == 0 if use_unpadded: q, k, v, cu_seqlen_q, cu_seqlen_k = self.gen_unpadded_data(dtype) else: q = ( paddle.randn( [ self.batch_size, self.seq_len, self.num_head, self.head_dim, ], dtype=dtype, ) / 1e2 ) k = ( paddle.randn( [ self.batch_size, self.seq_len, self.num_head // self.num_group, self.head_dim, ], dtype=dtype, ) / 1e2 ) v = ( paddle.randn( [ self.batch_size, self.seq_len, self.num_head // self.num_group, self.head_dim, ], dtype=dtype, ) / 1e2 ) cu_seqlen_q = None cu_seqlen_k = None out_grad = paddle.randn(q.shape, dtype=dtype) / 1e2 return q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad def clone_tensor(self, tensor): if tensor is None: return None elif isinstance(tensor, (list, tuple)): return [self.clone_tensor(t) for t in tensor] else: tensor = tensor.detach().clone() tensor.stop_gradient = False return tensor @paddle.no_grad() def convert_dtype(self, tensors): ret = [] for t in tensors: if t.dtype in [paddle.float16, paddle.bfloat16]: t = t.astype(paddle.float32) t = t.numpy() ret.append(t) return ret def calc_fa( self, q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, causal, use_unpadded ): q, k, v = self.clone_tensor([q, k, v]) if use_unpadded: scale = self.head_dim ** (-0.5) out = flash_attn_unpadded( q, k, v, cu_seqlens_q=cu_seqlen_q, cu_seqlens_k=cu_seqlen_k, max_seqlen_q=self.seq_len, max_seqlen_k=self.seq_len, scale=scale, causal=causal, ) else: out = flash_attention(q, k, v, causal=causal) out = out[0] out.backward(out_grad) return self.convert_dtype([out, q.grad, k.grad, v.grad]) def calc_raw_attn( self, q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, causal, use_unpadded ): q, k, v = self.clone_tensor([q, k, v]) if use_unpadded: qq, q_mask = self.pad(q, cu_seqlen_q, self.seq_len) kk, k_mask = self.pad(k, cu_seqlen_k, self.seq_len) vv, _ = self.pad(v, cu_seqlen_k, self.seq_len) qk_mask = paddle.matmul(q_mask, k_mask, transpose_y=True) qk_mask = qk_mask.reshape( [self.batch_size, 1, self.seq_len, self.seq_len] ) qk_mask[qk_mask == 0] = -1e6 qk_mask[qk_mask == 1] = 0 else: qq, kk, vv = q, k, v assert len(qq.shape) == 4, qq.shape assert len(kk.shape) == 4, kk.shape assert len(vv.shape) == 4, vv.shape perm = [0, 2, 1, 3] qq = paddle.transpose(qq, perm) kk = paddle.transpose(kk, perm) kk = paddle.stack([kk] * self.num_group, axis=2).reshape(qq.shape) vv = paddle.transpose(vv, perm) vv = paddle.stack([vv] * self.num_group, axis=2).reshape(qq.shape) scale = self.head_dim ** (-0.5) weight = paddle.matmul(qq * scale, kk, transpose_y=True) if use_unpadded: weight += qk_mask if causal: shape = weight.shape[-2:] mask = paddle.full(shape, -np.inf, dtype=weight.dtype) mask = paddle.triu(mask, diagonal=1) weight += mask weight = weight.astype(paddle.float32) weight = F.softmax(weight) out = paddle.matmul(weight.astype(vv.dtype), vv) out = paddle.transpose(out, perm) if use_unpadded: out = self.unpad(out, cu_seqlen_q) out.backward(out_grad) return self.convert_dtype([out, q.grad, k.grad, v.grad]) def pad(self, x, cu_seqlen, max_seqlen): cu_seqlen_cpu = cu_seqlen.numpy() split_sections = [] for i in range(len(cu_seqlen_cpu) - 1): split_sections.append(cu_seqlen_cpu[i + 1] - cu_seqlen_cpu[i]) tmp_xs = paddle.split(x, split_sections) batch_size = len(tmp_xs) tmp_masks = [] tmp_x_pads = [] for i in range(batch_size): tmp_mask = paddle.ones([max_seqlen], dtype=x.dtype) tmp_mask[split_sections[i] :] = 0 tmp_mask = tmp_mask.reshape([1, -1, 1]) tmp_masks.append(tmp_mask) tmp_shape = tmp_xs[i].shape tmp_pad = paddle.zeros( [max_seqlen - tmp_shape[0], *tmp_shape[1:]], dtype=x.dtype ) tmp_x = paddle.concat([tmp_xs[i], tmp_pad]).unsqueeze(0) tmp_x_pads.append(tmp_x) x_pad = paddle.concat(tmp_x_pads) mask = paddle.concat(tmp_masks) return x_pad, mask def unpad(self, x, cu_seqlen): cu_seqlen_cpu = cu_seqlen.numpy() xs = paddle.split(x, x.shape[0]) tmp_xs = [] for i in range(len(cu_seqlen_cpu) - 1): tmp = xs[i].squeeze(0)[: cu_seqlen_cpu[i + 1] - cu_seqlen_cpu[i]] tmp_xs.append(tmp) unpad_x = paddle.concat(tmp_xs) return unpad_x def test_main(self): # test dynamic paddle.disable_static() for causal in [False, True]: for use_unpadded in [False, True]: ( q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, ) = self.gen_test_data(self.dtype, use_unpadded) fa_out = self.calc_fa( q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, causal, use_unpadded, ) raw_out = self.calc_raw_attn( q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, causal, use_unpadded, ) assert len(fa_out) == len(raw_out) for t1, t2 in zip(fa_out, raw_out): np.testing.assert_allclose(t1, t2, atol=1e-2, rtol=1e-2) def generate_start_rows(bz, num_head, rows, cols, start_row): assert rows == cols, f"rows {rows} must be equal to cols {cols}." start_rows_list = [] for bz_idx in range(bz): for head_idx in range(num_head): start_rows = np.array([rows + 1] * cols) mask_pos = np.random.choice( cols - 1, cols - start_row, replace=False ) index = np.arange(start_row, rows) mask_pos = np.concatenate( [ mask_pos[mask_pos < index - 1], mask_pos[mask_pos >= index - 1], ] ) start_rows[mask_pos] = index start_rows_list.append(start_rows) start_rows_arr = np.array(start_rows_list).reshape([bz, num_head, rows]) return start_rows_arr def generate_mask_matrix_from_mask_indices(start_rows): bz, num_head, seq_len = start_rows.shape matrix = np.zeros((seq_len, seq_len)) matrix[np.triu_indices(seq_len, 1)] = -np.inf matrix = matrix[np.newaxis, np.newaxis, :, :] matrix = np.tile(matrix, (bz, num_head, 1, 1)) for bz_idx in range(bz): for head_idx in range(num_head): for j in range(seq_len): start_row = start_rows[bz_idx, head_idx, j] matrix[bz_idx, head_idx, start_row:, j] = -np.inf return matrix @unittest.skipIf( not is_flashattn_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.5 or 8.x", ) class TestFlashAttentionWithSparseMaskAPI(unittest.TestCase): def setUp(self): self.place = get_device_place() self.shape = (2, 128, 8, 32) self.dtype = 'float16' self.dropout = 0.0 self.causal = True def test_dot_scale_product(self): # test dynamic paddle.disable_static() query = np.random.random(self.shape) key = np.random.random(self.shape) value = np.random.random(self.shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) q_ = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) k_ = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) v_ = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) attn_mask_start_row = 48 start_row_indices = generate_start_rows( self.shape[0], self.shape[2], self.shape[1], self.shape[1], attn_mask_start_row, ) mask = generate_mask_matrix_from_mask_indices(start_row_indices) m = paddle.to_tensor( mask, place=self.place, dtype=self.dtype, stop_gradient=False ) attn_mask_start_row_indices = paddle.to_tensor( start_row_indices, dtype=paddle.int32 ) startend_row_indices = paddle.unsqueeze(attn_mask_start_row_indices, -1) out = flashmask_attention( q, k, v, startend_row_indices=startend_row_indices, dropout=self.dropout, causal=self.causal, ) out_ = attention_naive_with_mask(q_, k_, v_, m) out.backward() out_.backward() np.testing.assert_allclose(out.numpy(), out_, rtol=5e-03, atol=1e-03) class TestFlashAttentionWithSparseMaskAPITest( TestFlashAttentionWithSparseMaskAPI ): def setUp(self): self.place = get_device_place() self.shape = (8, 1024, 16, 128) self.dtype = 'float16' self.dropout = 0.0 self.causal = True class TestFlashAttentionWithSparseMaskBF16APITest( TestFlashAttentionWithSparseMaskAPI ): def setUp(self): self.place = get_device_place() self.shape = (8, 1024, 16, 128) self.dtype = 'bfloat16' self.dropout = 0.0 self.causal = True class TestFlashAttentionVarlenQKVPackedGQA(TestFlashAttentionGQA): def gen_unpadded_data(self, dtype): seq_len_q = np.random.randint( low=1, high=self.seq_len, size=[self.batch_size] ) seq_len_k = seq_len_q cu_seqlen_q = paddle.to_tensor( [0, *np.cumsum(seq_len_q).tolist()], dtype=paddle.int32 ) cu_seqlen_k = cu_seqlen_q qs, ks, vs = [], [], [] for i in range(self.batch_size): tmp_q = ( paddle.randn( [seq_len_q[i] * self.num_head * self.head_dim], dtype=dtype ) / 1e2 ) tmp_k = ( paddle.randn( [ seq_len_k[i] * self.num_head * self.head_dim // self.num_group ], dtype=dtype, ) / 1e2 ) tmp_v = ( paddle.randn( [ seq_len_k[i] * self.num_head * self.head_dim // self.num_group ], dtype=dtype, ) / 1e2 ) qs.append(tmp_q) ks.append(tmp_k) vs.append(tmp_v) q = paddle.concat(qs, axis=0).reshape( [-1, self.num_head, self.head_dim] ) k = paddle.concat(ks, axis=0).reshape( [-1, self.num_head // self.num_group, self.head_dim] ) v = paddle.concat(vs, axis=0).reshape( [-1, self.num_head // self.num_group, self.head_dim] ) return q, k, v, cu_seqlen_q, cu_seqlen_k def calc_qkvpackedfa( self, q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, causal, varlen_padded ): q, k, v = self.clone_tensor([q, k, v]) scale = self.head_dim ** (-0.5) if varlen_padded: tq = q.reshape( [ self.batch_size * self.seq_len, self.num_group, self.num_head // self.num_group, self.head_dim, ] ) tk = k.reshape( [ self.batch_size * self.seq_len, self.num_head // self.num_group, self.head_dim, ] ) tv = v.reshape( [ self.batch_size * self.seq_len, self.num_head // self.num_group, self.head_dim, ] ) kv = paddle.stack([tk, tv], axis=1) qkv = paddle.concat([tq, kv], axis=1) out = flash_attn_varlen_qkvpacked( qkv, cu_seqlens_q=cu_seqlen_q, cu_seqlens_k=cu_seqlen_k, max_seqlen_q=self.seq_len, max_seqlen_k=self.seq_len, scale=scale, causal=causal, varlen_padded=varlen_padded, ) out_grad = out_grad.reshape(out[0].shape) else: tq = q.reshape( [ 0, self.num_group, self.num_head // self.num_group, self.head_dim, ] ) kv = paddle.stack([k, v], axis=1) qkv = paddle.concat([tq, kv], axis=1) out = flash_attn_varlen_qkvpacked( qkv, cu_seqlens_q=cu_seqlen_q, cu_seqlens_k=cu_seqlen_k, max_seqlen_q=self.seq_len, max_seqlen_k=self.seq_len, scale=scale, causal=causal, varlen_padded=varlen_padded, ) out = out[0] grads = paddle.grad(outputs=out, inputs=qkv, grad_outputs=out_grad) qkvgrad = grads[0] out = out.reshape(q.shape) qgrad = qkvgrad[:, :-2].reshape(q.shape) kgrad = qkvgrad[:, -2].reshape(k.shape) vgrad = qkvgrad[:, -1].reshape(v.shape) if varlen_padded: out = self.unpad(out, cu_seqlen_q) qgrad = self.unpad(qgrad, cu_seqlen_q) kgrad = self.unpad(kgrad, cu_seqlen_k) vgrad = self.unpad(vgrad, cu_seqlen_k) return self.convert_dtype([out, qgrad, kgrad, vgrad]) def test_main(self): for causal in [False, True]: for varlen_padded in [False, True]: ( q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, ) = self.gen_test_data(self.dtype, True) if varlen_padded: q_pad, _ = self.pad(q, cu_seqlen_q, self.seq_len) k_pad, _ = self.pad(k, cu_seqlen_k, self.seq_len) v_pad, _ = self.pad(v, cu_seqlen_k, self.seq_len) out_grad_pad, _ = self.pad( out_grad, cu_seqlen_q, self.seq_len ) else: q_pad = q k_pad = k v_pad = v out_grad_pad = out_grad fa_out = self.calc_qkvpackedfa( q_pad, k_pad, v_pad, cu_seqlen_q, cu_seqlen_k, out_grad_pad, causal, varlen_padded, ) # if varlen_padded: # cu_seqlen_q = None # cu_seqlen_k = None raw_out = self.calc_raw_attn( q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, causal, True, ) assert len(fa_out) == len(raw_out) for t1, t2 in zip(fa_out, raw_out): np.testing.assert_allclose(t1, t2, atol=1e-2, rtol=1e-2) class TestFlashAttentionVarlenQKVPackedGQA2( TestFlashAttentionVarlenQKVPackedGQA ): def setUp(self): self.batch_size = 2 self.num_head = 16 self.seq_len = 2048 self.head_dim = 128 self.num_group = 4 self.dtype = 'bfloat16' class TestFlashAttentionVarlenQKVPacked(TestFlashAttentionVarlenQKVPackedGQA): def setUp(self): self.batch_size = 3 self.num_head = 7 self.seq_len = 563 self.head_dim = 64 self.num_group = 1 self.dtype = 'bfloat16' class TestFlashAttentionQKVPackedGQA(TestFlashAttentionGQA): def calc_qkvpackedfa(self, q, k, v, out_grad, causal): # q, k, v = self.clone_tensor([q, k, v]) tq = q.reshape( [ self.batch_size, self.seq_len, self.num_group, self.num_head // self.num_group, self.head_dim, ], ) kv = paddle.stack([k, v], axis=2) qkv = paddle.concat([tq, kv], axis=2) (qkv,) = self.clone_tensor([qkv]) out = flash_attn_qkvpacked(qkv, causal=causal) out = out[0] out.backward(out_grad) qkvgrad = qkv.grad qgrad = qkvgrad[:, :, :-2].reshape(q.shape) kgrad = qkvgrad[:, :, -2].reshape(k.shape) vgrad = qkvgrad[:, :, -1].reshape(v.shape) return self.convert_dtype([out, qgrad, kgrad, vgrad]) def test_main(self): for causal in [False, True]: ( q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, ) = self.gen_test_data(self.dtype, False) fa_out = self.calc_qkvpackedfa(q, k, v, out_grad, causal) raw_out = self.calc_raw_attn( q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, causal, False, ) assert len(fa_out) == len(raw_out) for t1, t2 in zip(fa_out, raw_out): np.testing.assert_allclose(t1, t2, atol=1e-2, rtol=1e-2) class TestFlashAttentionQKVPackedGQA2(TestFlashAttentionQKVPackedGQA): def setUp(self): self.batch_size = 2 self.num_head = 16 self.seq_len = 2048 self.head_dim = 128 self.num_group = 4 self.dtype = 'bfloat16' class TestFlashAttentionQKVPacked(TestFlashAttentionQKVPackedGQA): def setUp(self): self.batch_size = 3 self.num_head = 7 self.seq_len = 563 self.head_dim = 64 self.num_group = 1 self.dtype = 'bfloat16' class TestFlashAttentionVarlenQKVPackedGQADeter( TestFlashAttentionVarlenQKVPackedGQA ): def test_main(self): paddle.set_flags({'FLAGS_cudnn_deterministic': 1}) for causal in [False, True]: for varlen_padded in [False, True]: ( q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, ) = self.gen_test_data(self.dtype, True) if varlen_padded: q_pad, _ = self.pad(q, cu_seqlen_q, self.seq_len) k_pad, _ = self.pad(k, cu_seqlen_k, self.seq_len) v_pad, _ = self.pad(v, cu_seqlen_k, self.seq_len) out_grad_pad, _ = self.pad( out_grad, cu_seqlen_q, self.seq_len ) else: q_pad = q k_pad = k v_pad = v out_grad_pad = out_grad fa_out = self.calc_qkvpackedfa( q_pad, k_pad, v_pad, cu_seqlen_q, cu_seqlen_k, out_grad_pad, causal, varlen_padded, ) # cu_seqlen_q = None # cu_seqlen_k = None raw_out = self.calc_fa( q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, causal, True, ) assert len(fa_out) == len(raw_out) i = 0 for t1, t2 in zip(fa_out, raw_out): np.testing.assert_array_equal( t1, t2, err_msg=f"Tensor{i} causal={causal} varlen_padded={varlen_padded}", ) i += 1 paddle.set_flags({'FLAGS_cudnn_deterministic': 0}) # can't bit-match dk,dv now when num_group more than 2, since the sum kernel is different and sum sequence not defined # class TestFlashAttentionVarlenQKVPackedGQADeter2( # TestFlashAttentionVarlenQKVPackedGQADeter # ): # def setUp(self): # self.batch_size = 2 # self.num_head = 16 # self.seq_len = 2048 # self.head_dim = 128 # self.num_group = 4 # self.dtype = 'bfloat16' class TestFlashAttentionVarlenQKVPackedDeter( TestFlashAttentionVarlenQKVPackedGQADeter ): def setUp(self): self.batch_size = 3 self.num_head = 7 self.seq_len = 563 self.head_dim = 64 self.num_group = 1 self.dtype = 'bfloat16' class TestFlashAttentionQKVPackedGQADeter(TestFlashAttentionQKVPackedGQA): def test_main(self): paddle.set_flags({'FLAGS_cudnn_deterministic': 1}) for causal in [False, True]: ( q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, ) = self.gen_test_data(self.dtype, False) fa_out = self.calc_qkvpackedfa(q, k, v, out_grad, causal) raw_out = self.calc_fa( q, k, v, cu_seqlen_q, cu_seqlen_k, out_grad, causal, False, ) assert len(fa_out) == len(raw_out) i = 0 for t1, t2 in zip(fa_out, raw_out): np.testing.assert_array_equal( t1, t2, err_msg=f"Tensor{i} error, causal={causal}" ) i += 1 paddle.set_flags({'FLAGS_cudnn_deterministic': 0}) # can't bit-match dk,dv now when num_group more than 2, since the sum kernel is different and sum sequence not defined # class TestFlashAttentionQKVPackedDeter2(TestFlashAttentionQKVPackedGQADeter): # def setUp(self): # self.batch_size = 2 # self.num_head = 16 # self.seq_len = 2048 # self.head_dim = 128 # self.num_group = 4 # self.dtype = 'bfloat16' class TestFlashAttentionQKVPackedDeter(TestFlashAttentionQKVPackedGQADeter): def setUp(self): self.batch_size = 3 self.num_head = 7 self.seq_len = 563 self.head_dim = 64 self.num_group = 1 self.dtype = 'bfloat16' @unittest.skipIf( not is_flashattn_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 TestCalcReducedAttentionScores(unittest.TestCase): def setUp(self): self.place = get_device_place() self.batch_size = 1 self.num_head = 8 self.seqlen_q = 1024 self.seqlen_k = 10240 self.head_dim = 128 self.num_group = 1 self.dtype = 'bfloat16' def native_reduce(self, q, k): q_ref = paddle.cast(paddle.transpose(q, [0, 2, 1, 3]), 'float32') k_ref = paddle.cast(paddle.transpose(k, [0, 2, 1, 3]), 'float32') if self.num_group != 1: k_ref = paddle.stack([k_ref] * self.num_group, axis=2).reshape( [self.batch_size, self.num_head, self.seqlen_k, self.head_dim] ) scale = 1.0 / np.sqrt(q_ref.shape[-1]) product = paddle.matmul(x=q_ref, y=k_ref, transpose_y=True) product = paddle.scale(product, scale) product = product - paddle.max(product, axis=-1, keepdim=True) product = F.softmax(product, dtype='float32') product = paddle.sum(product, axis=-2, keepdim=True) return product def test_calc_reduced_attention_scores(self): paddle.disable_static() q_shape = [ self.batch_size, self.seqlen_q, self.num_head, self.head_dim, ] k_shape = [ self.batch_size, self.seqlen_k, self.num_head // self.num_group, self.head_dim, ] query = paddle.randn(q_shape) key = paddle.randn(k_shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=True ) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=True ) reduced_scores_ref = self.native_reduce(q, k) (_, _, softmax_lse, _) = paddle._C_ops.flash_attn( q, k, k, None, # fixed_seed_offset None, # attn_mask 0.0, # dropout False, # causal False, # return_softmax False, # is_test "", ) reduced_scores = calc_reduced_attention_scores(q, k, softmax_lse) np.testing.assert_allclose( reduced_scores.numpy(), reduced_scores_ref.numpy(), rtol=1e-05, atol=0, ) if self.dtype == 'float16': paddle.enable_static() with paddle.static.program_guard(paddle.static.Program()): qs = paddle.static.data( name="q", shape=q_shape, dtype=self.dtype ) ks = paddle.static.data( name="k", shape=k_shape, dtype=self.dtype ) softmax_lse_s = paddle.static.data( name="softmax_lse", shape=softmax_lse.shape, dtype='float32' ) reduced_scores = calc_reduced_attention_scores( qs, ks, softmax_lse_s ) exe = base.Executor(self.place) fetches_result = exe.run( feed={ "q": query.numpy().astype(self.dtype), "k": key.numpy().astype(self.dtype), "softmax_lse": softmax_lse.numpy(), }, fetch_list=[reduced_scores], ) np.testing.assert_allclose( fetches_result[0], reduced_scores_ref.numpy(), rtol=1e-05, atol=0, ) paddle.disable_static() @unittest.skipIf( not is_flashattn_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 TestCalcReducedAttentionScoresGQA(TestCalcReducedAttentionScores): def setUp(self): self.place = get_device_place() self.batch_size = 1 self.num_head = 8 self.seqlen_q = 1024 self.seqlen_k = 10240 self.head_dim = 128 self.num_group = 2 self.dtype = 'bfloat16' @unittest.skipIf( not is_flashattn_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 TestCalcReducedAttentionScoresFP16(TestCalcReducedAttentionScores): def setUp(self): self.place = get_device_place() self.batch_size = 1 self.num_head = 8 self.seqlen_q = 1024 self.seqlen_k = 10240 self.head_dim = 128 self.num_group = 1 self.dtype = 'float16' @unittest.skipIf( not is_flashattn_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 TestCalcReducedAttentionScoresNotEvenMN(TestCalcReducedAttentionScores): def setUp(self): self.place = get_device_place() self.batch_size = 1 self.num_head = 8 self.seqlen_q = 1023 self.seqlen_k = 10241 self.head_dim = 128 self.num_group = 1 self.dtype = 'bfloat16' @unittest.skipIf( not is_flashattn_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.5 or 8.x", ) class TestFlashAttentionAlignment(unittest.TestCase): def setUp(self): paddle.disable_static() self.bs = 1 self.seq_len = 8 self.num_head = 1 self.head_dim = 8 self.dtype = 'float16' self.query = np.array( [ # batch_size = 1 [[0.3, -0.7, 0.2, 0.5, -0.4, 0.8, -0.2, 0.1]], # seq position 0 [ [-0.5, 0.4, 0.7, -0.3, 0.6, -0.8, 0.3, -0.1] ], # seq position 1 [[0.2, 0.8, -0.4, 0.1, -0.6, 0.3, 0.7, -0.5]], # seq position 2 [[-0.8, 0.1, 0.6, 0.4, -0.2, -0.7, 0.5, 0.3]], # seq position 3 [[0.7, -0.3, -0.5, 0.8, 0.2, 0.4, -0.6, 0.1]], # seq position 4 [[-0.2, 0.5, 0.3, -0.7, 0.8, 0.1, -0.4, 0.6]], # seq position 5 [[0.4, -0.6, 0.8, -0.1, 0.3, 0.5, -0.8, 0.2]], # seq position 6 [[-0.4, 0.2, -0.8, 0.6, 0.1, -0.3, 0.7, 0.5]], # seq position 7 ], dtype=np.float16, ).reshape(1, 8, 1, 8) self.key = np.array( [ # batch_size = 1 [[0.6, -0.2, 0.8, -0.4, 0.3, 0.1, -0.7, 0.5]], # seq position 0 [[-0.3, 0.7, 0.1, 0.5, -0.8, 0.4, -0.2, 0.6]], # seq position 1 [[0.8, -0.5, 0.3, -0.1, 0.6, 0.2, -0.4, 0.7]], # seq position 2 [[-0.6, 0.4, -0.2, 0.7, 0.1, -0.8, 0.3, 0.5]], # seq position 3 [[0.2, 0.8, -0.6, 0.3, 0.5, -0.1, 0.7, -0.4]], # seq position 4 [[-0.7, 0.3, 0.5, 0.1, -0.4, 0.8, -0.2, 0.6]], # seq position 5 [[0.5, -0.8, 0.2, 0.6, -0.3, 0.7, 0.1, -0.5]], # seq position 6 [[-0.1, 0.6, 0.4, -0.7, 0.2, 0.5, -0.8, 0.3]], # seq position 7 ], dtype=np.float16, ).reshape(1, 8, 1, 8) self.value = np.array( [ # batch_size = 1 [[-0.4, 0.8, -0.1, 0.3, 0.6, -0.5, 0.2, 0.7]], # seq position 0 [[0.5, -0.3, 0.7, 0.2, -0.6, 0.4, -0.8, 0.1]], # seq position 1 [[-0.2, 0.6, 0.4, -0.7, 0.3, 0.8, -0.1, 0.5]], # seq position 2 [[0.7, -0.4, 0.1, 0.5, -0.8, 0.2, 0.6, -0.3]], # seq position 3 [[-0.5, 0.3, 0.8, -0.2, 0.4, 0.1, -0.7, 0.6]], # seq position 4 [[0.2, -0.6, 0.3, 0.7, -0.1, 0.5, -0.4, 0.8]], # seq position 5 [[-0.8, 0.1, 0.5, -0.3, 0.7, 0.4, -0.2, 0.6]], # seq position 6 [[0.3, -0.7, 0.2, 0.6, -0.4, 0.8, -0.5, 0.1]], # seq position 7 ], dtype=np.float16, ).reshape(1, 8, 1, 8) self.mask = paddle.zeros( [1, 1, self.seq_len, self.seq_len], dtype='float16' ) for i in range(self.bs): seq_len = self.seq_len mask = ( paddle.tril( paddle.ones(shape=(seq_len, seq_len), dtype=paddle.float32) ) - 1 ) self.mask[i, 0, :seq_len, :seq_len] = mask * 1e4 self.rtol = 1e-3 self.atol = 1e-3 self.expected_output_without_mask = np.array( [ [ [ [ -0.09814, 0.004566, 0.367, 0.0902, 0.09265, 0.3545, -0.2441, 0.4368, ] ], [ [ 0.02464, -0.04175, 0.339, 0.18, -0.0385, 0.3145, -0.197, 0.3508, ] ], [ [ -0.02863, -0.06235, 0.4292, 0.1333, -0.007267, 0.3306, -0.3108, 0.3796, ] ], [ [ 0.0829, -0.10266, 0.353, 0.2078, -0.1051, 0.323, -0.1888, 0.3223, ] ], [ [ -0.09283, 0.04092, 0.3728, 0.0602, 0.08417, 0.346, -0.2312, 0.4136, ] ], [ [ 0.01353, -0.035, 0.316, 0.1869, -0.01083, 0.352, -0.2344, 0.3857, ] ], [ [ -0.0946, 0.06775, 0.3074, 0.10254, 0.11365, 0.3347, -0.2047, 0.4473, ] ], [ [ 0.05087, -0.0742, 0.395, 0.1547, -0.0862, 0.3196, -0.2118, 0.3171, ] ], ] ], dtype=np.float16, ) self.expected_output = np.array( [ [ [ [ -3.9990e-01, 7.9980e-01, -9.9976e-02, 3.0005e-01, 6.0010e-01, -5.0000e-01, 1.9995e-01, 7.0020e-01, ] ], [ [ -6.1798e-03, 3.1860e-01, 2.5000e-01, 2.5610e-01, 7.5012e-02, -1.0626e-01, -2.3743e-01, 4.3750e-01, ] ], [ [ 1.0028e-01, 1.9958e-01, 4.2505e-01, 5.3787e-04, -7.5317e-02, 2.7441e-01, -3.7524e-01, 3.4985e-01, ] ], [ [ 2.9224e-01, 1.6373e-02, 2.7368e-01, 1.8188e-01, -3.0298e-01, 2.2412e-01, 3.4210e-02, 1.2610e-01, ] ], [ [ -1.6998e-02, 2.5220e-01, 3.7939e-01, -3.7048e-02, 3.0151e-02, 2.3108e-01, -1.6772e-01, 3.5327e-01, ] ], [ [ 1.1948e-02, 1.2378e-01, 3.2935e-01, 1.2390e-01, 2.6123e-02, 2.3279e-01, -1.6919e-01, 4.4019e-01, ] ], [ [ -1.6162e-01, 1.9812e-01, 3.2544e-01, 1.8021e-02, 2.0081e-01, 2.5586e-01, -1.5466e-01, 5.0635e-01, ] ], [ [ 5.0873e-02, -7.4219e-02, 3.9502e-01, 1.5466e-01, -8.6182e-02, 3.1958e-01, -2.1179e-01, 3.1714e-01, ] ], ] ], dtype=np.float16, ) def test_flash_attention(self): paddle.disable_static() query = paddle.to_tensor(self.query) key = paddle.to_tensor(self.key) value = paddle.to_tensor(self.value) mask = None with sdp_kernel( enable_flash=True, enable_math=False, enable_mem_efficient=False ): output = paddle.nn.functional.scaled_dot_product_attention( query, key, value, attn_mask=mask, dropout_p=0.0, is_causal=False, ) np.testing.assert_allclose( output.numpy(), self.expected_output_without_mask, rtol=self.rtol, atol=self.atol, err_msg='Flash attention output does not match expected values', ) def test_math_attention(self): paddle.disable_static() query = paddle.to_tensor(self.query) key = paddle.to_tensor(self.key) value = paddle.to_tensor(self.value) mask = paddle.to_tensor(self.mask) with sdp_kernel( enable_flash=False, enable_math=True, enable_mem_efficient=False ): output = paddle.nn.functional.scaled_dot_product_attention( query, key, value, attn_mask=mask, dropout_p=0.0, is_causal=False, ) np.testing.assert_allclose( output.numpy(), self.expected_output, rtol=self.rtol, atol=self.atol, err_msg='Math attention output does not match expected values', ) def test_mem_efficient_attention(self): paddle.disable_static() query = paddle.to_tensor(self.query) key = paddle.to_tensor(self.key) value = paddle.to_tensor(self.value) mask = paddle.to_tensor(self.mask) with sdp_kernel( enable_flash=False, enable_math=False, enable_mem_efficient=True ): output = paddle.nn.functional.scaled_dot_product_attention( query, key, value, attn_mask=mask, dropout_p=0.0, is_causal=False, ) np.testing.assert_allclose( output.numpy(), self.expected_output, rtol=self.rtol, atol=self.atol, err_msg='Memory efficient attention output does not match expected values', ) def test_auto_attention(self): paddle.disable_static() query = paddle.to_tensor(self.query) key = paddle.to_tensor(self.key) value = paddle.to_tensor(self.value) mask = None # auto-select the attention implementation output = paddle.nn.functional.scaled_dot_product_attention( query, key, value, attn_mask=mask, dropout_p=0.0, is_causal=False, ) np.testing.assert_allclose( output.numpy(), self.expected_output_without_mask, rtol=self.rtol, atol=self.atol, err_msg='Auto attention output does not match expected values', ) @unittest.skipIf( not is_flashattn_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.5 or 8.x", ) class TestSDPAttentionWithScale(unittest.TestCase): def setUp(self): self.place = get_device_place() self.shape = (2, 8, 8, 32) self.dtype = paddle.bfloat16 self.dropout = 0.0 self.causal = False self.scale = 0.5 self.rtol = 1e-3 self.atol = 5e-2 paddle.disable_static() def _prepare_tensors(self): """Helper to create q, k, v and reference q_, k_, v_""" query = np.random.random(self.shape) key = np.random.random(self.shape) value = np.random.random(self.shape) q = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) k = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) v = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) q_ = paddle.to_tensor( query, place=self.place, dtype=self.dtype, stop_gradient=False ) k_ = paddle.to_tensor( key, place=self.place, dtype=self.dtype, stop_gradient=False ) v_ = paddle.to_tensor( value, place=self.place, dtype=self.dtype, stop_gradient=False ) return q, k, v, q_, k_, v_ def _run_test(self, backends, attn_mask, scale, skip_grad=False): """Generic test runner""" q, k, v, q_, k_, v_ = self._prepare_tensors() with sdp_kernel(**backends): out = scaled_dot_product_attention( q, k, v, attn_mask=attn_mask, dropout_p=self.dropout, is_causal=self.causal, scale=scale, ) out_ = attention_naive_with_mask_and_scale(q_, k_, v_, attn_mask, scale) np.testing.assert_allclose( out.float().numpy(), out_.float().numpy(), rtol=self.rtol, atol=self.atol, ) if not skip_grad: out.backward() out_.backward() self.assertIsNotNone(q.grad, "q.grad is None, backward failed.") self.assertIsNotNone(k.grad, "k.grad is None, backward failed.") self.assertIsNotNone(v.grad, "v.grad is None, backward failed.") np.testing.assert_allclose( q.grad.float().numpy(), q_.grad.float().numpy(), rtol=self.rtol, atol=self.atol, ) np.testing.assert_allclose( k.grad.float().numpy(), k_.grad.float().numpy(), rtol=self.rtol, atol=self.atol, ) np.testing.assert_allclose( v.grad.float().numpy(), v_.grad.float().numpy(), rtol=self.rtol, atol=self.atol, ) def test_no_mask_with_scale_fallback(self): backends = { "enable_math": True, "enable_flash": True, "enable_mem_efficient": True, } self._run_test(backends, attn_mask=None, scale=self.scale) def test_mask_with_scale_math_only(self): backends = { "enable_math": True, "enable_flash": False, "enable_mem_efficient": False, } mask = paddle.randn( [self.shape[0], 1, self.shape[1], self.shape[1]], dtype=self.dtype, ) self._run_test(backends, attn_mask=mask, scale=self.scale) def test_mask_with_scale_full_fallback(self): backends = { "enable_math": True, "enable_flash": True, "enable_mem_efficient": True, } mask = paddle.randn( [self.shape[0], 1, self.shape[1], self.shape[1]], dtype=self.dtype, ) self._run_test(backends, attn_mask=mask, scale=self.scale) def test_mask_with_scale_none_math(self): backends = { "enable_math": True, "enable_flash": False, "enable_mem_efficient": False, } mask = paddle.randn( [self.shape[0], 1, self.shape[1], self.shape[1]], dtype=self.dtype, ) self._run_test(backends, attn_mask=mask, scale=None) if __name__ == '__main__': unittest.main()