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 布局

枚举器 NVTE_TH3D

TH3D 布局

枚举器 NVTE_THD_T2HD

THD_T2HD 布局

枚举器 NVTE_THD_TH2D

THD_TH2D 布局

枚举器 NVTE_THD_THD_THD

THD_THD_THD 布局

枚举 NVTE_QKV_Layout_Group

QKV 布局组。

值

枚举器 NVTE_3HD

3HD QKV 布局，例如 BS3HD、SB3HD、T3HD

枚举器 NVTE_H3D

H3D QKV 布局，例如 BSH3D、SBH3D、TH3D

枚举器 NVTE_HD_2HD

HD_2HD QKV 布局，例如 BSHD_BS2HD、SBHD_SB2HD、THD_T2HD

枚举器 NVTE_HD_H2D

HD_H2D QKV 布局，例如 BSHD_BSH2D、SBHD_SBH2D、THD_TH2D

枚举器 NVTE_HD_HD_HD

HD_HD_HD QKV 布局，例如 BSHD_BSHD_BSHD、SBHD_SBHD_SBHD、THD_THD_THD

枚举 NVTE_QKV_Format

QKV 格式。

值

枚举器 NVTE_SBHD

SBHD QKV 格式，例如 SB3HD、SBH3D、SBHD_SB2HD、SBHD_SBH2D、SBHD_SBHD_SBHD

枚举器 NVTE_BSHD

BSHD QKV 格式，例如 BS3HD、BSH3D、BSHD_BS2HD、BSHD_BSH2D、BSHD_BSHD_BSHD

枚举器 NVTE_THD

THD QKV 格式，例如 T3HD、TH3D、THD_T2HD、THD_TH2D、THD_THD_THD

枚举 NVTE_Bias_Type

偏置类型。

值

枚举器 NVTE_NO_BIAS

无偏置

枚举器 NVTE_PRE_SCALE_BIAS

缩放前偏置

枚举器 NVTE_POST_SCALE_BIAS

缩放后偏置

枚举器 NVTE_ALIBI

ALiBi

枚举 NVTE_Mask_Type

注意力掩码类型。

值

枚举器 NVTE_NO_MASK

无掩码

枚举器 NVTE_PADDING_MASK

填充注意力掩码

枚举器 NVTE_CAUSAL_MASK

因果注意力掩码

枚举器 NVTE_PADDING_CAUSAL_MASK

填充和因果注意力掩码

枚举 NVTE_Fused_Attn_Backend

融合注意力后端。

值

枚举器 NVTE_No_Backend

不支持的后端

枚举器 NVTE_F16_max512_seqlen

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

枚举器 NVTE_F16_arbitrary_seqlen

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

枚举器 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)

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

参数

  • 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 中的 head 数量。

  • num_gqa_groups – [in] K、V 中的 head 数量。

  • max_seqlen_q – [in] Q 的序列长度。

  • max_seqlen_kv – [in] K、V 的序列长度。

  • head_dim – [in] Q、K、V 的 head 维度。

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 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, NVTETensor workspace, cudaStream_t stream)

计算带有 packed 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            |

参数

  • QKV – [in] packed 格式的 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]。

  • 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] 注意力掩码类型。

  • workspace – [in] 工作区张量。

  • stream – [in] 用于此操作的 CUDA 流。

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, 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, NVTETensor workspace, cudaStream_t stream)

计算带有 packed QKV 输入的点积注意力的反向传播。

支持矩阵

| 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            |

参数

  • QKV – [in] packed 格式的 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]。

  • 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] 注意力掩码类型。

  • workspace – [in] 工作区张量。

  • stream – [in] 用于此操作的 CUDA 流。

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 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, NVTETensor workspace, cudaStream_t stream)

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

计算：

  • 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 |

参数

  • Q – [输入] Q 张量，采用 HD 布局。

  • KV – [输入] KV 张量，采用 2HD 或 H2D 布局。

  • Bias – [in] Bias 张量。

  • S – [inout] S 张量。

  • O – [out] 输出 O 张量。

  • Aux_CTX_Tensors – [out] 训练时的辅助输出张量，例如 M、ZInv、rng_state。

  • cu_seqlens_q – [输入] Q 的累积序列长度，[batch_size + 1]。

  • cu_seqlens_kv – [输入] KV 的累积序列长度，[batch_size + 1]。

  • rng_state – [in] CUDA 随机数生成器的种子和偏移量。

  • max_seqlen_q – [输入] 用于计算 Q 的最大序列长度。它可以 >= i=0,…batch_size-1 的 max(seqlen_q_i)。

  • max_seqlen_kv – [输入] 用于计算 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] 注意力掩码类型。

  • workspace – [in] 工作区张量。

  • stream – [in] 用于此操作的 CUDA 流。

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, 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, NVTETensor workspace, cudaStream_t stream)

