ops(8): self-attention 的 CUDA 实现及优化 (下)

原文链接: https://zhuanlan.zhihu.com/p/695898274
作者: 紫气东来
发布时间: 2024-05-09
文章ID: 695898274

重要说明

本文档对应的文章ID 695898274 实际上是该系列的"下篇"（ops(8)），而非"上篇"（ops(7)）。文章主要讨论：

1. 基于 cuDNN 的 self-attention 实现
2. self-attention 的反向传播实现及优化

一、使用 cuDNN 接口实现

1.1 cuDNN 概览及其 attention 实现

cuDNN 简介

cuDNN（NVIDIA CUDA Deep Neural Network Library）是深度神经网络算子层级GPU加速库集合，提供了深度学习算法中常见算子的高效实现。它是许多上层推理引擎（如TensorRT、TVM）底层调优的算子备选实现。

cuDNN 常见算子

• 卷积: 前向和反向卷积
• 矩阵运算: 矩阵乘法
• 池化: 前向和反向池化
• Softmax: 前向和反向 Softmax
• 激活函数: relu、tanh、sigmoid、elu、gelu、softplus、swish
• 归一化: BN、IN、LN、LRN、LCN
• 基础运算: 逐点数学计算、张量变换

Attention 实现特点

cuDNN 中的 Scaled Dot Product Attention 实现采用了 FlashAttention-2 算法，提供 Python 和 C++ 两种接口。

Python 接口参数

Args:
    q (cudnn_tensor): Query 数据
    k (cudnn_tensor): Key 数据
    v (cudnn_tensor): Value 数据
    is_inference (bool): 是否为推理步骤
    attn_scale (Optional[Union[float, cudnn_tensor]]): Attention 缩放因子
    bias (Optional[cudnn_tensor]): Attention 偏置数据
    use_alibi_mask (Optional[bool]): 是否使用 alibi mask
    use_padding_mask (Optional[bool]): 是否使用 padding mask
    seq_len_q (Optional[cudnn_tensor]): Query 序列长度
    seq_len_kv (Optional[cudnn_tensor]): Key 序列长度
    use_causal_mask (Optional[bool]): 是否使用因果 mask
    dropout (Optional): Dropout 配置
    compute_data_type (Optional[cudnn.data_type]): 计算数据类型
    name (Optional[str]): 操作名称

Returns:
    o (cudnn_tensor): 输出数据
    stats (Optional[cudnn_tensor]): Softmax 统计信息（训练时）

C++ API 接口

// 返回 [output, softmax_stats]
std::array<std::shared_ptr<Tensor_attributes>, 2> 
sdpa(std::shared_ptr<Tensor_attributes> q,
     std::shared_ptr<Tensor_attributes> k,
     std::shared_ptr<Tensor_attributes> v,
     SDPA_attributes options);

1.2 cuDNN 接口调用实现（V10）

实现步骤

步骤1: 构造输入输出 tensor

using graph_tensors_fwd = std::tuple<
    std::shared_ptr<fe::graph::Graph>,
    std::shared_ptr<fe::graph::Tensor_attributes>,  // Q
    std::shared_ptr<fe::graph::Tensor_attributes>,  // K
    std::shared_ptr<fe::graph::Tensor_attributes>,  // V
    std::shared_ptr<fe::graph::Tensor_attributes>,  // Attn_scale
    std::shared_ptr<fe::graph::Tensor_attributes>,  // O
    std::shared_ptr<fe::graph::Tensor_attributes>>; // Stats

using cache_type_fwd = std::unordered_map<std::size_t, graph_tensors_fwd>;

步骤2: 基于 cuDNN frontend 构造图

关键配置：

• 设置数据类型：IO 使用 FP16/BF16，中间计算使用 FP32
• QKV 张量维度：(B, H, T, HS)，支持直接处理交错存储的数据
• 设置 attention scale 和 causal mask
• 输出维度：O 为 (B, H, T, HS)，stats 为 (B, H, T, 1)

步骤3: 完整 kernel 实现

核心流程：

1. FP32 转换为 FP16/BF16（首次运行验证时）
2. 从缓存获取或构建计算图
3. 准备 tensor 指针和参数
4. 分配 workspace（按需动态分配，最大 256MB）
5. 执行计算图
6. 可选的 FP16/BF16 转回 FP32（验证时）

性能数据（RTX 4090 + CUDA 12.4）

cuDNN 实现性能：

block_size   32 | time 0.169061 ms
block_size   64 | time 0.165807 ms
block_size  128 | time 0.167423 ms
block_size  256 | time 0.165734 ms
block_size  512 | time 0.167426 ms

结论: cuDNN 的实现具有显著的性能优势，相比手工实现的 V1~V5 版本有数量级的提升。

