fused_attn.h

融合注意力的枚举和函数。

枚举

enum NVTE_QKV_Layout

QKV 张量的内存布局。 S, B, H, D, 和 T 分别代表序列长度、批次大小、头数、头大小以及批次中的令牌总数，即 t = sum(s_i) for i = 0...b-1。SBHD 和 BSHD-based 布局用于批次中序列长度相等或填充到相同长度的情况，而 THD-based 布局用于批次中序列长度不同的情况。

值

enumerator NVTE_SB3HD: SB3HD 布局

enumerator NVTE_SBH3D: SBH3D 布局

enumerator NVTE_SBHD_SB2HD: SBHD_SB2HD 布局

enumerator NVTE_SBHD_SBH2D: SBHD_SBH2D 布局

enumerator NVTE_SBHD_SBHD_SBHD: SBHD_SBHD_SBHD 布局

enumerator NVTE_BS3HD: BS3HD 布局

enumerator NVTE_BSH3D: BSH3D 布局

enumerator NVTE_BSHD_BS2HD: BSHD_BS2HD 布局

enumerator NVTE_BSHD_BSH2D: BSHD_BSH2D 布局

enumerator NVTE_BSHD_BSHD_BSHD: BSHD_BSHD_BSHD 布局

enumerator NVTE_T3HD: T3HD 布局

enumerator NVTE_TH3D: TH3D 布局

enumerator NVTE_THD_T2HD: THD_T2HD 布局

enumerator NVTE_THD_TH2D: THD_TH2D 布局

enumerator NVTE_THD_THD_THD: THD_THD_THD 布局

enum NVTE_QKV_Layout_Group

QKV 布局组。

值

enumerator NVTE_3HD: 3HD QKV 布局，即 BS3HD、SB3HD、T3HD

enumerator NVTE_H3D: H3D QKV 布局，即 BSH3D、SBH3D、TH3D

enumerator NVTE_HD_2HD: HD_2HD QKV 布局，即 BSHD_BS2HD、SBHD_SB2HD、THD_T2HD

enumerator NVTE_HD_H2D: HD_H2D QKV 布局，即 BSHD_BSH2D、SBHD_SBH2D、THD_TH2D

enumerator NVTE_HD_HD_HD: HD_HD_HD QKV 布局，即 BSHD_BSHD_BSHD、SBHD_SBHD_SBHD、THD_THD_THD

enum NVTE_QKV_Format

QKV 格式。

值

enumerator NVTE_SBHD: SBHD QKV 格式，即 SB3HD、SBH3D、SBHD_SB2HD、SBHD_SBH2D、SBHD_SBHD_SBHD

enumerator NVTE_BSHD: BSHD QKV 格式，即 BS3HD、BSH3D、BSHD_BS2HD、BSHD_BSH2D、BSHD_BSHD_BSHD

enumerator NVTE_THD: THD QKV 格式，即 T3HD、TH3D、THD_T2HD、THD_TH2D、THD_THD_THD

enum NVTE_Bias_Type

偏置类型。

值

enumerator NVTE_NO_BIAS: 无偏置

enumerator NVTE_PRE_SCALE_BIAS: 缩放前偏置

enumerator NVTE_POST_SCALE_BIAS: 缩放后偏置

enumerator NVTE_ALIBI: ALiBi

enum NVTE_Mask_Type

注意力掩码类型。

值

enumerator NVTE_NO_MASK: 无掩码

enumerator NVTE_PADDING_MASK: 填充注意力掩码

enumerator NVTE_CAUSAL_MASK: 因果注意力掩码 (与左上角对齐)

enumerator NVTE_PADDING_CAUSAL_MASK: 填充和因果注意力掩码 (与左上角对齐)

enumerator NVTE_CAUSAL_BOTTOM_RIGHT_MASK: 因果注意力掩码 (与右下角对齐)

enumerator NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK: 填充和因果注意力掩码 (与右下角对齐)

