ops(8)：self-attention 的 CUDA 实现及优化 (下) - 深度分析

原文作者: 紫气东来
发布时间: 2024-05-09
原文链接: https://zhuanlan.zhihu.com/p/696197013
分析日期: 2026-05-05

目录

1. 文章概述
2. cuDNN 实现详解
3. 反向传播实现
4. 性能对比分析
5. 核心要点总结

文章概述

本文是 self-attention CUDA 实现系列的下篇，承接上篇的基础实现与优化，主要聚焦于两个核心主题：

1. 使用 cuDNN 库实现高性能 Attention：基于 FlashAttention-2 算法的工业级实现
2. 反向传播的 CUDA 实现与优化：从基础版本到高度优化版本的完整演进（V1-V8）

文章展示了从 7000ms 优化到 1.7ms 的惊人性能提升过程，体现了 CUDA 优化的系统性方法论。

一、cuDNN 实现详解

1.1 cuDNN 库概览

什么是 cuDNN？

cuDNN (NVIDIA CUDA Deep Neural Network Library) 是深度神经网络算子层级的 GPU 加速库集合，提供了深度学习算法中常见算子的高效实现。

核心定位：

• 提供预先优化好的高效 CUDA C Kernel 集合
• 成为上层推理引擎（TensorRT、TVM）的底层算子实现
• 面向深度神经网络算子的工业级优化

cuDNN 支持的核心算子

1. 卷积运算：前向和反向卷积
2. 矩阵运算：高效矩阵乘法
3. 池化操作：前向和反向池化
4. Softmax：前向和反向 Softmax
5. 激活函数：ReLU、Tanh、Sigmoid、GELU、Swish、Softplus 等
6. 归一化：BN、IN、LN、LRN、LCN
7. 基础运算：逐点计算、张量变换

Attention 实现特点

• 算法基础：采用 FlashAttention-2 算法
• 接口支持：提供 Python 和 C++ 两种接口
• 性能优势：显著优于手工优化的 CUDA kernel

1.2 cuDNN API 接口详解

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/Value 序列长度
    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++ 接口结构

// 返回 [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);

SDPA_attributes 配置项：

• set_is_inference(bool): 设置推理/训练模式
• set_attn_scale(float/Tensor): 设置缩放因子
• set_bias(Tensor): 设置偏置
• set_causal_mask(bool): 设置因果 mask
• set_dropout(...): 配置 dropout
• set_compute_data_type(DataType_t): 设置计算精度

1.3 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

// 使用缓存机制，因为 graph->build_operation_graph() 很慢
using cache_type_fwd = std::unordered_map<std::size_t, graph_tensors_fwd>;

步骤 2：构造计算图

关键配置：

auto graph = std::make_shared<fe::graph::Graph>();
graph->set_io_data_type(CUDNN_16BIT)              // 输入输出：FP16/BF16
      .set_intermediate_data_typeFLOAT  // 中间计算：FP32
      .set_compute_data_typeFLOAT;      // 计算精度：FP32

Tensor 定义（支持非连续内存布局）：

// QKV 格式：(B, T, 3, NH, HS) - cuDNN 可直接处理无需外部 permute
auto Q = graph->tensorTensor_attributes(
    .set_name("Q")
    .set_dim({B, H, T, HS})
    .set_stride({3 * H * HS * T, HS, 3 * H * HS, 1}));

SDPA 配置：

auto sdpa_options = fe::graph::SDPA_attributes()
    .set_name("flash_attention")
    .set_is_inference(is_inference_only)
    .set_attn_scale(attn_scale)
    .set_causal_mask(true);  // 启用因果 mask

auto [O, stats] = graph->sdpa(Q, K, V, sdpa_options);

步骤 3：执行 Kernel

完整实现：

