wehub-resource-sync
2305 lines
87 KiB
Python
2305 lines
87 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Literal, overload
|
|
|
|
import paddle
|
|
import paddle.nn.functional as F
|
|
from paddle import _C_ops
|
|
from paddle.base.framework import in_dynamic_or_pir_mode
|
|
from paddle.base.layer_helper import LayerHelper
|
|
from paddle.base.wrapped_decorator import signature_safe_contextmanager
|
|
from paddle.nn.attention.sdpa import (
|
|
SDPBackend,
|
|
_get_enabled_backends,
|
|
sdpa_kernel,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Generator
|
|
|
|
from paddle import Tensor
|
|
from paddle.distributed.communication.group import Group
|
|
|
|
|
|
@signature_safe_contextmanager
|
|
def sdp_kernel(
|
|
enable_math: bool = False,
|
|
enable_flash: bool = True,
|
|
enable_mem_efficient: bool = True,
|
|
) -> Generator[None, None, None]:
|
|
r"""
|
|
With the sdp_kernel context manager, different algorithm implementations can
|
|
be selected for scaled_dot_product_attention.
|
|
"""
|
|
backend_list = []
|
|
if enable_flash:
|
|
backend_list.append(SDPBackend.FLASH_ATTENTION)
|
|
if enable_mem_efficient:
|
|
backend_list.append(SDPBackend.EFFICIENT_ATTENTION)
|
|
if enable_math:
|
|
backend_list.append(SDPBackend.MATH)
|
|
|
|
if not backend_list:
|
|
raise ValueError("At least one backend must be enabled")
|
|
|
|
with sdpa_kernel(backend_list) as context:
|
|
try:
|
|
yield context
|
|
finally:
|
|
pass
|
|
|
|
|
|
# special for XPU device
|
|
def get_triangle_upper_mask(x: Tensor) -> Tensor:
|
|
mask = paddle.full_like(x, -1e4)
|
|
mask.stop_gradient = True
|
|
mask = paddle.triu(mask, diagonal=1)
|
|
mask.stop_gradient = True
|
|
return mask
|
|
|
|
|
|
@overload
|
|
def _math_attention(
|
|
query: Tensor,
|
|
key: Tensor,
|
|
value: Tensor,
|
|
mask: Tensor,
|
|
dropout_rate: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: Literal[False] = ...,
|
|
training: bool = ...,
|
|
scale: float | None = ...,
|
|
) -> tuple[Tensor, None]: ...
|
|
|
|
|
|
@overload
|
|
def _math_attention(
|
|
query: Tensor,
|
|
key: Tensor,
|
|
value: Tensor,
|
|
mask: Tensor,
|
|
dropout_rate: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: Literal[True] = ...,
|
|
training: bool = ...,
|
|
scale: float | None = ...,
|
|
) -> tuple[Tensor, Tensor]: ...
|
|
|
|
|
|
@overload
|
|
def _math_attention(
|
|
query: Tensor,
|
|
key: Tensor,
|
|
value: Tensor,
|
|
mask: Tensor,
|
|
dropout_rate: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: bool = ...,
|
|
training: bool = ...,
|
|
scale: float | None = ...,
|
|
) -> tuple[Tensor, Tensor | None]: ...
|
|
|
|
|
|
def _math_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
mask=None,
|
|
dropout_rate=0.0,
|
|
causal=False,
|
|
return_softmax=False,
|
|
training=True,
|
|
scale=None,
|
|
):
|
|
r"""
|
|
This is a basic implementation of scaled dot product attention composed of
|
|
combinations of fundamental components.
|
|
"""
|
|
|
|
head_dim = query.shape[-1]
|
|
query = paddle.transpose(query, [0, 2, 1, 3])
|
|
key = paddle.transpose(key, [0, 2, 1, 3])
|
|
value = paddle.transpose(value, [0, 2, 1, 3])
|
|
# head_dim may be 0 in zero size case
|
|
scale = scale or (head_dim**-0.5 if head_dim != 0 else 1.0)
|
|
product = paddle.matmul(x=query * scale, y=key, transpose_y=True)
|
|
|
|
if not causal:
|
|
if mask is not None:
|
|
product = product + mask
|
|
weights = F.softmax(product)
|
|
else:
|
|
# special for XPU device
|
|
place = paddle.get_device()
|
|
if (
|
|
"xpu" in place
|
|
or "cpu" in place
|
|
or product.shape[-1] < 32
|
|
or product.shape[-1] > 16384
|
|
or product.shape[-1] != product.shape[-2]
|
|
):
|
|
# softmax_mask_fuse_upper_triangle is not supported on XPU, use plain implementation
|
|
mask = get_triangle_upper_mask(product)
|
|
product = product + mask
|
|
weights = F.softmax(product)
|
|
else:
|
|
weights = paddle.incubate.softmax_mask_fuse_upper_triangle(product)
|
|
if dropout_rate > 0.0:
|
|
weights = F.dropout(
|
|
weights, dropout_rate, training=training, mode="upscale_in_train"
|
|
)
|
|
|
|
out = paddle.matmul(weights, value)
|
|
out = paddle.transpose(out, [0, 2, 1, 3])
|
|
return out, weights if return_softmax else None
|
|
|
|
|
|
def _select_sdp_cuda(head_dim: int) -> str:
|
|
if head_dim <= 256:
|
|
return "flash_attn"
|
|
else:
|
|
return "mem_efficient"
|
|
|
|
|
|
def _select_sdp(head_dim: int) -> str:
|
|
r"""
|
|
There are currently three different implementation options available for
|
|
scaled dot product attention, and the chosen approach depends on whether it
|
|
is determined by the sdp_kernel configuration or specified through input values.
|
|
"""
|
|
place = paddle.get_device()
|
|
|
|
if "xpu" in place:
|
|
return "flash_attn"
|
|
|
|
enabled_backends = _get_enabled_backends()
|
|
if not enabled_backends:
|
|
raise AssertionError(
|
|
"No available backend for scaled_dot_product_attention was found."
|
|
)
|
|
|
|
enable_math = SDPBackend.MATH in enabled_backends
|
|
enable_flash = SDPBackend.FLASH_ATTENTION in enabled_backends
|
|
enable_mem_efficient = SDPBackend.EFFICIENT_ATTENTION in enabled_backends
|
|
|
|
if enable_math is True:
|
|
if enable_flash is False and enable_mem_efficient is False:
|
|
return "math"
|
|
if "gpu" not in place:
|
|
return "math"
|
|
if enable_flash is True and enable_mem_efficient is True:
|
|
return _select_sdp_cuda(head_dim)
|
|
if enable_flash is True:
|
|
return "flash_attn"
|
|
return "mem_efficient"
|
|
|
|
|
|
@overload
|
|
def flash_attention(
|
|
query: Tensor,
|
|
key: Tensor,
|
|
value: Tensor,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: Literal[False] = ...,
|
|
*,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
softmax_scale: float | None = ...,
|
|
) -> tuple[Tensor, None]: ...
|
|
|
|
|
|
@overload
|
|
def flash_attention(
|
|
query: Tensor,
|
|
key: Tensor,
|
|
value: Tensor,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: Literal[True] = ...,
|
|
*,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
softmax_scale: float | None = ...,
|
|
) -> tuple[Tensor, Tensor]: ...
|
|
|
|
|
|
@overload
|
|
def flash_attention(
|
|
query: Tensor,
|
|
key: Tensor,
|
|
value: Tensor,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: bool = ...,
|
|
*,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
softmax_scale: float | None = ...,
|
|
) -> tuple[Tensor, Tensor | None]: ...
|
|
|
|
|
|
def flash_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
dropout=0.0,
|
|
causal=False,
|
|
return_softmax=False,
|
|
*,
|
|
fixed_seed_offset=None,
|
|
rng_name="",
|
|
training=True,
|
|
name=None,
|
|
softmax_scale=None,
|
|
):
|
|
r"""
|
|
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 support inputs with dtype float16 and bfloat16.
|
|
|
|
Args:
|
|
query(Tensor): The query tensor in the Attention module.
|
|
4-D tensor with shape:
|
|
[batch_size, seq_len, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
key(Tensor): The key tensor in the Attention module.
|
|
4-D tensor with shape:
|
|
[batch_size, seq_len, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
value(Tensor): The value tensor in the Attention module.
|
|
4-D tensor with shape:
|
|
[batch_size, seq_len, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
dropout(float): The dropout ratio.
|
|
causal(bool): Whether enable causal mode.
|
|
return_softmax(bool): Whether to return softmax.
|
|
fixed_seed_offset(Tensor|None, optional): With fixed seed, offset for dropout mask.
|
|
training(bool): Whether it is in the training phase.
|
|
rng_name(str): The name to select Generator.
|
|
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:
|
|
out(Tensor): The attention tensor.
|
|
4-D tensor with shape: [batch_size, seq_len, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
softmax(Tensor): The softmax tensor. None if return_softmax is False.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
|
|
>>> paddle.seed(2023)
|
|
>>> q = paddle.rand((1, 128, 2, 16))
|
|
|
|
>>> output = paddle.nn.functional.flash_attention.flash_attention(q, q, q, 0.9, False, False)
|
|
>>> print(output)
|
|
(Tensor(shape=[1, 128, 2, 16], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
[[[[0.34992969, 0.34456205, 0.45826620, ..., 0.39883569,
|
|
0.42132437, 0.39157745],
|
|
[0.76687670, 0.65837246, 0.69117945, ..., 0.82817292,
|
|
0.76690865, 0.71485817]],
|
|
[[0.27992037, 0.45855168, 0.36554155, ..., 0.43579611,
|
|
0.32732859, 0.32411280],
|
|
[0.56813753, 0.49862429, 0.60471594, ..., 0.66300118,
|
|
0.63945228, 0.61648899]],
|
|
[[0.66428208, 0.69490069, 0.67286366, ..., 0.73747700,
|
|
0.57736880, 0.89188176],
|
|
[0.31396604, 0.27845621, 0.23340687, ..., 0.17776486,
|
|
0.29530025, 0.21338812]],
|
|
...,
|
|
[[0.72750765, 0.68531626, 0.84076148, ..., 0.63581419,
|
|
0.66560036, 0.74001575],
|
|
[0.51681775, 0.48662624, 0.47082719, ..., 0.41591716,
|
|
0.45305789, 0.46688002]],
|
|
[[0.52381468, 0.45015901, 0.49682677, ..., 0.49888846,
|
|
0.46920198, 0.37457168],
|
|
[0.56389368, 0.53894317, 0.57543570, ..., 0.48397997,
|
|
0.58046442, 0.55370426]],
|
|
[[0.71662450, 0.57275224, 0.57053083, ..., 0.48108250,
|
|
0.53336459, 0.54540110],
|
|
[0.59137988, 0.51350886, 0.50449550, ..., 0.38860252,
|
|
0.40526700, 0.60541761]]]]), None)
|
|
|
|
"""
|
|
head_dim = query.shape[3]
|
|
sdp_func_name = _select_sdp(head_dim)
|
|
|
|
if sdp_func_name == "flash_attn":
|
|
if "xpu" in paddle.get_device():
|
|
fa_version = 2
|
|
elif "iluvatar_gpu" in paddle.get_device():
|
|
fa_version = 2
|
|
elif paddle.get_flags(["FLAGS_cudnn_deterministic"])[
|
|
"FLAGS_cudnn_deterministic"
|
|
]:
|
|
fa_version = 2
|
|
else:
|
|
fa_version = paddle.base.framework.get_flags(
|
|
["FLAGS_flash_attn_version"]
|
|
)["FLAGS_flash_attn_version"]
|
|
assert in_dynamic_or_pir_mode() or fa_version == 2, (
|
|
"flash attention 3 only support dynamic or pir mode"
|
|
)
|
|
assert dropout == 0.0 or fa_version == 2, (
|
|
"flash attention 3 does not support dropout"
|
|
)
|
|
assert not return_softmax or fa_version == 2, (
|
|
"flash attention 3 does not support return softmax"
|
|
)
|
|
assert fixed_seed_offset is None or fa_version == 2, (
|
|
"flash attention 3 does not support setting seed_offset"
|
|
)
|
|
assert rng_name == "" or fa_version == 2, (
|
|
"flash attention 3 does not support setting rng_name"
|
|
)
|
|
assert training or fa_version == 2, (
|
|
"flash attention 3 does not support setting training"
|
|
)
|
|
assert name is None or fa_version == 2, (
|
|
"flash attention 3 does not support setting name"
|
|
)
|
|
assert softmax_scale is None or fa_version == 3, (
|
|
"flash attention 2 does not support setting softmax_scale"
|
|
)
|
|
if in_dynamic_or_pir_mode():
|
|
if fa_version == 2:
|
|
(result_attention, result_softmax, _, _) = _C_ops.flash_attn(
|
|
query,
|
|
key,
|
|
value,
|
|
fixed_seed_offset,
|
|
None,
|
|
dropout,
|
|
causal,
|
|
return_softmax,
|
|
not training,
|
|
rng_name,
|
|
)
|
|
return result_attention, (
|
|
result_softmax if return_softmax else None
|
|
)
|
|
elif fa_version == 3:
|
|
if softmax_scale is None:
|
|
softmax_scale = query.shape[-1] ** (-0.5)
|
|
|
|
out, softmax_lse = _C_ops.flash_attn_v3(
|
|
query,
|
|
key,
|
|
value,
|
|
None, # q_v_
|
|
None, # q_descale_
|
|
None, # k_descale_
|
|
None, # v_descale_
|
|
softmax_scale,
|
|
causal,
|
|
-1, # window_size_left
|
|
-1, # window_size_right
|
|
0.0, # softcap
|
|
1, # num_splits
|
|
False, # manual_set_pack_gqa
|
|
False, # pack_gqa_
|
|
0, # sm_margin
|
|
)
|
|
return out, None # return_softmax
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid flash attention version: {fa_version}"
|
|
)
|
|
|
|
helper = LayerHelper('flash_attn', **locals())
|
|
dtype = helper.input_dtype(input_param_name='q')
|
|
out = helper.create_variable_for_type_inference(dtype)
|
|
softmax = helper.create_variable_for_type_inference(dtype)
|
|
softmax_lse = helper.create_variable_for_type_inference(paddle.float32)
|
|
seed_offset = helper.create_variable_for_type_inference(paddle.int64)
|
|
inputs = {
|
|
'q': query,
|
|
'k': key,
|
|
'v': value,
|
|
'fixed_seed_offset': fixed_seed_offset,
|
|
}
|
|
outputs = {
|
|
'out': out,
|
|
'softmax': softmax,
|
|
'softmax_lse': softmax_lse,
|
|
'seed_offset': seed_offset,
|
|
}
|
|
helper.append_op(
|
|
type='flash_attn',
|
|
inputs=inputs,
|
|
outputs=outputs,
|
|
attrs={
|
|
'dropout': dropout,
|
|
'causal': causal,
|
|
'return_softmax': return_softmax,
|
|
'is_test': not training,
|
|
'rng_name': rng_name,
|
|
},
|
|
)
|
|
return out, softmax if return_softmax else None
|
|
else:
|
|
if sdp_func_name == "mem_efficient":
|
|
from paddle.incubate.nn.memory_efficient_attention import (
|
|
memory_efficient_attention,
|
|
)
|
|
|
|
output = memory_efficient_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
attn_bias=None,
|
|
p=dropout,
|
|
scale=None,
|
|
training=training,
|
|
)
|
|
return output, None
|
|
else:
|
|
return _math_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
dropout_rate=dropout,
|
|
causal=causal,
|
|
return_softmax=return_softmax,
|
|
training=training,
|
|
)
|
|
|
|
|
|
@overload
|
|
def flash_attn_qkvpacked(
|
|
qkv: Tensor,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: Literal[False] = ...,
|
|
*,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
) -> tuple[Tensor, None]: ...
|
|
|
|
|
|
@overload
|
|
def flash_attn_qkvpacked(
|
|
qkv: Tensor,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: Literal[True] = ...,
|
|
*,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
) -> tuple[Tensor, Tensor]: ...
|
|
|
|
|
|
@overload
|
|
def flash_attn_qkvpacked(
|
|
qkv: Tensor,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: bool = ...,
|
|
*,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
) -> tuple[Tensor, Tensor | None]: ...
|
|
|
|
|
|
def flash_attn_qkvpacked(
|
|
qkv,
|
|
dropout=0.0,
|
|
causal=False,
|
|
return_softmax=False,
|
|
*,
|
|
fixed_seed_offset=None,
|
|
rng_name="",
|
|
training=True,
|
|
name=None,
|
|
):
|
|
r"""
|
|
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 only supports inputs with dtype float16 and bfloat16.
|
|
Don't call this API if flash_attn is not supported.
|
|
|
|
Args:
|
|
qkv(Tensor): The query/key/value packed tensor in the Attention module.
|
|
5-D tensor with shape:
|
|
[batchsize, seqlen , num_heads/num_heads_k + 2, num_heads_k, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
dropout(float): The dropout ratio.
|
|
causal(bool): Whether enable causal mode.
|
|
return_softmax(bool): Whether to return softmax.
|
|
fixed_seed_offset(Tensor|None, optional): With fixed seed, offset for dropout mask.
|
|
training(bool): Whether it is in the training phase.
|
|
rng_name(str): The name to select Generator.
|
|
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:
|
|
- out(Tensor). The attention tensor. 4-D tensor with shape: [batch_size, seq_len, num_heads, head_dim]. The dtype can be float16 or bfloat16.
|
|
- softmax(Tensor). The softmax tensor. None if return_softmax is False.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +SKIP('flash_attn need A100 compile')
|
|
>>> import paddle
|
|
|
|
>>> paddle.seed(2023)
|
|
>>> q = paddle.rand((1, 128, 2, 16))
|
|
>>> qkv = paddle.stack([q, q, q], axis=2)
|
|
>>> output = paddle.nn.functional.flash_attn_qkvpacked(qkv, 0.9, False, False)
|
|
>>> print(output)
|
|
(Tensor(shape=[1, 128, 2, 16], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
[[[[0.34992966, 0.34456208, 0.45826620, ..., 0.39883569,
|
|
0.42132431, 0.39157745],
|
|
[0.76687670, 0.65837246, 0.69117945, ..., 0.82817286,
|
|
0.76690865, 0.71485823]],
|
|
...,
|
|
[[0.71662450, 0.57275224, 0.57053083, ..., 0.48108247,
|
|
0.53336465, 0.54540104],
|
|
[0.59137970, 0.51350880, 0.50449550, ..., 0.38860250,
|
|
0.40526697, 0.60541755]]]]), None)
|
|
>>> # doctest: -SKIP
|
|
|
|
"""
|
|
head_dim = qkv.shape[-1]
|
|
sdp_func_name = _select_sdp(head_dim)
|
|
|
|
if sdp_func_name == "flash_attn":
|
|
if in_dynamic_or_pir_mode():
|
|
(
|
|
result_attention,
|
|
result_softmax,
|
|
_,
|
|
_,
|
|
) = _C_ops.flash_attn_qkvpacked(
|
|
qkv,
|
|
fixed_seed_offset,
|
|
None,
|
|
dropout,
|
|
causal,
|
|
return_softmax,
|
|
not training,
|
|
rng_name,
|
|
)
|
|
return result_attention, result_softmax if return_softmax else None
|
|
|
|
helper = LayerHelper('flash_attn_qkvpacked', **locals())
|
|
dtype = helper.input_dtype(input_param_name='qkv')
|
|
out = helper.create_variable_for_type_inference(dtype)
|
|
softmax = helper.create_variable_for_type_inference(dtype)
|
|
softmax_lse = helper.create_variable_for_type_inference(paddle.float32)
|
|
seed_offset = helper.create_variable_for_type_inference(paddle.int64)
|
|
inputs = {
|
|
'qkv': qkv,
|
|
'fixed_seed_offset': fixed_seed_offset,
|
|
}
|
|
outputs = {
|
|
'out': out,
|
|
'softmax': softmax,
|
|
'softmax_lse': softmax_lse,
|
|
'seed_offset': seed_offset,
|
|
}
|
|
helper.append_op(
|
|
type='flash_attn_qkvpacked',
|
|
inputs=inputs,
|
|
outputs=outputs,
|
|
attrs={
|
|
'dropout': dropout,
|
|
'causal': causal,
|
|
'return_softmax': return_softmax,
|
|
'is_test': not training,
|
|
'rng_name': rng_name,
|
|
},
|
|
)
|
|
return out, softmax if return_softmax else None
|
|
else:
|
|
# don't call qkvpacked if not using flash_attn
|
|
query = qkv[:, :, :-2].reshape([0, 0, -1, qkv.shape[-1]])
|
|
key = qkv[:, :, -2]
|
|
value = qkv[:, :, -1]
|
|
if sdp_func_name == "mem_efficient":
|
|
from paddle.incubate.nn.memory_efficient_attention import (
|
|
memory_efficient_attention,
|
|
)
|
|
|
|
output = memory_efficient_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
attn_bias=None,
|
|
p=dropout,
|
|
scale=None,
|
|
training=training,
|
|
)
|
|
return output, None
|
|
else:
|
|
return _math_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
dropout_rate=dropout,
|
|
causal=causal,
|
|
return_softmax=return_softmax,
|
|
training=training,
|
|
)
|
|
|
|
|
|
@overload
|
|
def flash_attn_unpadded(
|
|
query: Tensor,
|
|
key: Tensor,
|
|
value: Tensor,
|
|
cu_seqlens_q: Tensor,
|
|
cu_seqlens_k: Tensor,
|
|
max_seqlen_q: int,
|
|
max_seqlen_k: int,
|
|
scale: float,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: Literal[False] = ...,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
) -> tuple[Tensor, None]: ...
|
|
|
|
|
|
@overload
|
|
def flash_attn_unpadded(
|
|
query: Tensor,
|
|
key: Tensor,
|
|
value: Tensor,
|
|
cu_seqlens_q: Tensor,
|
|
cu_seqlens_k: Tensor,
|
|
max_seqlen_q: int,
|
|
max_seqlen_k: int,
|
|
scale: float,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: Literal[True] = ...,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
) -> tuple[Tensor, Tensor]: ...
|
|
|
|
|
|
@overload
|
|
def flash_attn_unpadded(
|
|
query: Tensor,
|
|
key: Tensor,
|
|
value: Tensor,
|
|
cu_seqlens_q: Tensor,
|
|
cu_seqlens_k: Tensor,
|
|
max_seqlen_q: int,
|
|
max_seqlen_k: int,
|
|
scale: float,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: bool = ...,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
) -> tuple[Tensor, Tensor | None]: ...
|
|
|
|
|
|
def flash_attn_unpadded(
|
|
query,
|
|
key,
|
|
value,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
scale,
|
|
dropout=0.0,
|
|
causal=False,
|
|
return_softmax=False,
|
|
fixed_seed_offset=None,
|
|
rng_name='',
|
|
training=True,
|
|
name=None,
|
|
):
|
|
r"""
|
|
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 support inputs with dtype float16 and bfloat16.
|
|
|
|
Args:
|
|
query(Tensor): The query tensor in the Attention module.
|
|
3-D tensor with shape:
|
|
[total_seq_len, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
key(Tensor): The key tensor in the Attention module.
|
|
3-D tensor with shape:
|
|
[total_seq_len, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
value(Tensor): The value tensor in the Attention module.
|
|
3-D tensor with shape:
|
|
[total_seq_len, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
cu_seqlens_q(Tensor): The cumulative sequence lengths of the sequences in the batch,
|
|
used to index query.
|
|
cu_seqlens_k(Tensor): The cumulative sequence lengths of the sequences in the batch,
|
|
used to index key and value.
|
|
max_seqlen_q(int): Maximum sequence length of query in the batch.
|
|
max_seqlen_k(int): Maximum sequence length of key/value in the batch.
|
|
scale(float): The scaling of QK^T before applying softmax.
|
|
dropout(float, optional): The dropout ratio.
|
|
causal(bool, optional): Whether enable causal mode.
|
|
return_softmax(bool, optional): Whether to return softmax.
|
|
fixed_seed_offset(Tensor|None, optional): With fixed seed, offset for dropout mask.
|
|
rng_name(str, optional): The name to select Generator.
|
|
training(bool, optional): Whether it is in the training phase.
|
|
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:
|
|
out(Tensor): The attention tensor.
|
|
3-D tensor with shape: [total_seq_len, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
softmax(Tensor): The softmax tensor. None if return_softmax is False.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> paddle.seed(2023)
|
|
>>> q = paddle.rand((2, 128, 8, 16), dtype='float16')
|
|
>>> cu = paddle.arange(0, 384, 128, dtype='int32')
|
|
>>> qq = paddle.reshape(q, [256, 8, 16])
|
|
>>> output = paddle.nn.functional.flash_attention.flash_attn_unpadded(qq, qq, qq, cu, cu, 128, 128, 0.25, 0.0, False, False)
|
|
|
|
"""
|
|
if in_dynamic_or_pir_mode():
|
|
(
|
|
result_attention,
|
|
result_softmax,
|
|
) = _C_ops.flash_attn_unpadded(
|
|
query,
|
|
key,
|
|
value,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
fixed_seed_offset,
|
|
None,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
scale,
|
|
dropout,
|
|
causal,
|
|
return_softmax,
|
|
not training,
|
|
rng_name,
|
|
)
|
|
return result_attention, result_softmax if return_softmax else None
|
|
|
|
helper = LayerHelper('flash_attn_unpadded', **locals())
|
|
dtype = helper.input_dtype(input_param_name='q')
|
|
out = helper.create_variable_for_type_inference(dtype)
|
|
softmax = helper.create_variable_for_type_inference(dtype)
|
|
softmax_lse = helper.create_variable_for_type_inference(paddle.float32)
|
|
seed_offset = helper.create_variable_for_type_inference(paddle.int64)
|
|
inputs = {
|
|
'q': query,
|
|
'k': key,
|
|
'v': value,
|
|
'cu_seqlens_q': cu_seqlens_q,
|
|
'cu_seqlens_k': cu_seqlens_k,
|
|
'fixed_seed_offset': fixed_seed_offset,
|
|
}
|
|
outputs = {
|
|
'out': out,
|
|
'softmax': softmax,
|
|
'softmax_lse': softmax_lse,
|
|
'seed_offset': seed_offset,
|
|
}
|
|
helper.append_op(
|
|
type='flash_attn_unpadded',
|
|
inputs=inputs,
|
|
outputs=outputs,
|
|
attrs={
|
|
'max_seqlen_q': max_seqlen_q,
|
|
'max_seqlen_k': max_seqlen_k,
|
|
'scale': scale,
|
|
'dropout': dropout,
|
|
'causal': causal,
|
|
'return_softmax': return_softmax,
|
|
'is_test': not training,
|
|
'rng_name': rng_name,
|
|
},
|
|
)
|
|
return out, softmax if return_softmax else None
|
|
|
|
|
|
def flash_attention_v3_varlen(
|
|
query,
|
|
key,
|
|
value,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
seqused_q=None,
|
|
seqused_k=None,
|
|
softmax_scale=None,
|
|
causal=False,
|
|
qv=None,
|
|
q_descale=None,
|
|
k_descale=None,
|
|
v_descale=None,
|
|
window_size=(-1, -1),
|
|
softcap=0.0,
|
|
num_splits=1,
|
|
pack_gqa=None,
|
|
sm_margin=0,
|
|
):
|
|
return flash_attn_varlen_func(
|
|
query,
|
|
key,
|
|
value,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
seqused_q,
|
|
seqused_k,
|
|
softmax_scale,
|
|
causal,
|
|
qv,
|
|
q_descale,
|
|
k_descale,
|
|
v_descale,
|
|
window_size,
|
|
softcap,
|
|
num_splits,
|
|
pack_gqa,
|
|
sm_margin,
|
|
)
|
|
|
|
|
|
def flash_attn_varlen_func(
|
|
query,
|
|
key,
|
|
value,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
seqused_q=None,
|
|
seqused_k=None,
|
|
softmax_scale=None,
|
|
causal=False,
|
|
qv=None,
|
|
q_descale=None,
|
|
k_descale=None,
|
|
v_descale=None,
|
|
window_size=(-1, -1),
|
|
softcap=0.0,
|
|
num_splits=1,
|
|
pack_gqa=None,
|
|
sm_margin=0,
|
|
):
|
|
r"""
|
|
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.
|
|
This is the varlen version of flash attention.
|
|
|
|
Warning:
|
|
This API is only support inputs with dtype float16 and bfloat16.
|
|
Args:
|
|
query(Tensor): The query tensor in the Attention module.
|
|
3-D tensor with shape:
|
|
[token_num, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
key(Tensor): The key tensor in the Attention module.
|
|
3-D tensor with shape:
|
|
[token_num, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
value(Tensor): The value tensor in the Attention module.
|
|
3-D tensor with shape:
|
|
[token_num, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
cu_seqlens_q(Tensor): The cumulative sequence lengths of the sequences in the batch,
|
|
used to index query.
|
|
cu_seqlens_k(Tensor): The cumulative sequence lengths of the sequences in the batch,
|
|
used to index key and value.
|
|
causal(bool): Whether enable causal mode.
|
|
softmax_scale(float): The softmax scale of the attention.
|
|
max_seqlen_q(int): Maximum sequence length of query in the batch. Note it's the padding length, not the max actual seqlen.
|
|
max_seqlen_k(int): Maximum sequence length of key/value in the batch.
|
|
|
|
Returns:
|
|
out(Tensor): The attention tensor. 3-D tensor with shape: [token_num, num_heads, head_dim]. The dtype can be float16 or bfloat16.
|
|
softmax(Tensor): The softmax tensor. None if return_softmax is False.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +SKIP('flash_attn_v3 need H100 compile')
|
|
>>> import paddle
|
|
>>> paddle.seed(2023)
|
|
>>> q = paddle.rand((10, 2, 128), dtype="bfloat16")
|
|
>>> cu_seqlens_q = paddle.to_tensor([0, 10], dtype="int32")
|
|
>>> max_seq_len_q = 10
|
|
>>> output = paddle.nn.functional.flash_attention.flash_attention_v3_varlen(
|
|
... q, q, q, cu_seqlens_q, cu_seqlens_q, max_seqlen_q=max_seq_len_q, max_seqlen_k=max_seq_len_q, causal=True
|
|
... )
|
|
>>> # doctest: -SKIP
|
|
"""
|
|
assert "xpu" not in paddle.get_device(), (
|
|
"flash_attn_varlen_func is not supported on xpu"
|
|
)
|
|
|
|
assert not paddle.get_flags(["FLAGS_cudnn_deterministic"])[
|
|
"FLAGS_cudnn_deterministic"
|
|
], "flash_attn_varlen_func does not support deterministic"
|
|
|
|
assert (
|
|
paddle.base.framework.get_flags(["FLAGS_flash_attn_version"])[
|
|
"FLAGS_flash_attn_version"
|
|
]
|
|
== 3
|
|
), "FLAGS_flash_attn_version is 2, conflicts with flash_attn_varlen_func"
|
|
|
|
assert in_dynamic_or_pir_mode(), (
|
|
"flash_attn_varlen_func only support dynamic or pir mode"
|
|
)
|
|
|
|
assert qv is None, "flash_attn_varlen_func does not support setting qv"
|
|
|
|
if softmax_scale is None:
|
|
softmax_scale = (
|
|
query.shape[-1] + (qv.shape[-1] if qv is not None else 0)
|
|
) ** (-0.5)
|
|
|
|
out, softmax_lse = _C_ops.flash_attn_v3_varlen(
|
|
query,
|
|
key,
|
|
value,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
seqused_q,
|
|
seqused_k,
|
|
qv,
|
|
q_descale,
|
|
k_descale,
|
|
v_descale,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
softmax_scale,
|
|
causal,
|
|
window_size[0],
|
|
window_size[1],
|
|
softcap,
|
|
num_splits,
|
|
pack_gqa is not None,
|
|
pack_gqa if pack_gqa is not None else False,
|
|
sm_margin,
|
|
)
|
|
return out, softmax_lse
|
|
|
|
|
|
@overload
|
|
def flash_attn_varlen_qkvpacked(
|
|
qkv: Tensor,
|
|
cu_seqlens_q: Tensor,
|
|
cu_seqlens_k: Tensor,
|
|
max_seqlen_q: int,
|
|
max_seqlen_k: int,
|
|
scale: float,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: Literal[False] = ...,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
varlen_padded: bool = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
) -> tuple[Tensor, None]: ...
|
|
|
|
|
|
@overload
|
|
def flash_attn_varlen_qkvpacked(
|
|
qkv: Tensor,
|
|
cu_seqlens_q: Tensor,
|
|
cu_seqlens_k: Tensor,
|
|
max_seqlen_q: int,
|
|
max_seqlen_k: int,
|
|
scale: float,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: Literal[True] = ...,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
varlen_padded: bool = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
) -> tuple[Tensor, Tensor]: ...
|
|
|
|
|
|
@overload
|
|
def flash_attn_varlen_qkvpacked(
|
|
qkv: Tensor,
|
|
cu_seqlens_q: Tensor,
|
|
cu_seqlens_k: Tensor,
|
|
max_seqlen_q: int,
|
|
max_seqlen_k: int,
|
|
scale: float,
|
|
dropout: float = ...,
|
|
causal: bool = ...,
|
|
return_softmax: bool = ...,
|
|
fixed_seed_offset: Tensor | None = ...,
|
|
rng_name: str = ...,
|
|
varlen_padded: bool = ...,
|
|
training: bool = ...,
|
|
name: str | None = ...,
|
|
) -> tuple[Tensor, Tensor | None]: ...
|
|
|
|
|
|
def flash_attn_varlen_qkvpacked(
|
|
qkv,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
scale,
|
|
dropout=0.0,
|
|
causal=False,
|
|
return_softmax=False,
|
|
fixed_seed_offset=None,
|
|
rng_name="",
|
|
varlen_padded=True,
|
|
training=True,
|
|
name=None,
|
|
):
|
|
r"""
|
|
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 only supports inputs with dtype float16 and bfloat16.
|
|
|
|
Args:
|
|
qkv(Tensor): The padded query/key/value packed tensor in the Attention module. The padding part won't be computed
|
|
4-D tensor with shape:
|
|
[total_seq_len, num_heads/num_heads_k + 2, num_heads_k, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
cu_seqlens_q(Tensor): The cumulative sequence lengths of the sequences in the batch,
|
|
used to index query.
|
|
cu_seqlens_k(Tensor): The cumulative sequence lengths of the sequences in the batch,
|
|
used to index key and value.
|
|
max_seqlen_q(int): Maximum sequence length of query in the batch. Note it's the padding length, not the max actual seqlen
|
|
max_seqlen_k(int): Maximum sequence length of key/value in the batch.
|
|
scale(float): The scaling of QK^T before applying softmax.
|
|
dropout(float, optional): The dropout ratio.
|
|
causal(bool, optional): Whether enable causal mode.
|
|
return_softmax(bool, optional): Whether to return softmax.
|
|
fixed_seed_offset(Tensor|None, optional): With fixed seed, offset for dropout mask.
|
|
rng_name(str, optional): The name to select Generator.
|
|
training(bool, optional): Whether it is in the training phase.
|
|
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:
|
|
- out(Tensor). The attention tensor. The tensor is padded by zeros. 3-D tensor with shape: [total_seq_len, num_heads, head_dim]. The dtype can be float16 or bfloat16.
|
|
- softmax(Tensor). The softmax tensor. None if return_softmax is False.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +SKIP('flash_attn need A100 compile')
|
|
>>> import paddle
|
|
>>> paddle.seed(2023)
|
|
>>> q = paddle.rand((2, 128, 8, 16), dtype='float16')
|
|
>>> cu = paddle.arange(0, 384, 128, dtype='int32')
|
|
>>> qq = paddle.reshape(q, [256, 8, 16])
|
|
>>> qkv = paddle.stack([qq, qq, qq], axis=2)
|
|
>>> output = paddle.nn.functional.flash_attn_varlen_qkvpacked(qkv, cu, cu, 128, 128, 0.25, 0.0, False, False)
|
|
>>> # doctest: -SKIP
|
|
|
|
"""
|
|
if in_dynamic_or_pir_mode():
|
|
(
|
|
result_attention,
|
|
result_softmax,
|
|
) = _C_ops.flash_attn_varlen_qkvpacked(
|
|
qkv,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
fixed_seed_offset,
|
|
None,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
scale,
|
|
dropout,
|
|
causal,
|
|
return_softmax,
|
|
not training,
|
|
rng_name,
|
|
varlen_padded,
|
|
)
|
|
return result_attention, result_softmax if return_softmax else None
|
|
|
|
helper = LayerHelper('flash_attn_varlen_qkvpacked', **locals())
|
|
dtype = helper.input_dtype(input_param_name='qkv')
|
|
out = helper.create_variable_for_type_inference(dtype)
|
|
softmax = helper.create_variable_for_type_inference(dtype)
|
|
softmax_lse = helper.create_variable_for_type_inference(paddle.float32)
|
|
seed_offset = helper.create_variable_for_type_inference(paddle.int64)
|
|
inputs = {
|
|
'qkv': qkv,
|
|
'cu_seqlens_q': cu_seqlens_q,
|
|
'cu_seqlens_k': cu_seqlens_k,
|
|
'fixed_seed_offset': fixed_seed_offset,
|
|
}
|
|
outputs = {
|
|
'out': out,
|
|
'softmax': softmax,
|
|
'softmax_lse': softmax_lse,
|
|
'seed_offset': seed_offset,
|
|
}
|
|
helper.append_op(
|
|
type='flash_attn_varlen_qkvpacked',
|
|
inputs=inputs,
|
|
outputs=outputs,
|
|
attrs={
|
|
'max_seqlen_q': max_seqlen_q,
|
|
'max_seqlen_k': max_seqlen_k,
|
|
'scale': scale,
|
|
'dropout': dropout,
|
|
'causal': causal,
|
|
'return_softmax': return_softmax,
|
|
'is_test': not training,
|
|
'rng_name': rng_name,
|
|
},
|
|
)
|
|
return out, softmax if return_softmax else None
|
|
|
|
|
|
# Global cache: group_id -> unique_id tensor (for NVSHMEM bootstrap).
|
|
# Each distinct group needs its own unique_id, generated once by rank 0
|
|
# and broadcast to all members via all_gather_object.
|
|
_flashmask_unique_id_cache: dict[int, paddle.Tensor] = {}
|
|
|
|
|
|
def _get_or_create_unique_id(group):
|
|
"""Get or create the NVSHMEM unique_id for a communication group.
|
|
|
|
On first call for a given group: rank 0 generates the unique_id via
|
|
``flashmask_get_unique_id()``, then broadcasts it to all group members
|
|
using ``all_gather_object`` (CPU tensor, so regular ``all_gather`` would
|
|
raise). The result is cached by ``group.id``.
|
|
|
|
On subsequent calls: returns the cached unique_id directly.
|
|
|
|
Args:
|
|
group: A ``paddle.distributed.Group`` instance.
|
|
|
|
Returns:
|
|
tuple[Tensor, bool]: ``(unique_id, is_new)`` where ``unique_id`` is
|
|
a 128-byte uint8 CPU tensor, and ``is_new`` indicates whether this
|
|
is the first time (True) or a cache hit (False).
|
|
"""
|
|
import paddle.distributed as dist
|
|
|
|
gid = group.id
|
|
if gid in _flashmask_unique_id_cache:
|
|
return _flashmask_unique_id_cache[gid], False
|
|
|
|
# Rank 0 of this group generates the unique_id
|
|
if group.rank == 0:
|
|
unique_id = flashmask_get_unique_id()
|
|
else:
|
|
unique_id = paddle.zeros([128], dtype='uint8', device='cpu')
|
|
|
|
# Broadcast via all_gather_object (tensor is on CPU, dist env is GPU)
|
|
result_list = []
|
|
dist.all_gather_object(result_list, unique_id, group=group)
|
|
unique_id = result_list[0]
|
|
|
|
_flashmask_unique_id_cache[gid] = unique_id
|
|
return unique_id, True
|
|
|
|
|
|
def flashmask_attention(
|
|
query: Tensor,
|
|
key: Tensor,
|
|
value: Tensor,
|
|
startend_row_indices: Tensor | None = None,
|
|
*,
|
|
dropout: float = 0.0,
|
|
causal: bool = False,
|
|
window_size: int | tuple | None = None,
|
|
return_softmax_lse: bool = False,
|
|
return_seed_offset: bool = False,
|
|
fixed_seed_offset: Tensor | None = None,
|
|
rng_name: str = "",
|
|
training: bool = True,
|
|
name: str | None = None,
|
|
softmax_scale: float | None = None,
|
|
block_mask: Tensor | None = None,
|
|
group: Group | None = None,
|
|
) -> Tensor | list[Tensor]:
|
|
r"""
|
|
FlashMask: Official Implementation
|
|
|
|
This module provides the official implementation of the FlashMask algorithm as described in the paper. For more details, please refer to the paper available at: https://arxiv.org/abs/2410.01359.
|
|
|
|
The core equation utilized in FlashMask is as follows:
|
|
|
|
.. math::
|
|
|
|
\text{result} = \text{softmax}\left(\frac{Q \cdot K^T}{\sqrt{d}} + M\right) \cdot V
|
|
|
|
In this equation:
|
|
|
|
- ``Q``, ``K``, and ``V`` are the input tensors to the attention module.
|
|
- All these tensors share the same dimensions.
|
|
- ``d`` denotes the size of the last dimension of these tensors.
|
|
- ``M`` represents the column-wise sparse mask introduced by FlashMask.
|
|
|
|
Args:
|
|
query (Tensor): The query tensor in the attention module.
|
|
A 4-D tensor with shape [batch_size, q_seq_len, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
key (Tensor): The key tensor in the attention module.
|
|
A 4-D tensor with shape [batch_size, k_seq_len, k_num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
value (Tensor): The value tensor in the attention module.
|
|
A 4-D tensor with shape [batch_size, k_seq_len, k_num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
startend_row_indices(Tensor):
|
|
A column-wise sparse attention mask row indices tensor.
|
|
A 4-D tensor with shape [batch_size, k_num_heads, k_seq_len, {1, 2, 4}].
|
|
The dtype must be int32. k_num_heads can be 1 or the same as key's num_heads. When num_heads is 1, it will be broadcast to match key's num_heads.
|
|
Depending on the value of the causal parameter, startend_row_indices can take different shapes and meanings.
|
|
|
|
- When `causal=True` and the shape is [batch_size, k_num_heads, k_seq_len, 1],
|
|
indicating unidirectional attention. The value represents the starting row index of the left
|
|
lower triangular mask in the dense mask. The value startend_row_indices[..., 0] indicates that elements in the lower left triangle of the attention score matrix starting from the startend_row_indices[..., 0]-th row downwards (inclusive) will be masked.
|
|
- When `causal=True` and the shape is [batch_size, k_num_heads, k_seq_len, 2],
|
|
indicating unidirectional attention. The values represent the starting and ending row indices of
|
|
the left lower triangular mask in the dense mask. The values startend_row_indices[..., 0:2] in startend_row_indices indicate that elements in the lower left triangle of the attention score matrix starting from the startend_row_indices[..., 0]-th row downwards (inclusive) but above the startend_row_indices[..., 1]-th row (exclusive) will be masked.
|
|
- When `causal=False` and the shape is [batch_size, k_num_heads, k_seq_len, 2],
|
|
indicating bidirectional attention. The values represent the starting row index of the left
|
|
lower triangular mask and the ending row index of the right upper triangular mask in the dense mask. The values startend_row_indices[..., 0:2] in startend_row_indices indicate that elements in the lower left triangle of the attention score matrix starting from the startend_row_indices[..., 0]-th row downwards (inclusive) will be masked, and elements in the upper right triangle starting from the startend_row_indices[..., 1]-th row upwards (exclusive) will be masked.
|
|
- When `causal=False` and the shape is [batch_size, k_num_heads, k_seq_len, 4] ,
|
|
indicating bidirectional attention. The values represent the start and end row indices of the
|
|
left lower triangular mask and the start and end row indices of the right upper triangular mask in the dense mask. The values startend_row_indices[..., 0:4] in startend_row_indices indicate that elements in the lower left triangle of the attention score matrix starting from the startend_row_indices[..., 0]-th row downwards (inclusive) but above the startend_row_indices[..., 1] row (exclusive) will be masked, and elements in the upper right triangle starting from the startend_row_indices[..., 2]-th row downwards (inclusive) but above the startend_row_indices[..., 3] row (exclusive) will be masked.
|
|
|
|
dropout (float): The dropout ratio. Default is 0.0.
|
|
causal (bool): Whether to enable causal mode. Default is False.
|
|
window_size (int|tuple, optional): Indicates the window size of sliding window local attention.
|
|
If causal mode is enabled, Query at position i will only attend to keys between [i - window_size, i] or [i - window_size[0], i].
|
|
If causal mode is disabled, Query at position i will only attend to keys between [i - window_size, i + window_size] or [i - window_size[0], i + window_size[1]].
|
|
return_softmax_lse (bool): Whether to return the log-sum-exp of the softmax. Default is False.
|
|
return_seed_offset (bool): Whether to return the random seed offset. Default is False.
|
|
fixed_seed_offset(Tensor, optional): With fixed seed, offset for dropout mask.
|
|
rng_name (str): The name to select Generator.
|
|
training (bool): Whether the module is in training mode. Default is True.
|
|
name (str, optional): Name of the operation. Default is None. Normally, users do not need to set this property.
|
|
For more information, refer to :ref:`api_guide_Name` .
|
|
block_mask (tensor, optional):
|
|
A 4-D integer mask tensor indicating whether each block in the attention matrix should be kept or masked. Must be used together with flashmask.
|
|
The shape should be [batch_size, num_heads, blocklen_q, blocklen_k], where:
|
|
|
|
blocklen_q = ceil(seqlen_q / 128), i.e., block_mask.shape[2] must be (seqlen_q + 127) // 128
|
|
blocklen_k = ceil(seqlen_k / 128), i.e., block_mask.shape[3] must be (seqlen_k + 127) // 128
|
|
block_mask.shape[1] (number of heads) must match the num_heads dimension of the flashmask
|
|
Both seqlen_q and seqlen_k must be less than or equal to 128 * 1024
|
|
The dtype should be int32, and each element should be either 0 or 1.
|
|
A value of 1 indicates that the corresponding block is kept (not masked), while 0 means the block is masked.
|
|
|
|
Usage Notes:
|
|
|
|
Only supported when blockdim_q = blockdim_k = 128 now.
|
|
Only supported when headdim = 128 now.
|
|
This argument must be provided together with flashmask.
|
|
The mask will be applied at the block level: each [i, j] position in block_mask controls whether the corresponding [128 x 128] block in the attention matrix is masked.
|
|
Any mismatch in expected shape or head dimension will raise an error.
|
|
group (paddle.distributed.Group, optional):
|
|
The communication group for distributed context parallelism (CP) overlap.
|
|
When provided, ``rank`` and ``nranks`` are automatically extracted from the group,
|
|
and the NVSHMEM unique_id is managed internally (generated once per group by rank 0,
|
|
broadcast to all members, and cached for subsequent calls). Users only need to pass
|
|
the CP group without worrying about low-level NVSHMEM initialization.
|
|
Default: None (no distributed overlap).
|
|
|
|
|
|
Returns
|
|
Tensor. The computed attention result with the same shape as the input `query`.
|
|
|
|
Warning:
|
|
This API only supports inputs with dtype float16 and bfloat16.
|
|
|
|
Hint:
|
|
This API supports GQA.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +SKIP('flash_attn need A100 compile')
|
|
>>> import paddle
|
|
>>> paddle.seed(2023)
|
|
>>> q = paddle.rand((1, 10, 2, 32), dtype="bfloat16") # shape: [batch_size, seq_len, num_heads, head_dim]
|
|
>>> k = paddle.rand((1, 10, 2, 32), dtype="bfloat16") # shape: [batch_size, seq_len, num_heads, head_dim]
|
|
>>> v = paddle.rand((1, 10, 2, 32), dtype="bfloat16") # shape: [batch_size, seq_len, num_heads, head_dim]
|
|
>>> startend_row_indices = paddle.to_tensor([8] * 10 + [5] * 10, dtype="int32").reshape([1, 2, 10, 1])
|
|
>>> output = paddle.nn.functional.flashmask_attention(q, k, v, startend_row_indices, causal=True)
|
|
>>> print(output)
|
|
Tensor(shape=[1, 10, 2, 32], dtype=bfloat16, place=Place(gpu:0), stop_gradient=True,
|
|
[[[[0.82421875, 0.27539062, 0.80859375, 0.98046875, 0.00251770,
|
|
0.41992188, 0.17285156, 0.11767578, 0.42773438, 0.31250000,
|
|
0.34570312, 0.70312500, 0.29296875, 0.44531250, 0.51562500,
|
|
0.96093750, 0.85546875, 0.15625000, 0.34765625, 0.98437500,
|
|
0.96484375, 0.45312500, 0.33593750, 0.56640625, 0.07714844,
|
|
0.43750000, 0.83984375, 0.66796875, 0.93750000, 0.24804688,
|
|
0.51171875, 0.55468750],
|
|
[0.54687500, 0.74609375, 0.43164062, 0.32421875, 0.10693359,
|
|
0.37304688, 0.53906250, 0.17187500, 0.57421875, 0.75000000,
|
|
0.13378906, 0.57031250, 0.19531250, 0.01403809, 0.29101562,
|
|
0.14257812, 0.07568359, 0.88671875, 0.75390625, 0.17089844,
|
|
0.87109375, 0.93359375, 0.89843750, 0.58203125, 0.75390625,
|
|
0.27539062, 0.67968750, 0.24804688, 0.57812500, 0.67578125,
|
|
0.92578125, 0.98046875]],
|
|
|
|
[[0.59765625, 0.62890625, 0.62109375, 0.75781250, 0.03295898,
|
|
0.64062500, 0.27929688, 0.20800781, 0.72265625, 0.52343750,
|
|
0.53125000, 0.61718750, 0.57421875, 0.56640625, 0.65625000,
|
|
0.48242188, 0.68359375, 0.42968750, 0.26562500, 0.86718750,
|
|
0.83203125, 0.40820312, 0.38281250, 0.59765625, 0.43945312,
|
|
0.22851562, 0.86328125, 0.51562500, 0.89453125, 0.62500000,
|
|
0.50390625, 0.67968750],
|
|
[0.34765625, 0.61328125, 0.58593750, 0.60156250, 0.43164062,
|
|
0.41601562, 0.71093750, 0.59765625, 0.53515625, 0.78125000,
|
|
0.13867188, 0.30664062, 0.48828125, 0.04394531, 0.24316406,
|
|
0.18847656, 0.10644531, 0.71093750, 0.69140625, 0.35937500,
|
|
0.44531250, 0.81640625, 0.44140625, 0.64062500, 0.81640625,
|
|
0.61328125, 0.72265625, 0.53125000, 0.49414062, 0.59765625,
|
|
0.54296875, 0.61328125]],
|
|
|
|
[[0.65234375, 0.47656250, 0.71875000, 0.64843750, 0.23828125,
|
|
0.61328125, 0.29101562, 0.26562500, 0.54296875, 0.60937500,
|
|
0.67187500, 0.67578125, 0.64062500, 0.41406250, 0.47656250,
|
|
0.40820312, 0.66406250, 0.39453125, 0.39453125, 0.62109375,
|
|
0.58593750, 0.31054688, 0.31835938, 0.45703125, 0.52343750,
|
|
0.43164062, 0.64453125, 0.49804688, 0.82812500, 0.48242188,
|
|
0.38476562, 0.59375000],
|
|
[0.44921875, 0.62109375, 0.50390625, 0.51562500, 0.51953125,
|
|
0.57812500, 0.78515625, 0.73437500, 0.60546875, 0.55078125,
|
|
0.30273438, 0.23339844, 0.60546875, 0.33007812, 0.23242188,
|
|
0.30468750, 0.34570312, 0.70703125, 0.72656250, 0.58593750,
|
|
0.40234375, 0.62109375, 0.62109375, 0.69531250, 0.66796875,
|
|
0.51562500, 0.45898438, 0.67968750, 0.48828125, 0.50000000,
|
|
0.54687500, 0.71875000]],
|
|
|
|
[[0.67578125, 0.50000000, 0.58203125, 0.62109375, 0.43554688,
|
|
0.69531250, 0.30273438, 0.24023438, 0.57812500, 0.63671875,
|
|
0.51171875, 0.52734375, 0.60546875, 0.45507812, 0.42382812,
|
|
0.46093750, 0.55859375, 0.34960938, 0.39453125, 0.57031250,
|
|
0.55078125, 0.47265625, 0.24609375, 0.51953125, 0.46093750,
|
|
0.49218750, 0.49609375, 0.60156250, 0.76953125, 0.57421875,
|
|
0.40429688, 0.57031250],
|
|
[0.45703125, 0.71093750, 0.58984375, 0.43164062, 0.54296875,
|
|
0.57031250, 0.72265625, 0.61328125, 0.64453125, 0.50781250,
|
|
0.28125000, 0.19531250, 0.60546875, 0.40625000, 0.18554688,
|
|
0.33203125, 0.40039062, 0.58593750, 0.79687500, 0.45507812,
|
|
0.32812500, 0.58203125, 0.70703125, 0.64453125, 0.53906250,
|
|
0.57421875, 0.48828125, 0.53515625, 0.49804688, 0.50000000,
|
|
0.48437500, 0.55468750]],
|
|
|
|
[[0.64453125, 0.43164062, 0.54687500, 0.53125000, 0.42187500,
|
|
0.71484375, 0.30273438, 0.21484375, 0.50390625, 0.69531250,
|
|
0.58203125, 0.51562500, 0.61328125, 0.41992188, 0.40039062,
|
|
0.46679688, 0.58984375, 0.39062500, 0.41992188, 0.49023438,
|
|
0.47851562, 0.47070312, 0.30078125, 0.50390625, 0.47656250,
|
|
0.44921875, 0.43164062, 0.63671875, 0.78125000, 0.60156250,
|
|
0.48242188, 0.58203125],
|
|
[0.52343750, 0.69921875, 0.58984375, 0.35156250, 0.49218750,
|
|
0.58593750, 0.71093750, 0.59375000, 0.66406250, 0.49414062,
|
|
0.24023438, 0.18554688, 0.66796875, 0.50000000, 0.23144531,
|
|
0.29882812, 0.49414062, 0.57031250, 0.70312500, 0.42773438,
|
|
0.35351562, 0.47460938, 0.73437500, 0.53125000, 0.47070312,
|
|
0.49609375, 0.50000000, 0.55078125, 0.50000000, 0.45898438,
|
|
0.45703125, 0.61328125]],
|
|
|
|
[[0.63671875, 0.41210938, 0.52734375, 0.56640625, 0.44531250,
|
|
0.64843750, 0.37890625, 0.31250000, 0.56640625, 0.62890625,
|
|
0.53125000, 0.51562500, 0.54296875, 0.50781250, 0.35546875,
|
|
0.41601562, 0.55468750, 0.36914062, 0.35937500, 0.45117188,
|
|
0.46875000, 0.49609375, 0.28710938, 0.50000000, 0.49609375,
|
|
0.50000000, 0.51562500, 0.57031250, 0.77734375, 0.62109375,
|
|
0.43164062, 0.50781250],
|
|
[0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. ]],
|
|
|
|
[[0.62109375, 0.44531250, 0.46875000, 0.61328125, 0.39062500,
|
|
0.60156250, 0.41015625, 0.28710938, 0.58984375, 0.67968750,
|
|
0.55859375, 0.48632812, 0.51562500, 0.42382812, 0.37695312,
|
|
0.46679688, 0.54687500, 0.44921875, 0.33789062, 0.36328125,
|
|
0.49023438, 0.44140625, 0.25000000, 0.45312500, 0.43945312,
|
|
0.45507812, 0.46679688, 0.57812500, 0.65625000, 0.64062500,
|
|
0.42382812, 0.57031250],
|
|
[0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. ]],
|
|
|
|
[[0.62500000, 0.47070312, 0.51562500, 0.61328125, 0.36718750,
|
|
0.66406250, 0.37890625, 0.28320312, 0.65625000, 0.66015625,
|
|
0.48632812, 0.53906250, 0.46679688, 0.47851562, 0.43359375,
|
|
0.45703125, 0.47070312, 0.39843750, 0.32617188, 0.37304688,
|
|
0.49023438, 0.50390625, 0.27148438, 0.46679688, 0.37695312,
|
|
0.49023438, 0.47265625, 0.58593750, 0.64453125, 0.60156250,
|
|
0.38476562, 0.62109375],
|
|
[0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. ]],
|
|
|
|
[[0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. ],
|
|
[0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. ]],
|
|
|
|
[[0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. ],
|
|
[0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. , 0. , 0. , 0. ,
|
|
0. , 0. ]]]])
|
|
>>> # doctest: -SKIP
|
|
|
|
|
|
To convert FlashMask's `startend_row_indices` to `dense_mask`, use the code below:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> import numpy as np
|
|
>>> def flashmask_to_densemask(startend_row_indices, dtype, causal=True):
|
|
... if startend_row_indices is None:
|
|
... return None
|
|
... bz, num_head, seq_len, bound_num = startend_row_indices.shape
|
|
... m = paddle.zeros((bz, num_head, seq_len, seq_len), dtype=dtype)
|
|
... has_end = (causal and bound_num == 2) or ((not causal) and bound_num == 4)
|
|
... for bi in range(bz):
|
|
... for hi in range(num_head):
|
|
... for j in range(seq_len):
|
|
... downstart = startend_row_indices[bi, hi, j, 0]
|
|
... if has_end:
|
|
... downend = startend_row_indices[bi, hi, j, 1]
|
|
... m[bi, hi, downstart:downend, j] = -np.inf
|
|
... else:
|
|
... m[bi, hi, downstart:, j] = -np.inf
|
|
... if causal:
|
|
... m[bi, hi, :j, j] = -np.inf
|
|
... else:
|
|
... if has_end:
|
|
... upstart = startend_row_indices[bi, hi, j, 2]
|
|
... upend = startend_row_indices[bi, hi, j, 3]
|
|
... m[bi, hi, upstart:upend, j] = -np.inf
|
|
... else:
|
|
... upend = startend_row_indices[bi, hi, j, 1]
|
|
... m[bi, hi, :upend, j] = -np.inf
|
|
... return m
|
|
|
|
For `Causal Mask`, where `causal=True`, the values of `startend_row_indices` are as follows:
|
|
|
|
.. code-block:: pycon
|
|
|
|
[[[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]]])
|
|
|
|
>>> # doctest: +SKIP('Only example')
|
|
>>> import paddle
|
|
>>> startend_row_indices = paddle.to_tensor([8] * 10, dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> print(startend_row_indices)
|
|
Tensor(shape=[1, 1, 10, 1], dtype=int32, place=Place(gpu:0), stop_gradient=True,
|
|
[[[[8],
|
|
[8],
|
|
[8],
|
|
[8],
|
|
[8],
|
|
[8],
|
|
[8],
|
|
[8],
|
|
[8],
|
|
[8]]]])
|
|
>>> # doctest: -SKIP
|
|
|
|
|
|
For `Sliding Window Mask`, where `causal=True`, the values of `startend_row_indices` are as follows:
|
|
|
|
.. code-block:: pycon
|
|
|
|
[[[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 1, 1, 1, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 1, 1, 1, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 1, 1, 1, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 1, 1, 1, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 1, 1, 1, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 1, 1, 1, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 1, 1, 1]]]])
|
|
|
|
>>> # doctest: +SKIP('Only example')
|
|
>>> import paddle
|
|
>>> startend_row_indices = paddle.to_tensor([3, 4, 5, 6, 7, 8, 9, 10, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> print(startend_row_indices)
|
|
Tensor(shape=[1, 1, 10, 1], dtype=int32, place=Place(gpu:0), stop_gradient=True,
|
|
[[[[3 ],
|
|
[4 ],
|
|
[5 ],
|
|
[6 ],
|
|
[7 ],
|
|
[8 ],
|
|
[9 ],
|
|
[10],
|
|
[10],
|
|
[10]]]])
|
|
>>> # doctest: -SKIP
|
|
|
|
For `Causal Document Mask`, where `causal=True`, the values of `startend_row_indices` are as follows:
|
|
|
|
.. code-block:: pycon
|
|
|
|
[[[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 1, 1, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 1, 1, 1, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 1, 1, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 1, 1, 1]]]])
|
|
|
|
>>> # doctest: +SKIP('Only example')
|
|
>>> import paddle
|
|
>>> startend_row_indices = paddle.to_tensor([4, 4, 4, 4, 7, 7, 7, 10, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> print(startend_row_indices)
|
|
Tensor(shape=[1, 1, 10, 1], dtype=int32, place=Place(gpu:0), stop_gradient=True,
|
|
[[[[4 ],
|
|
[4 ],
|
|
[4 ],
|
|
[4 ],
|
|
[7 ],
|
|
[7 ],
|
|
[7 ],
|
|
[10],
|
|
[10],
|
|
[10]]]])
|
|
>>> # doctest: -SKIP
|
|
|
|
For `Document Mask`, where `causal=False`, the values of `startend_row_indices` are as follows:
|
|
|
|
.. code-block:: pycon
|
|
|
|
[[[[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 1, 1, 1, 0, 0, 0],
|
|
[0, 0, 0, 0, 1, 1, 1, 0, 0, 0],
|
|
[0, 0, 0, 0, 1, 1, 1, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 1, 1, 1],
|
|
[0, 0, 0, 0, 0, 0, 0, 1, 1, 1],
|
|
[0, 0, 0, 0, 0, 0, 0, 1, 1, 1]]]])
|
|
|
|
>>> # doctest: +SKIP('Only example')
|
|
>>> import paddle
|
|
>>> LTS = paddle.to_tensor([4, 4, 4, 4, 7, 7, 7, 10, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> UTE = paddle.to_tensor([0, 0, 0, 0, 4, 4, 4, 7, 7, 7], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> startend_row_indices = paddle.concat([LTS, UTE], axis=-1)
|
|
>>> print(startend_row_indices)
|
|
Tensor(shape=[1, 1, 10, 2], dtype=int32, place=Place(gpu:0), stop_gradient=True,
|
|
[[[[4 , 0 ],
|
|
[4 , 0 ],
|
|
[4 , 0 ],
|
|
[4 , 0 ],
|
|
[7 , 4 ],
|
|
[7 , 4 ],
|
|
[7 , 4 ],
|
|
[10, 7 ],
|
|
[10, 7 ],
|
|
[10, 7 ]]]])
|
|
>>> # doctest: -SKIP
|
|
|
|
For `Share Question Mask`, where `causal=True`, the values of `startend_row_indices` are as follows:
|
|
|
|
.. code-block:: pycon
|
|
|
|
[[[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
|
|
[1, 1, 1, 1, 0, 0, 0, 1, 0, 0],
|
|
[1, 1, 1, 1, 0, 0, 0, 1, 1, 0],
|
|
[1, 1, 1, 1, 0, 0, 0, 1, 1, 1]]]])
|
|
|
|
>>> # doctest: +SKIP('Only example')
|
|
>>> import paddle
|
|
>>> startend_row_indices = paddle.to_tensor([10, 10, 10, 10, 7, 7, 7, 10, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> print(startend_row_indices)
|
|
Tensor(shape=[1, 1, 10, 1], dtype=int32, place=Place(gpu:0), stop_gradient=True,
|
|
[[[[10],
|
|
[10],
|
|
[10],
|
|
[10],
|
|
[7 ],
|
|
[7 ],
|
|
[7 ],
|
|
[10],
|
|
[10],
|
|
[10]]]])
|
|
>>> # doctest: -SKIP
|
|
|
|
For `Global + Sliding Window Mask`, where `causal=False`, the values of `startend_row_indices` are as follows:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +SKIP('Only example')
|
|
|
|
[[[[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
|
|
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
|
|
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
|
|
[1, 1, 0, 1, 1, 1, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 1, 1, 1, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 1, 1, 1, 0, 0],
|
|
[1, 1, 0, 0, 0, 0, 1, 1, 1, 0],
|
|
[1, 1, 0, 0, 0, 0, 0, 1, 1, 1],
|
|
[1, 1, 0, 0, 0, 0, 0, 0, 1, 1]]]])
|
|
|
|
>>> import paddle
|
|
>>> LTS = paddle.to_tensor([10, 10, 4, 5, 6, 7, 8, 9, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> LTE = paddle.to_tensor([10, 10, 10, 10, 10, 10, 10, 10, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> UTS = paddle.to_tensor([0, 0, 0, 0, 2, 2, 2, 2, 2, 2], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> UTE = paddle.to_tensor([0, 0, 0, 0, 3, 4, 5, 6, 7, 8], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> startend_row_indices = paddle.concat([LTS, LTE, UTS, UTE], axis=-1)
|
|
>>> print(startend_row_indices)
|
|
Tensor(shape=[1, 1, 10, 4], dtype=int32, place=Place(gpu:0), stop_gradient=True,
|
|
[[[[10, 10, 0 , 0 ],
|
|
[10, 10, 0 , 0 ],
|
|
[4 , 10, 0 , 0 ],
|
|
[5 , 10, 0 , 0 ],
|
|
[6 , 10, 2 , 3 ],
|
|
[7 , 10, 2 , 4 ],
|
|
[8 , 10, 2 , 5 ],
|
|
[9 , 10, 2 , 6 ],
|
|
[10, 10, 2 , 7 ],
|
|
[10, 10, 2 , 8 ]]]])
|
|
>>> # doctest: -SKIP
|
|
|
|
For `Causal Blockwise Mask`, where `causal=True`, the values of `startend_row_indices` are as follows:
|
|
|
|
.. code-block:: pycon
|
|
|
|
[[[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 1, 1, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 1, 1, 1, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]]])
|
|
|
|
>>> # doctest: +SKIP('Only example')
|
|
>>> import paddle
|
|
>>> LTS = paddle.to_tensor([4, 4, 4, 4, 10, 10, 10, 10, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> LTE = paddle.to_tensor([7, 7, 7, 7, 10, 10, 10, 10, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> startend_row_indices = paddle.concat([LTS, LTE], axis=-1)
|
|
>>> print(startend_row_indices)
|
|
Tensor(shape=[1, 1, 10, 2], dtype=int32, place=Place(gpu:0), stop_gradient=True,
|
|
[[[[4 , 7 ],
|
|
[4 , 7 ],
|
|
[4 , 7 ],
|
|
[4 , 7 ],
|
|
[10, 10],
|
|
[10, 10],
|
|
[10, 10],
|
|
[10, 10],
|
|
[10, 10],
|
|
[10, 10]]]])
|
|
>>> # doctest: -SKIP
|
|
|
|
For `Prefix LM Document Mask`, where `causal=False`, the values of `startend_row_indices` are as follows:
|
|
|
|
.. code-block:: pycon
|
|
|
|
[[[[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 1, 1, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 1, 1, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 1, 1, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 1, 1, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 1, 1, 1, 0, 0],
|
|
[0, 0, 0, 0, 0, 1, 1, 1, 1, 0],
|
|
[0, 0, 0, 0, 0, 1, 1, 1, 1, 1]]]])
|
|
|
|
>>> # doctest: +SKIP('Only example')
|
|
>>> import paddle
|
|
>>> LTS = paddle.to_tensor([3, 3, 3, 5, 5, 10, 10, 10, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> UTE = paddle.to_tensor([0, 0, 2, 3, 3, 5, 5, 7, 8, 9], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> startend_row_indices = paddle.concat([LTS, UTE], axis=-1)
|
|
>>> print(startend_row_indices)
|
|
Tensor(shape=[1, 1, 10, 2], dtype=int32, place=Place(gpu:0), stop_gradient=True,
|
|
[[[[3 , 0 ],
|
|
[3 , 0 ],
|
|
[3 , 2 ],
|
|
[5 , 3 ],
|
|
[5 , 3 ],
|
|
[10, 5 ],
|
|
[10, 5 ],
|
|
[10, 7 ],
|
|
[10, 8 ],
|
|
[10, 9 ]]]])
|
|
>>> # doctest: -SKIP
|
|
|
|
For `Prefix LM Causal Mask`, where `causal=False`, the values of `startend_row_indices` are as follows:
|
|
|
|
.. code-block:: pycon
|
|
|
|
[[[[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]]])
|
|
|
|
>>> # doctest: +SKIP('Only example')
|
|
>>> import paddle
|
|
>>> LTS = paddle.to_tensor([10, 10, 10, 10, 10, 10, 10, 10, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> UTE = paddle.to_tensor([0, 0, 0, 0, 0, 5, 6, 7, 8, 9], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> startend_row_indices = paddle.concat([LTS, UTE], axis=-1)
|
|
>>> print(startend_row_indices)
|
|
Tensor(shape=[1, 1, 10, 2], dtype=int32, place=Place(gpu:0), stop_gradient=True,
|
|
[[[[10, 0 ],
|
|
[10, 0 ],
|
|
[10, 0 ],
|
|
[10, 0 ],
|
|
[10, 0 ],
|
|
[10, 5 ],
|
|
[10, 6 ],
|
|
[10, 7 ],
|
|
[10, 8 ],
|
|
[10, 9 ]]]])
|
|
|
|
For `QK-sparse Mask`, where `causal=True`, the values of `startend_row_indices` are as follows:
|
|
|
|
.. code-block:: pycon
|
|
|
|
[[[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
|
|
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]]])
|
|
|
|
>>> # doctest: +SKIP('Only example')
|
|
>>> import paddle
|
|
>>> LTS = paddle.to_tensor([10, 10, 2, 3, 4, 5, 6, 7, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> LTE = paddle.to_tensor([10, 10, 5, 5, 5, 5, 8, 8, 10, 10], dtype="int32").reshape([1, 1, 10, 1])
|
|
>>> startend_row_indices = paddle.concat([LTS, LTE], axis=-1)
|
|
>>> print(startend_row_indices)
|
|
Tensor(shape=[1, 1, 10, 2], dtype=int32, place=Place(gpu:0), stop_gradient=True,
|
|
[[[[10, 10],
|
|
[10, 10],
|
|
[2 , 5 ],
|
|
[3 , 5 ],
|
|
[4 , 5 ],
|
|
[5 , 5 ],
|
|
[6 , 8 ],
|
|
[7 , 8 ],
|
|
[10, 10],
|
|
[10, 10]]]])
|
|
|
|
>>> # doctest: -SKIP
|
|
"""
|
|
if window_size is not None:
|
|
if isinstance(window_size, int):
|
|
window_size = (window_size, window_size)
|
|
sq = query.shape[1]
|
|
bsz = query.shape[0]
|
|
assert startend_row_indices is None, (
|
|
"can't use window_size with startend_row_indices"
|
|
)
|
|
if causal:
|
|
startend_row_indices = paddle.arange(
|
|
window_size[0] + 1, sq + window_size[0] + 1, dtype="int32"
|
|
).reshape((1, 1, sq, 1))
|
|
startend_row_indices = paddle.clip(
|
|
startend_row_indices, max=sq
|
|
).repeat_interleave(bsz, 0)
|
|
|
|
else:
|
|
startend_row_indices = paddle.empty((1, 1, sq, 2), dtype="int32")
|
|
startend_row_indices[0, 0, :, 0] = paddle.arange(
|
|
window_size[0] + 1, sq + window_size[0] + 1, dtype="int32"
|
|
)
|
|
startend_row_indices[0, 0, :, 1] = paddle.arange(
|
|
-window_size[1], sq - window_size[1], dtype="int32"
|
|
)
|
|
startend_row_indices = paddle.clip(
|
|
startend_row_indices, min=0, max=sq
|
|
).repeat_interleave(bsz, 0)
|
|
|
|
# --- Distributed group resolution ---
|
|
if group is not None:
|
|
rank = group.rank
|
|
nranks = group.nranks
|
|
else:
|
|
rank = 0
|
|
nranks = 1
|
|
|
|
if block_mask is not None:
|
|
# xhy: can set a full startend_row_indices for block_mask_attn when using block_mask_attn?
|
|
assert startend_row_indices is not None, (
|
|
"must provide startend_row_indices when using block_mask_attn"
|
|
)
|
|
|
|
if startend_row_indices is None:
|
|
(
|
|
out,
|
|
result_softmax,
|
|
result_softmax_lse,
|
|
result_seed_offset,
|
|
) = _C_ops.flash_attn(
|
|
query,
|
|
key,
|
|
value,
|
|
fixed_seed_offset,
|
|
None,
|
|
dropout,
|
|
causal,
|
|
False,
|
|
not training,
|
|
rng_name,
|
|
)
|
|
|
|
else:
|
|
assert startend_row_indices.dtype == paddle.int32, (
|
|
f"startend_row_indices.dtype must be paddle.int32, but got {startend_row_indices.dtype}"
|
|
)
|
|
assert len(startend_row_indices.shape) == 4, (
|
|
f"startend_row_indices rank must be 4,but got {startend_row_indices.shape}"
|
|
)
|
|
|
|
assert startend_row_indices.shape[0] == key.shape[0], (
|
|
f"startend_row_indices.shape[0] must be equal to batch_size, but got {startend_row_indices.shape[0]} and {key.shape[0]}"
|
|
)
|
|
|
|
# for context parallel, seqlen of mask len can be cp_size * (local_key seqlen)
|
|
assert (
|
|
startend_row_indices.shape[2] == key.shape[1]
|
|
or startend_row_indices.shape[2] == key.shape[1] * nranks
|
|
), (
|
|
f"startend_row_indices.shape[2] must be equal to seqlen_k or seqlen_k * world_size, but got {startend_row_indices.shape[2]} and {key.shape[2]}. World size: {nranks}"
|
|
)
|
|
assert startend_row_indices.shape[1] in [
|
|
1,
|
|
query.shape[2],
|
|
key.shape[2],
|
|
], (
|
|
"startend_row_indices head_num must be equal to 1(broadcast) or head_num_q or head_num_k."
|
|
)
|
|
|
|
if block_mask is not None:
|
|
assert block_mask.dtype == paddle.int32, (
|
|
f"block_mask.dtype must be paddle.int32, but got {block_mask.dtype}"
|
|
)
|
|
|
|
assert block_mask.shape[0] == key.shape[0], (
|
|
f"block_mask.shape[0] must be equal to batch_size, but got {block_mask.shape[0]} and {key.shape[0]}"
|
|
)
|
|
|
|
assert block_mask.shape[1] == startend_row_indices.shape[1], (
|
|
f"block_mask.shape[1] must be equal to startend_row_indices.shape[1], but got {block_mask.shape[1]} and {key.shape[2]}"
|
|
)
|
|
|
|
assert block_mask.shape[2] == (query.shape[1] + 127) // 128, (
|
|
"block_size must be 128 when using block_mask_attn"
|
|
)
|
|
|
|
assert block_mask.shape[3] == (key.shape[1] + 127) // 128, (
|
|
"block_size must be 128 when using block_mask_attn"
|
|
)
|
|
|
|
assert key.shape[3] == 128, (
|
|
"headdim must be 128 when using block_mask_attn"
|
|
)
|
|
|
|
if causal:
|
|
if startend_row_indices.shape[-1] == 1:
|
|
has_end = False
|
|
elif startend_row_indices.shape[-1] == 2:
|
|
has_end = True
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid shape of startend_row_indices, when causal is True, the last dimension should be either 1 or 2 but got {startend_row_indices.shape[-1]}"
|
|
)
|
|
else:
|
|
if startend_row_indices.shape[-1] == 2:
|
|
has_end = False
|
|
elif startend_row_indices.shape[-1] == 4:
|
|
has_end = True
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid shape of startend_row_indices, when causal is False, the last dimension should be either 2 or 4 but got {startend_row_indices.shape[-1]}"
|
|
)
|
|
|
|
if (
|
|
"xpu" not in paddle.get_device()
|
|
and paddle.get_flags(["FLAGS_cudnn_deterministic"])[
|
|
"FLAGS_cudnn_deterministic"
|
|
]
|
|
):
|
|
assert block_mask is None, (
|
|
" blockmask attention no supports deterministic now ."
|
|
)
|
|
|
|
if "xpu" in paddle.get_device():
|
|
fa_version = 2
|
|
elif (
|
|
paddle.base.framework.get_flags(["FLAGS_flash_attn_version"])[
|
|
"FLAGS_flash_attn_version"
|
|
]
|
|
== 3
|
|
and paddle.base.framework.get_flags(["FLAGS_cudnn_deterministic"])[
|
|
"FLAGS_cudnn_deterministic"
|
|
]
|
|
and query.shape[3] > 128
|
|
):
|
|
fa_version = 2
|
|
else:
|
|
fa_version = paddle.base.framework.get_flags(
|
|
["FLAGS_flash_attn_version"]
|
|
)["FLAGS_flash_attn_version"]
|
|
|
|
if fa_version == 2:
|
|
assert softmax_scale is None, (
|
|
"flashmask_attention does not support setting softmax_scale, use flashmask_attention_v2 instead"
|
|
)
|
|
|
|
assert block_mask is None, (
|
|
" blockmask attention only supports sm >= 90 now."
|
|
)
|
|
|
|
(
|
|
out,
|
|
result_softmax,
|
|
result_softmax_lse,
|
|
result_seed_offset,
|
|
) = _C_ops.flashmask_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
startend_row_indices,
|
|
fixed_seed_offset,
|
|
dropout,
|
|
causal,
|
|
False,
|
|
not training,
|
|
rng_name,
|
|
)
|
|
|
|
elif fa_version == 3:
|
|
assert dropout == 0.0, (
|
|
"flashmask_attention_v2 does not support dropout"
|
|
)
|
|
assert not return_seed_offset, (
|
|
"flashmask_attention_v2 does not support return seed_offset"
|
|
)
|
|
assert fixed_seed_offset is None, (
|
|
"flashmask_attention_v2 does not support setting seed_offset"
|
|
)
|
|
assert rng_name == "", (
|
|
"flashmask_attention_v2 does not support setting rng_name"
|
|
)
|
|
assert training, (
|
|
"flashmask_attention_v2 does not support setting training to False"
|
|
)
|
|
assert name is None, (
|
|
"flashmask_attention_v2 does not support setting name"
|
|
)
|
|
|
|
# Obtain unique_id from group (first call per group triggers
|
|
# NVSHMEM bootstrap; subsequent calls reuse cached state).
|
|
unique_id = None
|
|
if group is not None and nranks > 1:
|
|
uid, is_new = _get_or_create_unique_id(group)
|
|
unique_id = uid if is_new else None
|
|
|
|
if softmax_scale is None:
|
|
softmax_scale = query.shape[-1] ** (-0.5)
|
|
|
|
(
|
|
out,
|
|
result_softmax_lse,
|
|
) = _C_ops.flashmask_attention_v2(
|
|
query,
|
|
key,
|
|
value,
|
|
startend_row_indices,
|
|
block_mask,
|
|
unique_id,
|
|
softmax_scale,
|
|
causal,
|
|
rank,
|
|
nranks,
|
|
)
|
|
else:
|
|
raise ValueError(f"Invalid flash attention version: {fa_version}")
|
|
|
|
outputs = [out]
|
|
if return_softmax_lse:
|
|
outputs += [result_softmax_lse]
|
|
if return_seed_offset:
|
|
outputs += [result_seed_offset]
|
|
if len(outputs) == 1:
|
|
return outputs[0]
|
|
else:
|
|
return outputs
|
|
|
|
|
|
def flashmask_get_unique_id() -> Tensor:
|
|
"""FlashMask distributed overlap: get the unique ID to initialize NVSHMEM.
|
|
|
|
Normally, this function only needs to be called once. After initializing NVSHMEM,
|
|
there is no need to pass the unique_id tensor again. Please refer to the doc of ``flashmask_attention``
|
|
and check the usage of ``unique_id`` for more detailed usage.
|
|
|
|
Return:
|
|
Tensor. CPU Tensor with exactly 128 uint8s (128B). If flashmask module is not compiled
|
|
with ``WITH_DISTRIBUTED_OVERLAP`` flag, this function returns a zero tensor.
|
|
"""
|
|
output = paddle.zeros([128], dtype=paddle.uint8, device='cpu')
|
|
paddle._C_ops.flashmask_get_unique_id_(output)
|
|
return output
|
|
|
|
|
|
def calc_reduced_attention_scores(
|
|
query: paddle.Tensor, key: paddle.Tensor, softmax_lse: paddle.Tensor
|
|
) -> paddle.Tensor:
|
|
r"""
|
|
The equation is:
|
|
|
|
.. math::
|
|
|
|
result=reduce\_sum(softmax(\frac{ Q * K^T }{\sqrt{d}}), dim=-2)
|
|
|
|
Warning:
|
|
This API only supports inputs with dtype float16 and bfloat16.
|
|
|
|
Args:
|
|
query(Tensor): The query tensor in the Attention module.
|
|
4-D tensor with shape:
|
|
[batch_size, seqlen_q, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
key(Tensor): The key tensor in the Attention module.
|
|
4-D tensor with shape:
|
|
[batch_size, seqlen_k, num_heads, head_dim].
|
|
The dtype can be float16 or bfloat16.
|
|
softmax_lse(Tensor): The logsumexp of each row returned by _C_ops.flash_attn().
|
|
3-D tensor with shape:
|
|
[batch_size, num_heads, seqlen_q_rounded], where seqlen_q_rounded = ceil(seqlen_q/128).
|
|
The dtype is float32.
|
|
Returns:
|
|
reduced_attention_scores(Tensor), The reduce sum of attention scores across seqlen_q.
|
|
4-D tensor with shape: [batch_size, num_heads, 1, seqlen_k]. The dtype is float32.
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +SKIP('reduce_attn_scores need A100 compile')
|
|
>>> import paddle
|
|
>>> import numpy as np
|
|
>>> import paddle._C_ops as _C_ops
|
|
>>> from paddle.nn.functional.flash_attention import (
|
|
>>> calc_reduced_attention_scores
|
|
>>> )
|
|
>>> np.random.seed(2024)
|
|
>>> q_shape = (5, 1024, 16, 128)
|
|
>>> k_shape = (5, 2048, 16, 128)
|
|
>>> dtype = 'float16'
|
|
>>> query = np.random.random(q_shape)
|
|
>>> key = np.random.random(k_shape)
|
|
>>> q = paddle.to_tensor(
|
|
>>> query, place=place, dtype=dtype, stop_gradient=True
|
|
>>> )
|
|
>>> k = paddle.to_tensor(
|
|
>>> key, place=place, dtype=dtype, stop_gradient=True
|
|
>>> )
|
|
>>> _, _, softmax_lse, _ = _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
|
|
>>> "" #rng_name
|
|
>>> )
|
|
>>> reduced_attn_scores = calc_reduced_attention_scores(
|
|
>>> q,
|
|
>>> k,
|
|
>>> softmax_lse,
|
|
>>> )
|
|
>>> # doctest: -SKIP
|
|
"""
|
|
assert query.stop_gradient and key.stop_gradient, (
|
|
'calc_reduced_attention_scores() is for inference only.'
|
|
)
|
|
|
|
if in_dynamic_or_pir_mode():
|
|
reduced_scores = _C_ops.calc_reduced_attn_scores(
|
|
query, key, softmax_lse
|
|
)
|
|
return reduced_scores
|
|
|
|
helper = LayerHelper('calc_reduced_attn_scores', **locals())
|
|
reduced_scores = helper.create_variable_for_type_inference(paddle.float32)
|
|
softmax = helper.create_variable_for_type_inference(paddle.float32)
|
|
inputs = {
|
|
'q': query,
|
|
'k': key,
|
|
'softmax_lse': softmax_lse,
|
|
}
|
|
outputs = {
|
|
'reduced_scores': reduced_scores,
|
|
}
|
|
helper.append_op(
|
|
type='calc_reduced_attn_scores',
|
|
inputs=inputs,
|
|
outputs=outputs,
|
|
)
|
|
return reduced_scores
|