enum NVTE_Fused_Attn_Backend

融合注意力后端。

值

enumerator NVTE_No_Backend: 无支持的后端

enumerator NVTE_F16_max512_seqlen: 基于 cuDNN 的 FP16/BF16 融合注意力，适用于 <= 512 序列长度

enumerator NVTE_F16_arbitrary_seqlen: 基于 cuDNN 的 FP16/BF16 融合注意力，适用于任意序列长度

enumerator NVTE_FP8: 基于 cuDNN 的 FP8 融合注意力，适用于 <= 512 序列长度

函数

NVTE_QKV_Layout_Group nvte_get_qkv_layout_group(NVTE_QKV_Layout qkv_layout)

获取给定 QKV 布局的 QKV 布局组。

参数:: qkv_layout – [in] QKV 布局，例如 sbh3d。
返回值:: qkv 布局组，例如 h3d。

NVTE_QKV_Format nvte_get_qkv_format(NVTE_QKV_Layout qkv_layout)

获取给定 QKV 布局的 QKV 格式。

参数:: qkv_layout – [in] QKV 布局，例如 sbh3d。
返回值:: qkv 格式，例如 sbhd。

NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend(NVTEDType q_dtype, NVTEDType kv_dtype, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, float dropout, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, int64_t window_size_left, int64_t window_size_right)

根据输入参数获取融合注意力后端。

参数:

q_dtype – [in] 张量 Q 的数据类型。
kv_dtype – [in] 张量 K、V 的数据类型。
qkv_layout – [in] 张量 Q、K、V 的布局。
bias_type – [in] 注意力偏置类型。
attn_mask_type – [in] 注意力掩码类型。
dropout – [in] dropout 概率。
num_attn_heads – [in] Q 中的头数。
num_gqa_groups – [in] K、V 中的头数。
max_seqlen_q – [in] Q 的序列长度。
max_seqlen_kv – [in] K、V 的序列长度。
head_dim_qk – [in] Q、K 的头维度。
head_dim_v – [in] V 的头维度。
window_size_left – [in] 滑动窗口大小 (左半部分)。
window_size_right – [in] 滑动窗口大小 (右半部分)。

void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, size_t max_seqlen, bool is_training, float attn_scale, float dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, int64_t window_size_left, int64_t window_size_right, NVTETensor workspace, cudaStream_t stream)

计算带有打包 QKV 输入的点积注意力。

计算

P = Q * Transpose(K) + Bias
S = ScaleMaskSoftmax(P)
D = Dropout(S)
O = D * Transpose(V)

支持矩阵

| backend | precision |        qkv layout       |           bias           |                 mask                  | dropout |  sequence length  | head_dim         |
|   0     | FP16/BF16 |       BS3HD,SB3HD       |   NO/POST_SCALE_BIAS     | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   | <= 512, % 64 == 0 |    64            |
|   1     | FP16/BF16 | BS3HD,SB3HD,BSH3D,SBH3D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   |  > 512, % 64 == 0 | <= 128, % 8 == 0 |
|   2     |   FP8     |          T3HD           |          NO_BIAS         |               PADDING_MASK            |   Yes   | <= 512, % 64 == 0 |    64            |

注释

Tensor cu_seqlens_padded 帮助识别张量 Q、K、V 和 O 中不同序列的正确偏移量。当 QKV 格式 (nvte_get_qkv_format(qkv_layout)) 为 bshd 或 sbhd 时，偏移张量不用于 attention 计算，可以设置为空 NVTETensor。当 QKV 格式为 thd 时，此张量应遵循以下规则。当序列之间没有填充时，偏移张量应等于 cu_seqlens。当序列之间存在填充时，用户有责任根据需要调整偏移量。例如，一个包含 4 个序列的张量 [a, PAD, b, b, c, PAD, PAD, d, d] 应该有 cu_seqlens = [0, 1, 3, 4, 6] 和 cu_seqlens_padded= [0, 2, 4, 7, 9]。