void attention_forward_cudnn(floatX* out,  // 输出: (B, T, NH, HS)
                             float* stats, // 反向传播统计: (B, NH, T)
                             floatX* inp,  // 输入: (B, T, 3, NH, HS) QKV
                             int B, int T, int C, int NH) {
    int HS = C / NH;
    bool is_inference_only = (stats == nullptr);
    
    // 从缓存获取或首次构建计算图
    auto [graph, Q, K, V, attn_scale, O, softmax_stats] =
        lookup_cache_or_build_graph_fwd(B, NH, T, HS, is_inference_only);
    
    // 准备 tensor 指针
    void* devPtrQ = inp;
    void* devPtrK = (inp + C);
    void* devPtrV = (inp + 2 * C);
    float attn_scale_cpu = 1.0 / sqrtf(HS);
    void* devPtrO = out;
    
    // 构建 variant pack（tensor 到指针的映射）
    std::unordered_map<std::shared_ptr<fe::graph::Tensor_attributes>, void*> 
        variant_pack = {
            {Q, devPtrQ}, {K, devPtrK}, {V, devPtrV}, 
            {attn_scale, &attn_scale_cpu}, {O, devPtrO}
        };
    
    if (!is_inference_only) {
        variant_pack[softmax_stats] = stats;
    }
    
    // 动态分配 workspace（cuDNN 默认最多使用 256MB）
    if (graph->get_workspace_size() > cudnn_workspace_size) {
        if (cudnn_workspace_size > 0) {
            cudaFree(cudnn_workspace);
        }
        cudnn_workspace_size = graph->get_workspace_size();
        cudaMalloc(&cudnn_workspace, cudnn_workspace_size);
    }
    
    // 执行计算图
    graph->execute(cudnn_handle, variant_pack, cudnn_workspace);
}

性能数据（RTX 4090）

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 实现相比手工优化版本具有显著性能优势（约 0.166ms）

二、self-attention 的反向实现

2.1 反向传播数学推导

矩阵乘法求导

对于矩阵乘法 Y = WX，设目标函数为 φ，则梯度为：

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

其中 dY、dW、dX 分别表示 ∂φ/∂Y、∂φ/∂W、∂φ/∂X

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ᵢ · yⱼ, if i ≠ j
}

推导过程（i = j 时）：

∂yᵢ/∂xᵢ = ∂/∂xᵢ (e^(xᵢ) / Σₖ e^(xₖ))
        = [e^(xᵢ) · Σₖ e^(xₖ) - e^(xᵢ) · e^(xᵢ)] / (Σₖ e^(xₖ))²
        = yᵢ - yᵢ²
        = yᵢ(1 - yᵢ)

推导过程（i ≠ j 时）：

∂yᵢ/∂xⱼ = ∂/∂xⱼ (e^(xᵢ) / Σₖ e^(xₖ))
        = [0 · Σₖ e^(xₖ) - e^(xᵢ) · e^(xⱼ)] / (Σₖ e^(xₖ))²
        = -yᵢ · yⱼ

Attention 反向传播流程

前向过程：

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

反向过程（链式法则）：

1. dV 计算：dV = P^T · dO
2. dP 计算：dP = dO · V^T
3. dS 计算（Softmax 反向）：dS = P ⊙ (dP - (dP ⊙ P) · 1)
4. dQ 计算：dQ = (dS · K) / √dₖ
5. dK 计算：dK = (dS^T · Q) / √dₖ

其中 ⊙ 表示逐元素乘法，1 表示全 1 向量

2.2 CPU 参考实现

void attention_backward_cpu(float* dinp, float* dpreatt, float* datt,
                            float* dout, float* inp, float* att,
                            int B, int T, int C, int NH) {
    int C3 = C * 3;
    int hs = C / NH;
    float scale = 1.0 / sqrtf(hs);
    
    for (int b = 0; b < B; b++) {
        for (int t = 0; t < T; t++) {
            for (int h = 0; h < NH; h++) {
                // 反向传播步骤 4：通过 value 累加
                for (int t2 = 0; t2 < T; t2++) {
                    for (int i = 0; i < hs; i++) {
                        datt_bth[t2] += value_t2[i] * dout_bth[i];
                        dvalue_t2[i] += att_bth[t2] * dout_bth[i];
                    }
                }
                
                // 反向传播步骤 2&3：Softmax 反向
                for (int t2 = 0; t2 <= t; t2++) {
                    for (int t3 = 0; t3 <= t; t3++) {
                        float indicator = (t2 == t3) ? 1.0f : 0.0f;
                        float local_derivative = att_bth[t2] * (indicator - att_bth[t3]);
                        dpreatt_bth[t3] += scale * local_derivative * datt_bth[t2];
                    }
                }
                
                // 反向传播步骤 1：Q @ K 矩阵乘法
                for (int t2 = 0; t2 <= t; t2++) {
                    for (int i = 0; i < hs; i++) {
                        dquery_t[i] += key_t2[i] * dpreatt_bth[t2];
                        dkey_t2[i] += query_t[i] * dpreatt_bth[t2];
                    }
                }
            }
        }
    }
}

