chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,179 @@
|
||||
# Copyright (c) 2021 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import paddle
|
||||
from paddle import _C_ops
|
||||
from paddle.base.layer_helper import LayerHelper
|
||||
|
||||
|
||||
def sparse_attention(
|
||||
query: paddle.Tensor,
|
||||
key: paddle.Tensor,
|
||||
value: paddle.Tensor,
|
||||
sparse_csr_offset: paddle.Tensor,
|
||||
sparse_csr_columns: paddle.Tensor,
|
||||
key_padding_mask: paddle.Tensor | None = None,
|
||||
attn_mask: paddle.Tensor | None = None,
|
||||
name: str | None = None,
|
||||
) -> paddle.Tensor:
|
||||
r"""
|
||||
This operator sparsify the Attention matrix in Transformer module
|
||||
to achieve the effect of reducing memory consumption and computation.
|
||||
The sparse layout is expressed in CSR format and contains two parameters,
|
||||
``offset`` and ``columns``. The equation is:
|
||||
|
||||
.. math::
|
||||
|
||||
result=softmax(\frac{ Q * K^T }{\sqrt{d}}) * V
|
||||
|
||||
where : ``Q``, ``K``, and ``V`` represent the three input parameters of the attention module.
|
||||
The dimensions of the three parameters are the same.
|
||||
``d`` represents the size of the last dimension of the three parameters.
|
||||
|
||||
Warning:
|
||||
This API is only used in ``CUDA 11.3`` and above versions.
|
||||
|
||||
Args:
|
||||
query(Tensor): The query tensor in the Attention module.
|
||||
4-D tensor with shape:
|
||||
[batch_size, num_heads, seq_len, head_dim].
|
||||
The dtype can be float32 and float64.
|
||||
key(Tensor): The key tensor in the Attention module.
|
||||
4-D tensor with shape:
|
||||
[batch_size, num_heads, seq_len, head_dim].
|
||||
The dtype can be float32 and float64.
|
||||
value(Tensor): The value tensor in the Attention module.
|
||||
4-D tensor with shape:
|
||||
[batch_size, num_heads, seq_len, head_dim].
|
||||
The dtype can be float32 and float64.
|
||||
sparse_csr_offset(Tensor): The sparsity feature in the Attention module
|
||||
is expressed in the CSR format, and the offset represents
|
||||
the number of non-zero elements in each row of the matrix.
|
||||
3-D tensor with shape:
|
||||
[batch_size, num_heads, seq_len + 1].
|
||||
The dtype should be int32.
|
||||
sparse_csr_columns(Tensor): The sparsity feature in the Attention module
|
||||
is expressed in the CSR format, and the columns represent
|
||||
the column index values of non-zero elements in the matrix.
|
||||
3-D tensor with shape:
|
||||
[batch_size, num_heads, sparse_nnz].
|
||||
The dtype should be int32.
|
||||
key_padding_mask(Tensor|None, optional):The key padding mask tensor in the Attention module.
|
||||
2-D tensor with shape: [batch_size, seq_len].
|
||||
The dtype can be float32 and float64.
|
||||
A value of 0 means that the position is masked.
|
||||
attn_mask(Tensor|None, optional):The attention mask tensor in the Attention module.
|
||||
2-D tensor with shape: [seq_len, seq_len].
|
||||
The dtype can be float32 and float64.
|
||||
A value of 0 means that the position is masked.
|
||||
name(str|None, optional): The default value is None. Normally there is no need for user
|
||||
to set this property. For more information, please refer to
|
||||
:ref:`api_guide_Name`.
|
||||
|
||||
Returns:
|
||||
Tensor, 4-D tensor with shape:
|
||||
[batch_size, num_heads, seq_len, head_dim].
|
||||
The dtype can be float32 or float64.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('This API is only used in CUDA11.3 and above.')
|
||||
>>> import paddle
|
||||
|
||||
>>> paddle.disable_static()
|
||||
|
||||
>>> # `query`, `key` and `value` all have shape [1, 1, 4, 2]
|
||||
>>> query = paddle.to_tensor(
|
||||
... [[[[0, 1], [2, 3], [0, 1], [2, 3]]]],
|
||||
... dtype="float32",
|
||||
... )
|
||||
>>> key = paddle.to_tensor([[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32")
|
||||
>>> value = paddle.to_tensor([[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32")
|
||||
>>> offset = paddle.to_tensor([[[0, 2, 4, 6, 8]]], dtype="int32")
|
||||
>>> columns = paddle.to_tensor([[[0, 1, 0, 1, 2, 3, 2, 3]]], dtype="int32")
|
||||
>>> print(offset.shape)
|
||||
paddle.Size([1, 1, 5])
|
||||
>>> print(columns.shape)
|
||||
paddle.Size([1, 1, 8])
|
||||
...
|
||||
>>> key_padding_mask = paddle.to_tensor([[1, 1, 1, 0]], dtype="float32")
|
||||
>>> attention_mask = paddle.to_tensor(
|
||||
... [
|
||||
... [1, 0, 1, 1],
|
||||
... [1, 1, 1, 1],
|
||||
... [1, 1, 1, 1],
|
||||
... [1, 1, 1, 1],
|
||||
... ],
|
||||
... dtype="float32",
|
||||
... )
|
||||
>>> output_mask = paddle.nn.functional.sparse_attention(
|
||||
... query,
|
||||
... key,
|
||||
... value,
|
||||
... offset,
|
||||
... columns,
|
||||
... key_padding_mask=key_padding_mask,
|
||||
... attn_mask=attention_mask,
|
||||
... )
|
||||
>>> print(output_mask)
|
||||
Tensor(shape=[1, 1, 4, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
|
||||
[[[[0. , 1. ],
|
||||
[1.99830270, 2.99830270],
|
||||
[0. , 1. ],
|
||||
[0. , 1. ]]]])
|
||||
|
||||
>>> output = paddle.nn.functional.sparse_attention(query, key, value, offset, columns)
|
||||
>>> print(output)
|
||||
Tensor(shape=[1, 1, 4, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
|
||||
[[[[1.60885942, 2.60885954],
|
||||
[1.99830270, 2.99830270],
|
||||
[1.60885942, 2.60885954],
|
||||
[1.99830270, 2.99830270]]]])
|
||||
"""
|
||||
if paddle.framework.in_dynamic_or_pir_mode():
|
||||
res = _C_ops.sparse_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
sparse_csr_offset,
|
||||
sparse_csr_columns,
|
||||
key_padding_mask,
|
||||
attn_mask,
|
||||
)
|
||||
return res
|
||||
|
||||
helper = LayerHelper('sparse_attention', **locals())
|
||||
dtype = helper.input_dtype(input_param_name='Q')
|
||||
out = helper.create_variable_for_type_inference(dtype)
|
||||
result_sdd = helper.create_variable_for_type_inference(dtype)
|
||||
result_softmax = helper.create_variable_for_type_inference(dtype)
|
||||
inputs = {
|
||||
'Q': query,
|
||||
'K': key,
|
||||
'V': value,
|
||||
'Offset': sparse_csr_offset,
|
||||
'Columns': sparse_csr_columns,
|
||||
'KeyPaddingMask': key_padding_mask,
|
||||
'AttnMask': attn_mask,
|
||||
}
|
||||
outputs = {
|
||||
'Out': out,
|
||||
'SparseDotSdd': result_sdd,
|
||||
'Softmax': result_softmax,
|
||||
}
|
||||
helper.append_op(type='sparse_attention', inputs=inputs, outputs=outputs)
|
||||
return out
|
||||
Reference in New Issue
Block a user