参数:

QKV – [in] 打包格式的 QKV 张量，H3D 或 3HD。
Bias – [in] Bias 张量。
S – [inout] S 张量。
O – [out] 输出 O 张量。
Aux_CTX_Tensors – [out] 训练时的辅助输出张量，例如 M、ZInv、rng_state。
cu_seqlens – [in] 累积序列长度，[batch_size + 1]。
cu_seqlens_padded – [in] QKV 的累积序列偏移量，[batch_size + 1]。
rng_state – [in] CUDA 随机数生成器的种子和偏移量。
max_seqlen – [in] 用于计算的最大序列长度，它可能 >= i=0,…batch_size-1 的 max(seqlen_i)。
is_training – [in] 是否处于训练模式或推理模式。
attn_scale – [in] Q * K.T 的缩放因子。
dropout – [in] Dropout 概率。
qkv_layout – [in] QKV 张量的布局。
bias_type – [in] Bias 类型。
attn_mask_type – [in] Attention mask 类型。
window_size_left – [in] 滑动窗口大小 (左半部分)。
window_size_right – [in] 滑动窗口大小 (右半部分)。
workspace – [in] Workspace 张量。
stream – [in] 用于此操作的 CUDA stream。

void nvte_fused_attn_bwd_qkvpacked(const NVTETensor QKV, const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQKV, NVTETensor dBias, const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, size_t max_seqlen, float attn_scale, float dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, int64_t window_size_left, int64_t window_size_right, bool deterministic, NVTETensor workspace, cudaStream_t stream)

Compute the backward of the dot product attention with packed QKV input.

支持矩阵

| backend | precision |        qkv layout       |           bias           |                 mask                  | dropout |  sequence length  | head_dim         |
|   0     | FP16/BF16 |       BS3HD,SB3HD       |   NO/POST_SCALE_BIAS     | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   | <= 512, % 64 == 0 |    64            |
|   1     | FP16/BF16 | BS3HD,SB3HD,BSH3D,SBH3D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   |  > 512, % 64 == 0 | <= 128, % 8 == 0 |
|   2     |   FP8     |          T3HD           |          NO_BIAS         |               PADDING_MASK            |   Yes   | <= 512, % 64 == 0 |    64            |

注释

Tensor cu_seqlens_padded 帮助识别张量 Q、K、V 和 O 中不同序列的正确偏移量。当 QKV 格式 (nvte_get_qkv_format(qkv_layout)) 为 bshd 或 sbhd 时，偏移张量不用于 attention 计算，可以设置为空 NVTETensor。当 QKV 格式为 thd 时，此张量应遵循以下规则。当序列之间没有填充时，偏移张量应等于 cu_seqlens。当序列之间存在填充时，用户有责任根据需要调整偏移量。例如，一个包含 4 个序列的张量 [a, PAD, b, b, c, PAD, PAD, d, d] 应该有 cu_seqlens = [0, 1, 3, 4, 6] 和 cu_seqlens_padded= [0, 2, 4, 7, 9]。

参数:

QKV – [in] 打包格式的 QKV 张量，H3D 或 3HD。
O – [in] 前向传播的 O 张量。
dO – [in] O 张量的梯度。
S – [in] S 张量。
dP – [inout] P 张量的梯度。
Aux_CTX_Tensors – [in] 训练模式下上下文的辅助张量，例如 M、ZInv、rng_state。
dQKV – [out] QKV 张量的梯度。
dBias – [out] Bias 张量的梯度。
cu_seqlens – [in] 累积序列长度，[batch_size + 1]。
cu_seqlens_padded – [in] QKV 的累积序列偏移量，[batch_size + 1]。
max_seqlen – [in] 用于计算的最大序列长度，它可能 >= i=0,…batch_size-1 的 max(seqlen_i)。
attn_scale – [in] Q * K.T 的缩放因子。
dropout – [in] Dropout 概率。
qkv_layout – [in] QKV 张量的布局。
bias_type – [in] Bias 类型。
attn_mask_type – [in] Attention mask 类型。
window_size_left – [in] 滑动窗口大小 (左半部分)。
window_size_right – [in] 滑动窗口大小 (右半部分)。
deterministic – [in] 是否以确定性行为执行。
workspace – [in] Workspace 张量。
stream – [in] 用于此操作的 CUDA stream。