2.3 CUDA 优化演进（V1-V8）

V1：基础 CUDA 实现

策略：利用 cuBLAS 处理矩阵乘法，手工实现 Softmax 反向

template<class SoftmaxKernel>
void attention_backward1(float* dinp, float* dqkvr, float* dpreatt, float* datt,
                        const float* dout, const float* inp, const float* att,
                        int B, int T, int C, int NH) {
    int HS = C / NH;
    const float alpha = 1.0f, beta = 1.0f;  // beta=1.0 用于梯度累加
    
    // 1. 反向通过 unpermute
    unpermute_kernel_backward<<<num_blocks, block_size>>>(dvaccum, dout, B, T, NH, HS);
    
    // 2. 反向到 datt（使用 cuBLAS）
    cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N,
                              T, T, HS, &alpha, v, HS, T*HS,
                              dvaccum, HS, T*HS, &beta, datt, T, T*T, B*NH);
    
    // 3. 反向到 dv
    cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T,
                              HS, T, T, &alpha, dvaccum, HS, T*HS,
                              att, T, T*T, &beta, dv, HS, T*HS, B*NH);
    
    // 4. Softmax 反向
    softmax_autoregressive_backward(dpreatt, datt, att, B, T, C, NH, block_size);
    
    // 5. 反向到 dq 和 dk
    cublasSgemmStridedBatched(...);  // dq
    cublasSgemmStridedBatched(...);  // dk
    
    // 6. 反向到 dinp
    permute_kernel_backward<<<num_blocks, block_size>>>(dinp, dq, dk, dv, B, T, NH, HS);
}

性能（RTX 4090）：

block_size   64 | time 7067.52 ms  ❌ 极慢

瓶颈：Softmax 反向计算效率低下

V2：并行化优化

策略：在 (t, b, h) 维度上并行

性能：

block_size   64 | time 271.53 ms  ✅ 26x 加速

提升原因：充分利用 GPU 并行能力

V3：协作组优化

策略：在 (t, b, h) 维度并行 + 使用 Cooperative Groups

性能：

block_size  256 | time 14.37 ms  ✅ 18.9x 加速

关键技术：利用协作组进行高效的线程间通信和同步

V4：循环展开

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

性能：

block_size  256 | time 10.38 ms  ✅ 1.4x 加速

优化点：减少循环开销，提高指令级并行

V5：特殊情况优化

策略：针对 V4 的特殊情况进行优化

性能：

block_size   32 | time 6.03 ms  ✅ 1.7x 加速

优化点：处理边界条件，减少分支预测失败

V6：内存访问优化

策略：循环重构 + 内存访问模式优化

性能：

block_size   64 | time 2.99 ms  ✅ 2.0x 加速

优化点：

• 改善内存访问的局部性
• 减少全局内存访问次数
• 优化缓存命中率

V7：数学简化 + 协作组规约

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

性能：

block_size  256 | time 1.76 ms  ✅ 1.7x 加速

关键优化：

• 数学公式简化，减少计算量
• 使用高效的协作组规约（warp-level primitives）
• 减少共享内存使用

V8：综合优化

策略：在 V7 基础上新增多种 tricks

性能：

block_size  512 | time 1.76 ms  ≈ 持平

结论：已接近硬件理论性能上限

2.4 性能对比总结

惊人的优化成果：从 7067ms 优化到 1.76ms，实现了 4015 倍的性能提升！

三、性能对比分析

3.1 前向传播性能对比

根据文章数据，在 RTX 4090 + CUDA 12.4 环境下：

结论：cuDNN 的工业级实现在前向传播中表现出色，且无需手工调优。

3.2 反向传播性能对比

关键洞察：反向传播的优化空间远大于前向传播，主要瓶颈在 Softmax 反向计算。

3.3 优化技术总结

并行化策略

• 粗粒度并行：在 batch、head、sequence 维度并行
• 细粒度并行：warp 级别的协作计算

内存优化

• 访问模式：合并访问、减少全局内存访问
• 缓存利用：提高 L1/L2 缓存命中率
• 共享内存：减少共享内存使用，避免 bank conflict

计算优化

• 循环展开：减少循环开销
• 数学简化：减少浮点运算次数
• 协作组规约：使用 warp shuffle 等高效原语

四、核心要点总结

4.1 FlashAttention 原理（cuDNN 实现基础）

