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 TFLOPSFP16。当计算密度低于 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_i0, l_i0, 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在线 SoftmaxOnline 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 容量。典型设置为 Br128, Bc128。三、简化的 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_partition32(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 决定了典型 Br128, Bc128。移植到非 NVIDIA 平台的核心工作量为 warp-level 原语的平台适配。直接影响 30%50% 的移植周期。长序列16384场景需要在 FP16 的 critical 路径上使用 FP32 中间缓冲来防止 exp 溢出。