void nvte_fused_attn_fwd_kvpacked(const NVTETensor Q, const NVTETensor KV, const NVTETensor Bias, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, const NVTETensor rng_state, size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, float attn_scale, float dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, int64_t window_size_left, int64_t window_size_right, NVTETensor workspace, cudaStream_t stream)

Compute dot product attention with packed KV input.

计算

P = Q * Transpose(K) + Bias
S = ScaleMaskSoftmax(P)
D = Dropout(S)
O = D * Transpose(V)

支持矩阵

| backend | precision |                 qkv layout                  |           bias           |                 mask                  | dropout |  sequence length  | head_dim         |
|   0     | FP16/BF16 |            BSHD_BS2HD,SBHD_SB2HD            |   NO/POST_SCALE_BIAS     | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   | <= 512, % 64 == 0 |    64            |
|   1     | FP16/BF16 | BSHD_BS2HD,BSHD_BSH2D,SBHD_SB2HD,SBHD_SBH2D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   |  > 512, % 64 == 0 | <= 128, % 8 == 0 |

注释

张量 cu_seqlens_q_padded 和 cu_seqlens_kv_padded 帮助识别张量 Q、K、V 和 O 中不同序列的正确偏移量。当 QKV 格式 (nvte_get_qkv_format(qkv_layout)) 为 bshd 或 sbhd 时，偏移张量不用于 attention 计算，可以设置为空 NVTETensor。当 QKV 格式为 thd 时，这些张量应遵循以下规则。当序列之间没有填充时，偏移张量应分别等于 cu_seqlens_q 和 cu_seqlens_kv。当序列之间存在填充时，用户有责任根据需要调整偏移量。例如，一个包含 4 个序列的张量 [a, PAD, b, b, c, PAD, PAD, d, d] 应该有 cu_seqlens = [0, 1, 3, 4, 6] 和 cu_seqlens_padded= [0, 2, 4, 7, 9]。

参数:

Q – [in] Q 张量，HD 布局。
KV – [in] KV 张量，2HD 或 H2D 布局。
Bias – [in] Bias 张量。
S – [inout] S 张量。
O – [out] 输出 O 张量。
Aux_CTX_Tensors – [out] 训练时的辅助输出张量，例如 M、ZInv、rng_state。
cu_seqlens_q – [in] Q 的累积序列长度，[batch_size + 1]。
cu_seqlens_kv – [in] KV 的累积序列长度，[batch_size + 1]。
cu_seqlens_q_padded – [in] Q 的累积序列偏移量，[batch_size + 1]。
cu_seqlens_kv_padded – [in] KV 的累积序列偏移量，[batch_size + 1]。
rng_state – [in] CUDA 随机数生成器的种子和偏移量。
max_seqlen_q – [in] 用于计算 Q 的最大序列长度。它可能 >= i=0,…batch_size-1 的 max(seqlen_q_i)。
max_seqlen_kv – [in] 用于计算 KV 的最大序列长度。它可能 >= i=0,…batch_size-1 的 max(seqlen_kv_i)。
is_training – [in] 是否处于训练模式或推理模式。
attn_scale – [in] Q * K.T 的缩放因子。
dropout – [in] Dropout 概率。
qkv_layout – [in] QKV 张量的布局。
bias_type – [in] Bias 类型。
attn_mask_type – [in] Attention mask 类型。
window_size_left – [in] 滑动窗口大小 (左半部分)。
window_size_right – [in] 滑动窗口大小 (右半部分)。
deterministic – [in] 是否以确定性行为执行。
workspace – [in] Workspace 张量。
stream – [in] 用于此操作的 CUDA stream。