虽然文章未详细展开 FlashAttention-2 原理，但其核心思想是：

1. 分块计算：将大矩阵分成小块，逐块计算
2. 在线 Softmax：避免完整物化 attention 矩阵
3. 重计算策略：反向传播时重新计算而非存储中间结果
4. IO 优化：最小化 HBM（高带宽内存）访问

关键优势：

• 内存占用：O(N) vs O(N²)
• IO 复杂度：显著降低
• 支持长序列：可处理更长的上下文

4.2 分块计算策略

cuDNN 的实现特点：

• 支持非连续内存布局（通过 stride 配置）
• 自动处理 workspace 分配（默认最多 256MB）
• 缓存计算图以加速后续调用

内存布局优化：

QKV 输入格式：(B, T, 3, NH, HS)
Q stride: {3*H*HS*T, HS, 3*H*HS, 1}
K stride: {3*H*HS*T, HS, 3*H*HS, 1}
V stride: {3*H*HS*T, HS, 3*H*HS, 1}

这种布局避免了额外的 permute 操作。

4.3 IO 复杂度分析

传统实现的 IO 复杂度：

• 读取 Q, K, V：3 × B × T × C
• 写入 S (QK^T)：B × NH × T²
• 读取 S，写入 P (Softmax)：2 × B × NH × T²
• 读取 P, V，写入 O：2 × B × T × C + B × NH × T²

总计：O(B × T² × NH) - 对长序列非常不友好

FlashAttention 的 IO 复杂度：

• 通过分块和重计算，降低到 O(B × T × C)
• 对长序列友好

4.4 实现细节要点

cuDNN 实现关键点

1. 混合精度计算

   • 输入/输出：FP16/BF16（节省内存带宽）
   • 中间计算：FP32（保证数值稳定性）

2. 计算图缓存

   • build_operation_graph() 非常耗时
   • 使用 hash map 缓存已构建的图
   • 相同形状的输入可复用计算图

3. 动态 Workspace 管理

   • 根据实际需求动态分配
   • 避免预分配过大内存
   • 默认上限 256MB

4. 训练 vs 推理模式

   • 推理模式：不计算 softmax_stats，节省内存
   • 训练模式：保存统计信息用于反向传播

反向传播优化关键点

1. Softmax 反向是瓶颈

   • V1 → V2：并行化带来 26x 加速
   • V2 → V7：持续优化带来 154x 加速

2. 协作组的威力

   • V3 引入协作组：18.9x 加速
   • V7 使用协作组规约：额外 1.7x 加速

3. 内存访问模式

   • V6 的内存优化：2x 加速
   • 合并访问、提高缓存命中率

4. 数学简化

   • V7 简化 Softmax 反向公式
   • 减少不必要的计算

4.5 性能提升效果

前向传播：

• 手工优化版本：~0.17 ms
• cuDNN 版本：0.166 ms
• 提升：约 2.4%

反向传播：

• 基础版本：7067 ms
• 高度优化版本：1.76 ms
• 提升：4015 倍（99.975% 性能提升）

总体结论：

• 反向传播的优化空间远大于前向传播
• 系统性优化方法论至关重要
• 工业级库（cuDNN）在前向传播中更具优势

五、技术洞察与最佳实践

5.1 CUDA 优化方法论

文章展示了一套系统性的 CUDA 优化方法论：

第一阶段：并行化（V1 → V2）

• 识别可并行的维度
• 充分利用 GPU 的并行能力
• 通常能带来数十倍加速

第二阶段：协作优化（V2 → V3）

• 引入协作组（Cooperative Groups）
• 优化线程间通信和同步
• 利用 warp 级别的原语

第三阶段：细节优化（V3 → V6）

• 循环展开（减少控制流开销）
• 特殊情况处理（减少分支）
• 内存访问模式优化（提高缓存命中率）

第四阶段：算法优化（V6 → V7）

• 数学公式简化
• 高效的规约操作
• 接近硬件理论性能

5.2 何时使用 cuDNN vs 手工优化

使用 cuDNN 的场景：

• ✅ 标准算子（卷积、Attention、归一化等）
• ✅ 需要快速部署
• ✅ 追求稳定性和可维护性
• ✅ 前向推理为主

手工优化的场景：

• ✅ 非标准算子
• ✅ 特定硬件或特定场景
• ✅ 需要极致性能
• ✅ 反向传播优化

