[Sparse] add Fused Attention kernel and API for SparseCsrTensor

[zhwesky2010]

PR types

New features

PR changes

OPs

Describe

paddle.incubate.sparse.nn.functional.attention(query, key, value, sparse_mask, key_padding_mask, attn_mask)

import paddle

batch_size = 16
num_heads = 16
seq_len = 512
head_dim = 32

query = paddle.rand([batch_size, num_heads, seq_len, head_dim])
key = paddle.rand([batch_size, num_heads, seq_len, head_dim])
value = paddle.rand([batch_size, num_heads, seq_len, head_dim])

query.stop_gradient = False
key.stop_gradient = False
value.stop_gradient = False

mask = paddle.nn.functional.dropout(paddle.ones([seq_len, seq_len])).expand([batch_size*num_heads, seq_len, seq_len])
sp_mask = mask.to_sparse_csr()

kp_mask = paddle.randint(0, 2, [batch_size, seq_len]).astype('float32')
attn_mask = paddle.randint(0, 2, [seq_len, seq_len]).astype('float32')

output = paddle.incubate.sparse.nn.functional.attention(query, key, value, sp_mask, kp_mask, attn_mask)
# kp_mask, attn_mask 是optional
output = paddle.incubate.sparse.nn.functional.attention(query, key, value, sp_mask)
output.backward()

由于该API调用的 cusparseDnMatSetStridedBatch、cusparseCsrSetStridedBatch等NV接口需要在CUDA11.7上才支持，因此CI无法直接运行，本地运行单测结果如下：

[zkh2016]

GPU？

[zhwesky2010]

Done

[zkh2016]

2021->2022

[zhwesky2010]

Done