void nvte_fused_attn_bwd_kvpacked(const NVTETensor Q, const NVTETensor KV, const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, NVTETensor dKV, NVTETensor dBias, const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, int64_t window_size_left, int64_t window_size_right, bool deterministic, NVTETensor workspace, cudaStream_t stream)

Compute the backward of the dot product attention with packed KV input.

支持矩阵

| backend | precision |                 qkv layout                  |           bias           |                 mask                  | dropout |  sequence length  | head_dim         |
|   0     | FP16/BF16 |            BSHD_BS2HD,SBHD_SB2HD            |   NO/POST_SCALE_BIAS     | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   | <= 512, % 64 == 0 |    64            |
|   1     | FP16/BF16 | BSHD_BS2HD,BSHD_BSH2D,SBHD_SB2HD,SBHD_SBH2D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   |  > 512, % 64 == 0 | <= 128, % 8 == 0 |

注释

张量 cu_seqlens_q_padded 和 cu_seqlens_kv_padded 帮助识别张量 Q、K、V 和 O 中不同序列的正确偏移量。当 QKV 格式 (nvte_get_qkv_format(qkv_layout)) 为 bshd 或 sbhd 时，偏移张量不用于 attention 计算，可以设置为空 NVTETensor。当 QKV 格式为 thd 时，这些张量应遵循以下规则。当序列之间没有填充时，偏移张量应分别等于 cu_seqlens_q 和 cu_seqlens_kv。当序列之间存在填充时，用户有责任根据需要调整偏移量。例如，一个包含 4 个序列的张量 [a, PAD, b, b, c, PAD, PAD, d, d] 应该有 cu_seqlens = [0, 1, 3, 4, 6] 和 cu_seqlens_padded= [0, 2, 4, 7, 9]。

参数:

Q – [in] Q 张量，HD 布局。
KV – [in] KV 张量，H2D 或 2HD 布局。
O – [in] 前向传播的 O 张量。
dO – [in] O 张量的梯度。
S – [in] S 张量。
dP – [inout] P 张量的梯度。
Aux_CTX_Tensors – [in] 训练模式下上下文的辅助张量，例如 M、ZInv、rng_state。
dQ – [out] Q 张量的梯度。
dKV – [out] KV 张量的梯度。
dBias – [out] Bias 张量的梯度。
cu_seqlens_q – [in] Q 的累积序列长度，[batch_size + 1]。
cu_seqlens_kv – [in] KV 的累积序列长度，[batch_size + 1]。
cu_seqlens_q_padded – [in] Q 的累积序列偏移量，[batch_size + 1]。
cu_seqlens_kv_padded – [in] KV 的累积序列偏移量，[batch_size + 1]。
max_seqlen_q – [in] 用于计算 Q 的最大序列长度。它可能 >= i=0,…batch_size-1 的 max(seqlen_q_i)。
max_seqlen_kv – [in] 用于计算 KV 的最大序列长度。它可能 >= i=0,…batch_size-1 的 max(seqlen_kv_i)。
attn_scale – [in] Q * K.T 的缩放因子。
dropout – [in] Dropout 概率。
qkv_layout – [in] QKV 张量的布局。
bias_type – [in] Bias 类型。
attn_mask_type – [in] Attention mask 类型。
window_size_left – [in] 滑动窗口大小 (左半部分)。
window_size_right – [in] 滑动窗口大小 (右半部分)。
deterministic – [in] 是否以确定性行为执行。
workspace – [in] Workspace 张量。
stream – [in] 用于此操作的 CUDA stream。