二、self-attention 的反向实现

2.1 反向过程的数学推导

矩阵乘法求导

对于矩阵乘法 Y = WX，目标函数为 φ，记：

• dY = ∂φ/∂Y
• dW = ∂φ/∂W
• dX = ∂φ/∂X

则有：

dW = dY · X^T
dX = W^T · dY

Softmax 求导

设 X = [x₁, x₂, ..., xₙ]，Y = softmax(X) = [y₁, y₂, ..., yₙ]

其中 yᵢ = e^(xᵢ) / Σⱼ e^(xⱼ)，且 Σᵢ yᵢ = 1

求导结果：

∂yᵢ/∂xⱼ = {
    yᵢ(1 - yⱼ),  if i = j
    yᵢ(0 - yⱼ),  if i ≠ j
}

Attention 反向传播公式

令：

• P = Softmax(S)
• S = (Q·K^T) / √dₖ
• O = P·V

则 attention 的反向过程为：

1. dP = dO·V^T
2. dV = P^T·dO
3. dS = Softmax 反向(dP, P)
4. dQ = dS·K
5. dK = dS^T·Q

2.2 反向过程的 CPU 实现

CPU 实现按照以上推导逻辑，分为四个主要步骤：

步骤1: 反向通过 value 累积（backward pass 4）

• 计算 datt 和 dvalue

步骤2: Softmax 反向（backward pass 2 & 3）

• 利用 softmax 导数公式计算 dpreatt
• 应用 scale 因子

步骤3: Query @ Key 矩阵乘反向（backward pass 1）

• 计算 dquery 和 dkey

步骤4: 应用 causal mask

• 只计算 t2 <= t 的部分

2.3 CUDA 简单实现（V1）

V1 版本利用 cuBLAS 库函数计算矩阵乘法，主要步骤：

1. Unpermute 反向: 将 dout 转换为 dvaccum
2. 计算 datt: 使用 cublasSgemmStridedBatched
3. 计算 dv: 使用 cublasSgemmStridedBatched
4. Softmax 反向: 计算 dpreatt
5. 计算 dq 和 dk: 使用 cublasSgemmStridedBatched
6. Permute 反向: 将梯度写回 dinp

V1 性能数据（RTX 4090）:

block_size   32 | time 7084.399902 ms
block_size   64 | time 7067.519531 ms
block_size  128 | time 7077.885254 ms
block_size  256 | time 9231.899414 ms
block_size  512 | time 8673.948242 ms
block_size 1024 | time 14843.697266 ms

2.4 反向过程的优化（V2~V8）

由于矩阵乘法已使用 cuBLAS 优化，后续优化主要围绕 Softmax 反向传播展开。

V2: 基础并行化

策略: 在 t, b, h 维度上并行

性能:

block_size   32 | time 273.880554 ms
block_size   64 | time 271.534363 ms
block_size  128 | time 271.852448 ms
block_size  256 | time 287.613922 ms
block_size  512 | time 339.183411 ms
block_size 1024 | time 388.895203 ms

提升: 相比 V1 提升约 26x

V3: 协作组优化

策略: 在 t, b, h 维度并行 + 使用协作组（Cooperative Groups）

性能:

block_size   32 | time 17.532518 ms
block_size   64 | time 15.095194 ms
block_size  128 | time 14.864792 ms
block_size  256 | time 14.368258 ms
block_size  512 | time 14.173902 ms
block_size 1024 | time 14.197658 ms

提升: 相比 V2 提升约 18x，相比 V1 提升约 470x

V4: 循环展开优化

策略: V3 基础上增加 UNROLL 操作

性能:

block_size   32 | time 10.701917 ms
block_size   64 | time 11.321651 ms
block_size  128 | time 10.746265 ms
block_size  256 | time 10.383463 ms
block_size  512 | time 10.818865 ms
block_size 1024 | time 10.808935 ms

提升: 相比 V3 提升约 1.4x

V5: 特殊情况优化

策略: 优化 V4 版本的特殊情况处理

性能:

block_size   32 | time 6.028288 ms
block_size   64 | time 6.770285 ms
block_size  128 | time 6.253971 ms
block_size  256 | time 6.313370 ms
block_size  512 | time 6.458060 ms
block_size 1024 | time 6.442598 ms

提升: 相比 V4 提升约 1.7x

V6: 循环重构和内存访问优化

策略: 通过循环重构和内存访问模式优化提高性能

性能:

block_size   32 | time 3.683847 ms
block_size   64 | time 2.985779 ms
block_size  128 | time 2.968576 ms
block_size  256 | time 3.310278 ms
block_size  512 | time 4.242841 ms
block_size 1024 | time 6.089725 ms

