【AI Infra 核心】图解 FlashAttention：长上下文推理背后的“内存墙”革命与底层优化思维

【AI Infra 核心】图解 FlashAttention长上下文推理背后的“内存墙”革命与底层优化思维摘要现在的大模型上下文窗口从 8K 卷到 128K甚至百万级别如 Gemini 1.5 Pro。但如果你用过早期的 Transformer会知道当序列长度翻倍时显存占用会呈**平方级Quadratic**爆炸这一切的救星就是大模型基建领域的“神作”——FlashAttention。今天我们不堆砌晦涩的论文公式而是从 GPU 底层硬件架构出发结合 Python 代码模拟带你彻底搞懂 FlashAttention 是如何通过“分块计算”打破内存墙的一、 标准 Attention 的原罪O(N2)O(N^2)O(N2)的“内存墙”困境在原始的 Transformer 中Self-Attention 的核心公式大家倒背如流Attention(Q,K,V)softmax(QKTd)V \text{Attention}(Q, K, V) \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)VAttention(Q,K,V)softmax(d​QKT​)V假设序列长度Context Length为NNNHead 的维度为ddd。我们来看看标准深度学习框架如早期的 PyTorch在底层是怎么执行这行公式的计算注意力分数SSQ×KTS Q \times K^TSQ×KT。这里会生成一个大小为N×NN \times NN×N的庞大矩阵并写入显存。计算 SoftmaxPPsoftmax(S)P \text{softmax}(S)Psoftmax(S)。从显存中读取SSS进行指数计算和归一化再将N×NN \times NN×N的PPP矩阵写入显存。计算输出OOP×VO P \times VOP×V。再次从显存读取PPP和VVV相乘后得到最终结果。灾难发生了当你的NNN达到 32K 时这个中间矩阵N×NN \times NN×N会变得极其巨大。更可怕的不是“算不过来”而是**“搬运不过来”**二、 认清现实GPU 计算极快但“读写”极慢要理解 FlashAttention必须先懂 GPU 的内存层级架构SRAM静态随机存储器在 GPU 计算核心旁边速度极快19TB/s但容量极小比如 A100 每块 SM 上只有 192KB。HBM高带宽内存/显存就是我们常说的 40G/80G 显存。容量大但速度相对较慢1.5TB/s - 2TB/s。标准 Attention 的致命问题在于它在 HBM 和 SRAM 之间频繁地来回搬运那个巨大的N×NN \times NN×N矩阵GPU 核心的算力FLOPS被白白闲置全在等 HBM 慢吞吐的内存读写。在学术上这叫Memory-Bound访存瓶颈。三、 FlashAttention 的破局Tiling分块计算与 SRAM 常驻既然 HBM 读写太慢而 SRAM 太小装不下整个矩阵那该怎么办FlashAttention 祭出了系统优化里最经典的招式分块Tiling。它的核心理念极其狂野我不存中间结果了我把Q,K,VQ, K, VQ,K,V切成小块Block每次只把一小块加载到超快的 SRAM 中一口气算完Q×KTQ \times K^TQ×KT、Softmax 和×V\times V×V直接输出最终结果写回 HBM完全绕过对中间矩阵SSS和PPP的保存。这就好比以前你是把整车砖头卸在村口HBM再用小推车一车车拉到工地SRAM盖房子现在你直接让砖车开进工地现拿现砌绝不二次搬运。四、 核心难点局部数据怎么算全局 Softmax“分块”听起来简单但遇到了一个巨大的数学阻碍Softmax。标准的 Softmax 公式要求你必须拥有完整的一整行数据才能求出分母所有元素的指数和softmax(xi)exi∑j1Nexj \text{softmax}(x_i) \frac{e^{x_i}}{\sum_{j1}^{N} e^{x_j}}softmax(xi​)∑j1N​exj​exi​​如果你只把矩阵切了一小块放进 SRAM你根本不知道这一行的其他元素是多少怎么求和怎么做归一化 绝妙的数学 TrickOnline SoftmaxFlashAttention 使用了Safe Online Softmax算法。它维护了两个变量局部最大值Max和局部指数和Sum。当新的块Block进来时它会通过极其精妙的数学换算动态更新全局的最大值和和值并对之前已经算出来的结果进行“缩放补偿”。我们用一段 Python 代码徒手模拟这个只占用O(1)O(1)O(1)额外空间且支持分块计算的 Online Attention五、 硬核实战Python 模拟 FlashAttention 底层分块逻辑代码说明这段代码舍弃了标准 PyTorch 的矩阵直乘而是用for循环模拟 GPU 中分块加载Tiling和状态更新Online Softmax的过程。importtorchdefflash_attention_simulate(Q,K,V,block_size2): 模拟 FlashAttention 的分块计算与 Online Softmax (前向传播) Q, K, V shape: (seq_len, head_dim) - 这里简化为单头 seq_len,head_dimQ.shape# 最终的输出 O (存放在 HBM)Otorch.zeros((seq_len,head_dim))# 将 Q 切分成块 (模拟外层循环加载到 SRAM)foriinrange(0,seq_len,block_size):q_blockQ[i:iblock_size,:]# 维护这一块的全局最大值和全局指数和 (存放在超快的 SRAM)m_itorch.full((q_block.shape[0],1),float(-inf))# 局部最大值l_itorch.zeros((q_block.shape[0],1))# 局部指数和# 累加结果o_itorch.zeros((q_block.shape[0],head_dim))# 内层循环遍历 K 和 V 的所有块forjinrange(0,seq_len,block_size):k_blockK[j:jblock_size,:]v_blockV[j:jblock_size,:]# 1. 计算局部 Score: S_{ij} Q_i * K_j^Ts_ijtorch.matmul(q_block,k_block.transpose(0,1))/(head_dim**0.5)# 2. 更新最大值 m_new max(m_old, row_max(s_ij))m_ij_maxtorch.max(s_ij,dim1,keepdimTrue).values m_newtorch.maximum(m_i,m_ij_max)# 3. 计算指数并按新的最大值进行数值稳定 (Safe Softmax)# P_{ij} exp(S_{ij} - m_new)p_ijtorch.exp(s_ij-m_new)# 4. 更新局部的和 l_new l_old * exp(m_old - m_new) sum(P_{ij})# 核心 trick: 之前的和必须缩小 exp(m_old - m_new) 倍才能对齐当前的最大值l_newl_i*torch.exp(m_i-m_new)torch.sum(p_ij,dim1,keepdimTrue)# 5. 更新输出 O_i (O_i * l_i * exp(m_old - m_new) P_{ij} * V_j) / l_new# 这里先不急着除以 l_new先维护未归一化的输出分子o_io_i*torch.exp(m_i-m_new)torch.matmul(p_ij,v_block)# 更新状态为下一次迭代准备m_im_new l_il_new# 6. 所有 K, V 遍历完后对当前 Q_block 进行最终的归一化O[i:iblock_size,:]o_i/l_ireturnO# 测试验证 torch.manual_seed(42)N,d8,4# 序列长度 8维度 4Qtorch.randn(N,d)Ktorch.randn(N,d)Vtorch.randn(N,d)# 1. 标准 Attention (占用 O(N^2) 内存)scorestorch.matmul(Q,K.transpose(0,1))/(d**0.5)attn_weightstorch.softmax(scores,dim-1)out_standardtorch.matmul(attn_weights,V)# 2. FlashAttention 模拟 (分块大小设为 2)out_flashflash_attention_simulate(Q,K,V,block_size2)# 对比结果差距difftorch.max(torch.abs(out_standard-out_flash))print(f标准 Attention 与 FlashAttention 模拟的输出差异:{diff.item():.6e})# 输出极小证明逻辑完全等价六、 总结从算法突破到系统工程的胜利这段模拟代码完美揭示了 FlashAttention 的精髓它不仅没有减少计算量甚至增加了一点点指数计算但它通过消除海量的显存读写延迟让整体运算速度飙升了 2 到 4 倍并且将显存占用从O(N2)O(N^2)O(N2)强行降维到了线性O(N)O(N)O(N)。可以说如果没有 FlashAttention 的底层重构目前大行其道的 128K 超长上下文模型根本不可能被普通企业部署。在 AI Infra 领域算法的数学优雅与底层硬件的脾气同样重要。懂 CUDA 层级的 IO 优化才是当下 AI 系统工程师的核心竞争力。