void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, const NVTETensor rng_state, size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, float attn_scale, float dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, int64_t window_size_left, int64_t window_size_right, NVTETensor workspace, cudaStream_t stream)

Compute dot product attention with separate Q, K and V.

计算

P = Q * Transpose(K) + Bias
S = ScaleMaskSoftmax(P)
D = Dropout(S)
O = D * Transpose(V)

支持矩阵

| backend | precision |                qkv layout                   |           bias           |                 mask                  | dropout |  sequence length  | head_dim         |
|   0     | FP16/BF16 |     BS3HD,SB3HD,BSHD_BS2HD,SBHD_SB2HD       |   NO/POST_SCALE_BIAS     | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   | <= 512, % 64 == 0 |    64            |
|   1     | FP16/BF16 |          BS3HD,SB3HD,BSH3D,SBH3D            | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   |  > 512, % 64 == 0 | <= 128, % 8 == 0 |
|         |           | BSHD_BS2HD,BSHD_BSH2D,SBHD_SB2HD,SBHD_SBH2D |                          |                                       |         |                   |                  |
|         |           |       BSHD_BSHD_BSHD,SBHD_SBHD_SBHD         |                          |                                       |         |                   |                  |
|   2     |   FP8     |                 T3HD                        |          NO_BIAS         |               PADDING_MASK            |   Yes   | <= 512, % 64 == 0 |    64            |

注释

张量 cu_seqlens_q_padded 和 cu_seqlens_kv_padded 帮助识别张量 Q、K、V 和 O 中不同序列的正确偏移量。当 QKV 格式 (nvte_get_qkv_format(qkv_layout)) 为 bshd 或 sbhd 时，偏移张量不用于 attention 计算，可以设置为空 NVTETensor。当 QKV 格式为 thd 时，这些张量应遵循以下规则。当序列之间没有填充时，偏移张量应分别等于 cu_seqlens_q 和 cu_seqlens_kv。当序列之间存在填充时，用户有责任根据需要调整偏移量。例如，一个包含 4 个序列的张量 [a, PAD, b, b, c, PAD, PAD, d, d] 应该有 cu_seqlens = [0, 1, 3, 4, 6] 和 cu_seqlens_padded= [0, 2, 4, 7, 9]。

参数:

Q – [in] Q 张量。
K – [in] K 张量。
V – [in] V 张量。
Bias – [in] Bias 张量。
S – [inout] S 张量。
O – [out] 输出 O 张量。
Aux_CTX_Tensors – [out] 训练时的辅助输出张量，例如 M、ZInv、rng_state。
cu_seqlens_q – [in] Q 的累积序列长度，[batch_size + 1]。
cu_seqlens_kv – [in] K 和 V 的累积序列长度，[batch_size + 1]。
cu_seqlens_q_padded – [in] Q 的累积序列偏移量，[batch_size + 1]。
cu_seqlens_kv_padded – [in] KV 的累积序列偏移量，[batch_size + 1]。
rng_state – [in] CUDA 随机数生成器的种子和偏移量。
max_seqlen_q – [in] 用于计算 Q 的最大序列长度。它可能 >= i=0,…batch_size-1 的 max(seqlen_q_i)。
max_seqlen_kv – [in] 用于计算 K 和 V 的最大序列长度。它可能 >= i=0,…batch_size-1 的 max(seqlen_kv_i)。
is_training – [in] 是否处于训练模式或推理模式。
attn_scale – [in] Q * K.T 的缩放因子。
dropout – [in] Dropout 概率。
qkv_layout – [in] QKV 张量的布局。
bias_type – [in] Bias 类型。
attn_mask_type – [in] Attention mask 类型。
window_size_left – [in] 滑动窗口大小 (左半部分)。
window_size_right – [in] 滑动窗口大小 (右半部分)。
workspace – [in] Workspace 张量。
stream – [in] 用于此操作的 CUDA stream。