提升: 相比 V5 提升约 2.1x

V7: 数学简化和协作组规约

策略: 简化数学计算，使用协作组的规约操作

性能:

block_size   32 | time 1.766189 ms
block_size   64 | time 1.761990 ms
block_size  128 | time 1.760563 ms
block_size  256 | time 1.758925 ms
block_size  512 | time 1.760154 ms
block_size 1024 | time 1.897673 ms

提升: 相比 V6 提升约 1.7x，相比 V1 提升约 4000x

V8: 额外 Tricks

策略: 在 V7 基础上新增优化技巧

性能:

block_size   32 | time 1.781242 ms
block_size   64 | time 1.778285 ms
block_size  128 | time 1.768653 ms
block_size  256 | time 1.762298 ms
block_size  512 | time 1.760768 ms
block_size 1024 | time 1.875651 ms

提升: 与 V7 性能基本持平

三、性能总结与分析

3.1 优化版本性能对比

3.2 关键优化技术分析

1. 并行化策略

• V1 问题: 串行计算，未充分利用 GPU 并行能力
• V2 改进: 在 batch、head、sequence 维度并行
• 效果: 26x 性能提升

2. 协作组（Cooperative Groups）

• 作用: 提供更灵活的线程同步和通信机制
• 优势: 相比传统 __syncthreads()，支持更细粒度的控制
• 效果: 18x 性能提升（V2→V3）

3. 循环展开（Loop Unrolling）

• 作用: 减少循环控制开销，增加指令级并行
• 实现: 使用 #pragma unroll 指令
• 效果: 1.4x 性能提升（V3→V4）

4. 内存访问优化

• 策略:
  • 合并内存访问（Coalesced Memory Access）
  • 减少全局内存访问次数
  • 优化访问模式以提高缓存命中率
• 效果: 2.1x 性能提升（V5→V6）

5. 数学计算简化

• 策略:
  • 利用 softmax 导数的数学性质简化计算
  • 减少冗余计算
  • 使用协作组规约操作
• 效果: 1.7x 性能提升（V6→V7）

3.3 性能瓶颈分析

Softmax 反向传播是主要瓶颈

• V1 中矩阵乘法已使用 cuBLAS 优化
• V2~V8 的优化都集中在 Softmax 反向
• 最终实现了 4000x 的性能提升

Block Size 选择

• 大多数版本在 block_size=256 或 512 时性能最佳
• 过小的 block size 导致占用率不足
• 过大的 block size 导致寄存器压力和共享内存限制

四、实现要点总结

4.1 cuDNN 实现要点

1. 使用 FlashAttention-2 算法: 高效的内存访问模式
2. 支持多种 mask 类型: causal、padding、alibi
3. 混合精度计算: IO 用 FP16/BF16，计算用 FP32
4. 动态 workspace 管理: 按需分配，最大 256MB
5. 计算图缓存: 避免重复构建图的开销

4.2 反向传播实现要点

1. 数学推导正确性:

   • 矩阵乘法求导: dW = dY·X^T, dX = W^T·dY
   • Softmax 求导: 区分 i=j 和 i≠j 两种情况

2. 计算顺序:

   • 从输出到输入反向传播
   • 先计算 dV 和 datt，再计算 dQ 和 dK

3. 内存管理:

   • 复用中间结果
   • 避免不必要的内存分配

4. 并行策略:

   • 在 batch、head、sequence 维度并行
   • 使用协作组进行线程间通信

五、代码资源

完整代码实现：

• GitHub 仓库: https://github.com/ifromeast/cuda_learning
• 前向实现: /blob/main/04_transformer/ops/attention_forward.cu
• 反向实现: /blob/main/04_transformer/ops/attention_backward.cu

六、参考资料

1. Karpathy llm.c - attention_forward.cu
2. Karpathy llm.c - attention_backward.cu
3. cuDNN Frontend - Attention Operations
4. cuDNN Frontend - MHA Sample
5. Matrix-Matrix Product Derivative
6. Matrix Cookbook
7. 反向传播之一：softmax函数

七、关键收获

1. cuDNN 的优势: 使用 FlashAttention-2 实现，性能显著优于手工实现
2. 优化的重要性: 从 V1 到 V7，通过系统化优化实现了 4000x 的性能提升
3. 优化方向: 并行化 → 协作组 → 循环展开 → 内存优化 → 数学简化
4. 工程实践:
   • 先用库函数（cuBLAS）优化矩阵乘法
   • 再针对特定算子（Softmax）进行深度优化
   • 逐步迭代，每次优化都有明确的性能提升

注: 本文为系列文章的下篇，主要讨论 cuDNN 实现和反向传播。上篇内容请参考 ops(7) 文章。