From b09f1557802e904176b9d897e1edfa733ca1da16 Mon Sep 17 00:00:00 2001 From: kircle <945428667@qq.com> Date: Tue, 9 Apr 2024 14:36:40 +0800 Subject: [PATCH 1/2] Add API flash_attn_qkvpacked and flash_attn_varlen_qkvpacked --- .../nn/functional/flash_attn_qkvpacked_cn.rst | 34 ++++++++++++++++ .../flash_attn_varlen_qkvpacked_cn.rst | 40 +++++++++++++++++++ 2 files changed, 74 insertions(+) create mode 100644 docs/api/paddle/nn/functional/flash_attn_qkvpacked_cn.rst create mode 100644 docs/api/paddle/nn/functional/flash_attn_varlen_qkvpacked_cn.rst diff --git a/docs/api/paddle/nn/functional/flash_attn_qkvpacked_cn.rst b/docs/api/paddle/nn/functional/flash_attn_qkvpacked_cn.rst new file mode 100644 index 00000000000..df047e53a48 --- /dev/null +++ b/docs/api/paddle/nn/functional/flash_attn_qkvpacked_cn.rst @@ -0,0 +1,34 @@ +.. _cn_api_paddle_nn_functional_flash_attn_qkvpacked: + +flash_attn_qkvpacked +------------------------------- + +.. py:function:: paddle.nn.functional.flash_attn_qkvpacked(qkv, dropout=0.0, causal=False, return_softmax=False, return_softmax_lse=False, return_seed_offset=False, training=True, name=None) + +QKVPacked的 flash_attention。 + +.. math:: + + result = softmax(\frac{ Q * K^T }{\sqrt{d}} + mask) * V + +参数 +:::::::::::: + + - **qkv** (Tensor) - 输入 Query+Key+Value Tensor，shape =[batch_size, seq_len, num_group + 2, num_heads_k, head_dim]，数据类型为 float16 或 bfloat16，其中num_group=num_heads_q/num_heads_k，并且qkv[:,:,:num_group,:,:]是Query Tensor，qkv[:,:,num_group,:,:]是Key Tensor，qkv[:,:,num_group+1,:,:]是Value Tensor + - **dropout** (bool，可选) – dropout 概率值，默认值为 0。 + - **causal** (bool，可选) - 是否使用 causal 模式，默认值：False。 + - **return_softmax** (bool，可选) - 是否返回 softmax 的结果。默认值 False。 + - **return_softmax_lse** (bool，可选) - 是否返回 return_softmax_lse 的结果。默认值为 False。 + - **return_seed_offset** (bool，可选) - 是否返回 return_seed_offset 的结果。默认值为 False。 + - **training** (bool，可选) - 指示是否为训练模式。默认值为 True。 + - **name** (str，可选) - 具体用法请参见 :ref:`api_guide_Name`，一般无需设置，默认值为 None。 + + +返回 +:::::::::::: +`Tensor`，attention 的结果。 + + +代码示例 +:::::::::::: +COPY-FROM: paddle.nn.functional.flash_attn_qkvpacked diff --git a/docs/api/paddle/nn/functional/flash_attn_varlen_qkvpacked_cn.rst b/docs/api/paddle/nn/functional/flash_attn_varlen_qkvpacked_cn.rst new file mode 100644 index 00000000000..e179a02dbb3 --- /dev/null +++ b/docs/api/paddle/nn/functional/flash_attn_varlen_qkvpacked_cn.rst @@ -0,0 +1,40 @@ +.. _cn_api_paddle_nn_functional_flash_attn_varlen_qkvpacked: + +flash_attn_varlen_qkvpacked +------------------------------- + +.. py:function:: paddle.nn.functional.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, return_softmax_lse=False, return_seed_offset=False, varlen_padded=True, training=True, name=None) + +支持变长batch的QKVPacked的 flash_attention。 + +.. math:: + + result = softmax(\frac{ Q * K^T }{\sqrt{d}} + mask) * V + +参数 +:::::::::::: + + - **qkv** (Tensor) - 输入 Query+Key+Value Tensor，shape =[total_seqlen, num_group + 2, num_heads_k, head_dim]，数据类型为 float16 或 bfloat16，其中num_group=num_heads_q/num_heads_k，并且qkv[:,:num_group,:,:]是Query Tensor，qkv[:,num_group,:,:]是Key Tensor，qkv[:,num_group+1,:,:]是Value Tensor + - **cu_seqlens_q** (Tensor) - batch中Query序列的累积序列长度 + - **cu_seqlens_k** (Tensor) - batch中Key/Value序列的累积序列长度 + - **max_seqlen_q** (int) - batch中单个Query序列最大长度 + - **max_seqlen_k** (int) - batch中单个Key/Value序列最大长度 + - **scale** (float) - QK^T在执行Softmax前的缩放因子 + - **dropout** (bool，可选) – dropout 概率值，默认值为 0。 + - **causal** (bool，可选) - 是否使用 causal 模式，默认值：False。 + - **return_softmax** (bool，可选) - 是否返回 softmax 的结果。默认值 False。 + - **return_softmax_lse** (bool，可选) - 是否返回 return_softmax_lse 的结果。默认值为 False。 + - **return_seed_offset** (bool，可选) - 是否返回 return_seed_offset 的结果。默认值为 False。 + - **varlen_padded** (bool，可选) - 输入输出是否采用Padded模式，当设置为True时输入输出应按照max_seqlen序列长度Padding，否则输入输出为Unpad格式，默认值为 True + - **training** (bool，可选) - 指示是否为训练模式。默认值为 True。 + - **name** (str，可选) - 具体用法请参见 :ref:`api_guide_Name`，一般无需设置，默认值为 None。 + + +返回 +:::::::::::: +`Tensor`，attention 的结果。 + + +代码示例 +:::::::::::: +COPY-FROM: paddle.nn.functional.flash_attn_varlen_qkvpacked From f253d29963d1b210107a923e61d6b6284809d4a5 Mon Sep 17 00:00:00 2001 From: kircle <945428667@qq.com> Date: Thu, 25 Apr 2024 12:12:30 +0800 Subject: [PATCH 2/2] update --- .../nn/functional/flash_attn_qkvpacked_cn.rst | 10 ++++----- .../flash_attn_varlen_qkvpacked_cn.rst | 22 +++++++++---------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/docs/api/paddle/nn/functional/flash_attn_qkvpacked_cn.rst b/docs/api/paddle/nn/functional/flash_attn_qkvpacked_cn.rst index df047e53a48..39feecaacd7 100644 --- a/docs/api/paddle/nn/functional/flash_attn_qkvpacked_cn.rst +++ b/docs/api/paddle/nn/functional/flash_attn_qkvpacked_cn.rst @@ -3,9 +3,9 @@ flash_attn_qkvpacked ------------------------------- -.. py:function:: paddle.nn.functional.flash_attn_qkvpacked(qkv, dropout=0.0, causal=False, return_softmax=False, return_softmax_lse=False, return_seed_offset=False, training=True, name=None) +.. py:function:: paddle.nn.functional.flash_attn_qkvpacked(qkv, dropout=0.0, causal=False, return_softmax=False, fixed_seed_offset=None, rng_name="", training=True, name=None) -QKVPacked的 flash_attention。 +QKVPacked 的 flash_attention。 .. math:: @@ -14,12 +14,12 @@ QKVPacked的 flash_attention。 参数 :::::::::::: - - **qkv** (Tensor) - 输入 Query+Key+Value Tensor，shape =[batch_size, seq_len, num_group + 2, num_heads_k, head_dim]，数据类型为 float16 或 bfloat16，其中num_group=num_heads_q/num_heads_k，并且qkv[:,:,:num_group,:,:]是Query Tensor，qkv[:,:,num_group,:,:]是Key Tensor，qkv[:,:,num_group+1,:,:]是Value Tensor + - **qkv** (Tensor) - 输入 Query+Key+Value Tensor，shape =[batch_size, seq_len, num_group + 2, num_heads_k, head_dim]，数据类型为 float16 或 bfloat16，其中 num_group=num_heads_q/num_heads_k，并且 qkv[:,:,:num_group,:,:]是 Query Tensor，qkv[:,:,num_group,:,:]是 Key Tensor，qkv[:,:,num_group+1,:,:]是 Value Tensor - **dropout** (bool，可选) – dropout 概率值，默认值为 0。 - **causal** (bool，可选) - 是否使用 causal 模式，默认值：False。 - **return_softmax** (bool，可选) - 是否返回 softmax 的结果。默认值 False。 - - **return_softmax_lse** (bool，可选) - 是否返回 return_softmax_lse 的结果。默认值为 False。 - - **return_seed_offset** (bool，可选) - 是否返回 return_seed_offset 的结果。默认值为 False。 + - **fixed_seed_offset** (Tensor，可选) - Dropout mask 的偏移量，默认值：None。 + - **rng_name** (str，可选) - 随机数生成器的名称。默认值：""。 - **training** (bool，可选) - 指示是否为训练模式。默认值为 True。 - **name** (str，可选) - 具体用法请参见 :ref:`api_guide_Name`，一般无需设置，默认值为 None。 diff --git a/docs/api/paddle/nn/functional/flash_attn_varlen_qkvpacked_cn.rst b/docs/api/paddle/nn/functional/flash_attn_varlen_qkvpacked_cn.rst index e179a02dbb3..27e93b19eef 100644 --- a/docs/api/paddle/nn/functional/flash_attn_varlen_qkvpacked_cn.rst +++ b/docs/api/paddle/nn/functional/flash_attn_varlen_qkvpacked_cn.rst @@ -3,9 +3,9 @@ flash_attn_varlen_qkvpacked ------------------------------- -.. py:function:: paddle.nn.functional.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, return_softmax_lse=False, return_seed_offset=False, varlen_padded=True, training=True, name=None) +.. py:function:: paddle.nn.functional.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) -支持变长batch的QKVPacked的 flash_attention。 +支持变长 batch 的 QKVPacked 的 flash_attention。 .. math:: @@ -14,18 +14,18 @@ flash_attn_varlen_qkvpacked 参数 :::::::::::: - - **qkv** (Tensor) - 输入 Query+Key+Value Tensor，shape =[total_seqlen, num_group + 2, num_heads_k, head_dim]，数据类型为 float16 或 bfloat16，其中num_group=num_heads_q/num_heads_k，并且qkv[:,:num_group,:,:]是Query Tensor，qkv[:,num_group,:,:]是Key Tensor，qkv[:,num_group+1,:,:]是Value Tensor - - **cu_seqlens_q** (Tensor) - batch中Query序列的累积序列长度 - - **cu_seqlens_k** (Tensor) - batch中Key/Value序列的累积序列长度 - - **max_seqlen_q** (int) - batch中单个Query序列最大长度 - - **max_seqlen_k** (int) - batch中单个Key/Value序列最大长度 - - **scale** (float) - QK^T在执行Softmax前的缩放因子 + - **qkv** (Tensor) - 输入 Query+Key+Value Tensor，shape =[total_seqlen, num_group + 2, num_heads_k, head_dim]，数据类型为 float16 或 bfloat16，其中 num_group=num_heads_q/num_heads_k，并且 qkv[:,:num_group,:,:]是 Query Tensor，qkv[:,num_group,:,:]是 Key Tensor，qkv[:,num_group+1,:,:]是 Value Tensor + - **cu_seqlens_q** (Tensor) - batch 中 Query 序列的累积序列长度 + - **cu_seqlens_k** (Tensor) - batch 中 Key/Value 序列的累积序列长度 + - **max_seqlen_q** (int) - batch 中单个 Query 序列最大长度 + - **max_seqlen_k** (int) - batch 中单个 Key/Value 序列最大长度 + - **scale** (float) - QK^T 在执行 Softmax 前的缩放因子 - **dropout** (bool，可选) – dropout 概率值，默认值为 0。 - **causal** (bool，可选) - 是否使用 causal 模式，默认值：False。 - **return_softmax** (bool，可选) - 是否返回 softmax 的结果。默认值 False。 - - **return_softmax_lse** (bool，可选) - 是否返回 return_softmax_lse 的结果。默认值为 False。 - - **return_seed_offset** (bool，可选) - 是否返回 return_seed_offset 的结果。默认值为 False。 - - **varlen_padded** (bool，可选) - 输入输出是否采用Padded模式，当设置为True时输入输出应按照max_seqlen序列长度Padding，否则输入输出为Unpad格式，默认值为 True + - **fixed_seed_offset** (Tensor，可选) - Dropout mask 的偏移量，默认值：None。 + - **rng_name** (str，可选) - 随机数生成器的名称。默认值：""。 + - **varlen_padded** (bool，可选) - 输入输出是否采用 Padded 模式，当设置为 True 时输入输出应按照 max_seqlen 序列长度 Padding，否则输入输出为 Unpad 格式，默认值为 True - **training** (bool，可选) - 指示是否为训练模式。默认值为 True。 - **name** (str，可选) - 具体用法请参见 :ref:`api_guide_Name`，一般无需设置，默认值为 None。