GPU Kernel 级别的 Attention 融合实现：FlashAttention 在自定义推理引擎中的移植记录
GPU Kernel 级别的 Attention 融合实现：FlashAttention 在自定义推理引擎中的移植记录
一、标准 Attention 在长序列下的显存墙问题
在部署一个长文本摘要模型时。输入序列长度为 8192 tokens。标准 Scaled Dot-Product Attention 的显存占用达到了 12GB。仅用于存储中间注意力矩阵 S = QK^T。而这个矩阵的大小是 8192 × 8192 × 4 字节（FP32）= 256MB。乘以多头注意力（假设 40 头）就是 10GB 以上。
更严重的是，这个矩阵还需要经历 softmax 和与 V 的二次矩阵乘法。每一步都涉及从 HBM 读取、计算、再写回 HBM。在 A100 上，HBM 带宽约 1.5TB/s。而计算吞吐是 312 TFLOPS（FP16）。当计算密度低于 HBM 带宽能支撑的水平时。GPU 处于内存带宽瓶颈状态。大量的计算单元处于空闲。
FlashAttention 的核心洞察是：注意力矩阵不需要显式地在 HBM 中展开。可以通过分块（Tiling）和在线 Softmax 在 SRAM 中完成全部计算。只将最终结果写回 HBM。这同时解决了显存占用和内存带宽瓶颈两个问题。
二、FlashAttention 的分块计算与在线 Softmax 原理
FlashAttention 将 Q、K、V 矩阵分别切分为小块。每次加载一个 Q 块和所有 K、V 块到 SRAM 中。在 SRAM 中计算局部注意力并累积。处理完所有块后将最终结果写回 HBM。
sequenceDiagram participant HBM as HBM显存 participant SRAM as SRAM片上缓存 participant SM as SM计算单元 Note over HBM,SM: FlashAttention 分块计算流程 loop 遍历 Q 的所有分块 HBM->>SRAM: 加载 Q_i 块 (Br × d) Note over SRAM: 初始化累积器 O_i=0, l_i=0, m_i=-inf loop 遍历 K,V 的所有分块 HBM->>SRAM: 加载 K_j 块 (Bc × d) HBM->>SRAM: 加载 V_j 块 (Bc × d) SRAM->>SM: S_ij = Q_i × K_j^T SM->>SRAM: 计算局部 softmax Note over SM: m_new = max(m_old, rowmax(S_ij)) Note over SM: l_new = exp(m_old-m_new)*l_old Note over SM: + sum(exp(S_ij-m_new)) Note over SM: O_i = diag(exp(m_old-m_new)) * O_i Note over SM: + exp(S_ij-m_new) × V_j end SRAM->>HBM: 写入最终 O_i 结果 end在线 Softmax（Online Softmax）是实现分块计算的关键。标准 Softmax 需要先遍历所有元素求最大值和总和。然后再次遍历做归一化。这对于分块计算不可行，因为后续块在计算当前块时还不存在。
在线 Softmax 用增量更新的方式解决这个问题。维护三个累积量：当前最大值 m、归一化分母 l、累积注意力输出 O。每处理一个新分块时，如果发现更大的最大值。就用缩放因子调整之前的 m、l 和 O。数学上可以证明这等价于全局 Softmax。
SRAM 的大小是分块策略的根本约束。A100 的 SRAM 约为 192KB/SM。分块大小 Br 和 Bc 需要保证 Q_i、K_j、V_j、S_ij、O_i 的总和不超过 SRAM 容量。典型设置为 Br=128, Bc=128。
三、简化的 CUDA Kernel 实现
以下代码展示了 FlashAttention 前向传播的 CUDA Kernel 核心逻辑。
#include <cuda_runtime.h> #include <cuda_fp16.h> #include <cooperative_groups.h> namespace cg = cooperative_groups; // FlashAttention 前向传播的简化 CUDA Kernel // // 约束条件： // - Q, K, V 的维度为 (seq_len, head_dim)，head_dim 通常为 64/128 // - 使用 FP16 进行计算和存储 // - Br 和 Bc 为分块大小，取决于 SRAM 容量 // - 此 Kernel 将每个 Q 块分配给一个 thread block 处理 __global__ void flash_attention_fwd_kernel( const half* __restrict__ Q, // (seq_len, head_dim) const half* __restrict__ K, // (seq_len, head_dim) const half* __restrict__ V, // (seq_len, head_dim) half* __restrict__ O, // (seq_len, head_dim) 输出 const int seq_len, const int head_dim, const float softmax_scale // 1/sqrt(head_dim) ) { // 使用 cooperative_groups 获取同步范围 // thread_block_tile 支持 warp 级别的 shuffle 操作 auto block = cg::this_thread_block(); auto warp = cg::tiled_partition<32>(block); // 每个 thread block 处理一个 Q 块 const int q_block_idx = blockIdx.x; const int q_start = q_block_idx * Br; // 寄存器中的累积状态 // 使用 __shared__ 仅在 block 内共享，HBM 仅用于加载和最终存储 extern __shared__ half shared_mem[]; half* Q_shared = shared_mem; half* K_shared = Q_shared + Br * head_dim; half* V_shared = K_shared + Bc * head_dim; // 将 Q 块加载到 shared memory for (int i = threadIdx.x; i < Br * head_dim; i += blockDim.x) { int q_idx = (q_start + i / head_dim) * head_dim + (i % head_dim); Q_shared[i] = (q_idx < seq_len * head_dim) ? Q[q_idx] : __float2half(0.0f); } block.sync(); // 为每个线程初始化累积器 // m: 行最大值, l: softmax 分母累积, O_acc: 加权 V 累积 const int row = threadIdx.x / head_dim; float m = -INFINITY; float l = 0.0f; float O_acc[head_dim] = {0.0f}; // 遍历 K, V 的所有分块 for (int j = 0; j < seq_len; j += Bc) { // 加载 K_j, V_j 到 shared memory for (int i = threadIdx.x; i < Bc * head_dim; i += blockDim.x) { int kv_idx = (j + i / head_dim) * head_dim + (i % head_dim); K_shared[i] = (kv_idx < seq_len * head_dim) ? K[kv_idx] : __float2half(0.0f); V_shared[i] = (kv_idx < seq_len * head_dim) ? V[kv_idx] : __float2half(0.0f); } block.sync(); // 计算 S = Q_i × K_j^T // 此处为简化示意，实际使用 warp 级矩阵乘法 float S_ij[Bc]; #pragma unroll for (int k = 0; k < Bc; ++k) { float dot = 0.0f; #pragma unroll for (int d = 0; d < head_dim; ++d) { dot += __half2float(Q_shared[row * head_dim + d]) * __half2float(K_shared[k * head_dim + d]); } S_ij[k] = dot * softmax_scale; } // 在线 Softmax 更新 float m_new = m; #pragma unroll for (int k = 0; k < Bc; ++k) { m_new = fmaxf(m_new, S_ij[k]); } // 计算新的归一化分母 float l_new = expf(m - m_new) * l; #pragma unroll for (int k = 0; k < Bc; ++k) { l_new += expf(S_ij[k] - m_new); } // 更新累积输出 float scale = expf(m - m_new); #pragma unroll for (int d = 0; d < head_dim; ++d) { O_acc[d] *= scale; } #pragma unroll for (int k = 0; k < Bc; ++k) { float weight = expf(S_ij[k] - m_new); #pragma unroll for (int d = 0; d < head_dim; ++d) { O_acc[d] += weight * __half2float(V_shared[k * head_dim + d]); } } m = m_new; l = l_new; block.sync(); } // 最终归一化并写回 HBM #pragma unroll for (int d = 0; d < head_dim; ++d) { int out_idx = (q_start + row) * head_dim + d; O[out_idx] = __float2half(O_acc[d] / l); } }实现中的核心性能优化在于 Shared Memory 的使用。Q、K、V 的分块都缓存在 SRAM 中。避免了对 HBM 的重复访问。在线 Softmax 使用三个累积变量 m、l、O_acc。它们的更新公式确保了等价于全局 Softmax 的结果。这是数学正确性和分块计算可行性的关键保证。
四、FlashAttention 移植的生产级考量与限制
将 FlashAttention 移植到自定义推理引擎涉及明确的工程挑战。
首先是指令集适配。原版 FlashAttention 依赖于 CUDA 的 warp-level 原语（如__shfl_xor_sync）。如果目标平台是非 NVIDIA GPU（如 Apple Metal、AMD ROCm）。需要重写这些原语。这属于跨平台移植的最大工作量。
其次是头维度（head_dim）的约束。FlashAttention 对 head_dim 有预设要求。通常支持 64 和 128。对于非标准的 head_dim（如某些小模型的 32 或大模型的 256）。需要修改 Shared Memory 布局。这会导致分块大小 Br、Bc 的改变。
第三是精度验证。在线 Softmax 使用 FP16 计算。对于 seq_len 超过 16384 的极长序列。softmax 中的 exp 操作可能溢出 FP16 的表示范围。需要在加载时自动提升为 FP32 计算。
第四是反向传播的加倍复杂度。如果推理引擎需要支持训练。FlashAttention 的反向传播实现比前向复杂 3 倍以上。需要重新计算中间量（Recomputation）来节省显存。对于纯推理场景则无需关注。
五、总结
- FlashAttention 通过分块计算将 Attention 的 HBM 读写从 O(N^2) 降低到 O(N)。彻底解决了长序列 Attention 的显存墙问题。
- 在线 Softmax 是分块计算正确性的关键。三个累积变量 m、l、O_acc 的增量更新公式保证了数学等价性。
- Shared Memory 容量是分块大小的根本约束。A100 的 192KB SRAM 决定了典型 Br=128, Bc=128。
- 移植到非 NVIDIA 平台的核心工作量为 warp-level 原语的平台适配。直接影响 30%～50% 的移植周期。
- 长序列（>16384）场景需要在 FP16 的 critical 路径上使用 FP32 中间缓冲来防止 exp 溢出。