void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, NVTETensor dK, NVTETensor dV, NVTETensor dBias, const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, int64_t window_size_left, int64_t window_size_right, bool deterministic, NVTETensor workspace, cudaStream_t stream)

Compute the backward of the dot product attention with separate Q, K and V.

支持矩阵

| backend | precision |                qkv layout                   |           bias           |                 mask                  | dropout |  sequence length  | head_dim         |
|   0     | FP16/BF16 |     BS3HD,SB3HD,BSHD_BS2HD,SBHD_SB2HD       |   NO/POST_SCALE_BIAS     | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   | <= 512, % 64 == 0 |    64            |
|   1     | FP16/BF16 |          BS3HD,SB3HD,BSH3D,SBH3D            | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK |   Yes   |  > 512, % 64 == 0 | <= 128, % 8 == 0 |
|         |           | BSHD_BS2HD,BSHD_BSH2D,SBHD_SB2HD,SBHD_SBH2D |                          |                                       |         |                   |                  |
|         |           |       BSHD_BSHD_BSHD,SBHD_SBHD_SBHD         |                          |                                       |         |                   |                  |
|   2     |   FP8     |                 T3HD                        |          NO_BIAS         |               PADDING_MASK            |   Yes   | <= 512, % 64 == 0 |    64            |

注释

张量 cu_seqlens_q_padded 和 cu_seqlens_kv_padded 帮助识别张量 Q、K、V 和 O 中不同序列的正确偏移量。当 QKV 格式 (nvte_get_qkv_format(qkv_layout)) 为 bshd 或 sbhd 时，偏移张量不用于 attention 计算，可以设置为空 NVTETensor。当 QKV 格式为 thd 时，这些张量应遵循以下规则。当序列之间没有填充时，偏移张量应分别等于 cu_seqlens_q 和 cu_seqlens_kv。当序列之间存在填充时，用户有责任根据需要调整偏移量。例如，一个包含 4 个序列的张量 [a, PAD, b, b, c, PAD, PAD, d, d] 应该有 cu_seqlens = [0, 1, 3, 4, 6] 和 cu_seqlens_padded= [0, 2, 4, 7, 9]。

参数:

Q – [in] Q 张量。
K – [in] K 张量。
V – [in] V 张量。
O – [in] 前向传播的 O 张量。
dO – [in] O 张量的梯度。
S – [in] S 张量。
dP – [inout] P 张量的梯度。
Aux_CTX_Tensors – [in] 训练模式下上下文的辅助张量，例如 M、ZInv、rng_state。
dQ – [out] Q 张量的梯度。
dK – [out] K 张量的梯度。
dV – [out] V 张量的梯度。
dBias – [out] Bias 张量的梯度。
cu_seqlens_q – [in] Q 的累积序列长度，[batch_size + 1]。
cu_seqlens_kv – [in] K 和 V 的累积序列长度，[batch_size + 1]。
cu_seqlens_q_padded – [in] Q 的累积序列偏移量，[batch_size + 1]。
cu_seqlens_kv_padded – [in] KV 的累积序列偏移量，[batch_size + 1]。
max_seqlen_q – [in] 用于计算 Q 的最大序列长度。它可能 >= i=0,…batch_size-1 的 max(seqlen_q_i)。
max_seqlen_kv – [in] 用于计算 K 和 V 的最大序列长度。它可能 >= i=0,…batch_size-1 的 max(seqlen_kv_i)。
attn_scale – [in] Q * K.T 的缩放因子。
dropout – [in] Dropout 概率。
qkv_layout – [in] QKV 张量的布局。
bias_type – [in] Bias 类型。
attn_mask_type – [in] Attention mask 类型。
window_size_left – [in] 滑动窗口大小 (左半部分)。
window_size_right – [in] 滑动窗口大小 (右半部分)。
deterministic – [in] 是否以确定性行为执行。
workspace – [in] Workspace 张量。
stream – [in] 用于此操作的 CUDA stream。