5.3 关键技术要点

1. Softmax 反向传播公式

标准形式：

dS = P ⊙ (dP - (dP ⊙ P) · 1)

简化形式（V7 使用）：

利用 Softmax 的性质进行数学简化
减少浮点运算次数

2. 协作组规约

// 使用 warp shuffle 进行高效规约
float warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

3. 内存访问模式

优化前：

// 跨步访问，缓存命中率低
for (int i = 0; i < N; i++) {
    result += data[i * stride];
}

优化后：

// 连续访问，缓存友好
for (int i = 0; i < N; i++) {
    result += data[i];
}

六、代码资源与参考资料

6.1 代码仓库

作者代码：

• GitHub: https://github.com/ifromeast/cuda_learning/blob/main/04_transformer/ops/attention_backward.cu
• 包含 V1-V8 所有版本的实现

参考实现：

1. Andrej Karpathy 的 llm.c 项目

   • 前向：https://github.com/karpathy/llm.c/blob/master/dev/cuda/attention_forward.cu
   • 反向：https://github.com/karpathy/llm.c/blob/master/dev/cuda/attention_backward.cu

2. NVIDIA cuDNN Frontend

   • 文档：https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md
   • 示例：https://github.com/NVIDIA/cudnn-frontend/blob/main/samples/cpp/mha.cpp
   • Python 示例：https://github.com/NVIDIA/cudnn-frontend/blob/main/samples/python/50_scaled_dot_product_attention.ipynb

6.2 理论参考

1. 矩阵求导

   • Stack Exchange: Not understanding derivative of a matrix-matrix product
   • Matrix Cookbook: https://www.math.uwaterloo.ca/~hwolkowi/matrixcookbook.pdf

2. Softmax 反向传播

   • 知乎文章：反向传播之一：softmax函数

3. FlashAttention 论文

   • FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

6.3 相关文章

• 上篇：ops(7)：self-attention 的 CUDA 实现及优化 (上)
  • 链接：https://zhuanlan.zhihu.com/p/695898274
  • 内容：基础实现与前向优化（V1-V5）

七、总结与展望

7.1 核心收获

技术层面：

1. cuDNN 的价值：工业级库在标准算子上具有显著优势，无需重复造轮子
2. 优化的系统性：从并行化 → 协作优化 → 细节优化 → 算法优化的完整路径
3. 反向传播的复杂性：反向传播的优化空间远大于前向传播
4. Softmax 是关键：Softmax 反向计算是主要性能瓶颈

性能数据：

• 前向传播：cuDNN 实现 0.166ms（接近最优）
• 反向传播：从 7067ms 优化到 1.76ms（4015 倍提升）

方法论：

• 先并行化（粗粒度）
• 再协作优化（细粒度）
• 然后内存优化（访问模式）
• 最后算法优化（数学简化）

7.2 实践建议

对于工程实践：

1. 优先使用成熟库（cuDNN、cuBLAS）
2. 性能瓶颈明确后再考虑手工优化
3. 使用 profiler 工具识别热点
4. 渐进式优化，每步验证正确性

对于学习研究：

1. 理解数学原理是基础
2. 掌握 CUDA 编程模型
3. 学习协作组等高级特性
4. 研究开源实现（llm.c、cuDNN）

7.3 未来方向

硬件层面：

• 新一代 GPU 架构（Hopper、Blackwell）
• 更大的共享内存和 L2 缓存
• 更高的内存带宽

算法层面：

• FlashAttention-3 及后续版本
• 稀疏 Attention 优化
• 长序列 Attention 优化

工程层面：

• 自动调优框架
• 编译器优化
• 算子融合

附录：关键公式速查

A.1 前向传播

S = (Q · K^T) / √d_k
P = Softmax(S)
O = P · V

A.2 反向传播

dV = P^T · dO
dP = dO · V^T
dS = P ⊙ (dP - (dP ⊙ P) · 1)
dQ = (dS · K) / √d_k
dK = (dS^T · Q) / √d_k

A.3 Softmax 导数

∂y_i/∂x_j = {
  y_i(1 - y_j),  if i = j
  -y_i · y_j,     if i ≠ j
}

文章引用：

休言万事转头空，未转头时皆梦。——苏轼《西江月·平山堂》

分析完成时间：2026-05-05
原文链接：https://zhuanlan.zhihu.com/p/696197013
代码仓库：https://github.com/ifromeast/cuda_learning