计算带有打包 KV 输入的点积注意力的反向传播。

支持矩阵

| 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 |

参数

  • Q – [输入] Q 张量，采用 HD 布局。

  • KV – [输入] KV 张量，采用 H2D 或 2HD 布局。

  • O – [in] 来自正向传播的 O 张量。

  • dO – [in] O 张量的梯度。

  • S – [in] S 张量。

  • dP – [inout] P 张量的梯度。

  • Aux_CTX_Tensors – [in] 训练模式下来自上下文的辅助张量，例如 M、ZInv、rng_state。

  • dQ – [输出] Q 张量的梯度。

  • dKV – [输出] KV 张量的梯度。

  • dBias – [out] Bias 张量的梯度。

  • cu_seqlens_q – [输入] Q 的累积序列长度，[batch_size + 1]。

  • cu_seqlens_kv – [输入] KV 的累积序列长度，[batch_size + 1]。

  • max_seqlen_q – [输入] 用于计算 Q 的最大序列长度。它可以 >= i=0,…batch_size-1 的 max(seqlen_q_i)。

  • max_seqlen_kv – [输入] 用于计算 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] 注意力掩码类型。

  • workspace – [in] 工作区张量。

  • stream – [in] 用于此操作的 CUDA 流。

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 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, NVTETensor workspace, cudaStream_t stream)

计算带有分离的 Q、K 和 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            |

参数

  • Q – [输入] Q 张量。

  • K – [输入] K 张量。

  • V – [输入] V 张量。

  • Bias – [in] Bias 张量。

  • S – [inout] S 张量。

  • O – [out] 输出 O 张量。

  • Aux_CTX_Tensors – [out] 训练时的辅助输出张量，例如 M、ZInv、rng_state。

  • cu_seqlens_q – [输入] Q 的累积序列长度，[batch_size + 1]。

  • cu_seqlens_kv – [输入] K 和 V 的累积序列长度，[batch_size + 1]。

  • rng_state – [in] CUDA 随机数生成器的种子和偏移量。

  • max_seqlen_q – [输入] 用于计算 Q 的最大序列长度。它可以 >= i=0,…batch_size-1 的 max(seqlen_q_i)。

  • max_seqlen_kv – [输入] 用于计算 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 – [输入] QKV 张量的布局。

  • bias_type – [in] Bias 类型。

  • attn_mask_type – [in] 注意力掩码类型。

  • workspace – [in] 工作区张量。

  • stream – [in] 用于此操作的 CUDA 流。

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, 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, NVTETensor workspace, cudaStream_t stream)

计算具有独立 Q、K 和 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            |

参数

  • Q – [输入] Q 张量。

  • K – [输入] K 张量。

  • V – [输入] V 张量。

  • O – [in] 来自正向传播的 O 张量。

  • dO – [in] O 张量的梯度。

  • S – [in] S 张量。

  • dP – [inout] P 张量的梯度。

  • Aux_CTX_Tensors – [in] 训练模式下来自上下文的辅助张量，例如 M、ZInv、rng_state。

  • dQ – [输出] Q 张量的梯度。

  • dK – [输出] K 张量的梯度。

  • dV – [输出] V 张量的梯度。

  • dBias – [out] Bias 张量的梯度。

  • cu_seqlens_q – [输入] Q 的累积序列长度，[batch_size + 1]。

  • cu_seqlens_kv – [输入] K 和 V 的累积序列长度，[batch_size + 1]。

  • max_seqlen_q – [输入] 用于计算 Q 的最大序列长度。它可以 >= i=0,…batch_size-1 的 max(seqlen_q_i)。

  • max_seqlen_kv – [输入] 用于计算 K 和 V 的最大序列长度。对于 i=0,…batch_size-1，它可能 >= max(seqlen_kv_i)。

  • attn_scale – [in] Q * K.T 的缩放因子。

  • dropout – [in] Dropout 概率。

  • qkv_layout – [输入] QKV 张量的布局。

  • bias_type – [in] Bias 类型。

  • attn_mask_type – [in] 注意力掩码类型。

  • workspace – [in] 工作区张量。

  • stream – [in] 用于此操作的 CUDA 流。