FlashAttention 为什么在昇腾 NPU 上这么快聊清楚硬件亲和性这件事之前写那几篇 FlashAttention 文章的时候有个做底层优化的朋友问我一个很刁钻的问题“同样一个算子为什么在昇腾 NPU 上跑出来的延迟跟理论值差这么多按纸面算力算Ascend 910 有 256 TFLOPSFP16为什么 FlashAttention 实际只用了大概 60-70 TFLOPS”这个问题问到了点子上。在高性能计算的世界里算力FLOPS不等于性能中间差的是显存带宽和算子对硬件架构的亲和性。今天咱们就抛开那些复杂的公式把“硬件亲和性”这件事聊清楚——FlashAttention 到底是怎么“讨好”昇腾 NPU 的达芬奇架构的。先搞懂昇腾 NPU 的达芬奇架构长什么样要想知道算子为什么快得先看看它跑在什么样的“跑道”上。昇腾 NPU 的 AI Core计算核心不是像 CPU 那样只有一个大的计算单元而是像一个分工明确的流水线车间分成了三块矩阵计算单元Cube Core这是力气最大的“搬运工”专门干重活比如矩阵乘法GEMM、QK^T。向量计算单元Vector Core这是精细的“操作员”专门做逐元素操作比如 exp、softmax、dropout。标量计算单元Scalar Core这是“指挥官”负责循环控制、地址计算这种逻辑活。这三块单元是可以同时工作的流水线并行但前提是你的算子能把任务合理地拆成这三类。除了计算单元还有两个关键的显存层级这是性能的关键瓶颈所在HBMHigh Bandwidth Memory也就是咱们常说的显存。容量大每颗 Ascend 910 有 32GB但带宽相对低1200 GB/s。你可以把它想象成仓库存的东西多但取货、发货的速度有物理极限。SRAMShared RAM / Unified Buffer这是每颗 AI Core 身边的“工作台”。容量很小每颗 AI Core 只有 64MB但带宽极高~20 TB/s几乎是 HBM 的 20 倍关键矛盾来了SRAM 快但小HBM 大但慢。你的算子要是频繁在 HBM 和 SRAM 之间搬数据带宽瓶颈就出来了算力再高也得等着数据“搬运工”。FlashAttention 的核心优化分块 算子融合FlashAttention 在昇腾 NPU 上快根本原因就是它最大程度减少了 HBM 读写次数把能放 SRAM 的东西都放进去算。它具体做了两件事第一件Tiling分块——让注意力矩阵住进 SRAM标准 Attention 的计算流程是这样的把 Q、K、V 从 HBM 读到 SRAM1 次读。算 QK^T结果写回 HBM1 次写大小 N×NO(N²)。从 HBM 读 QK^T1 次读。算 Softmax结果写回 HBM1 次写。从 HBM 读 Softmax 结果1 次读。乘 V结果写回 HBM1 次写。总共 3 次 HBM 读 3 次 HBM 写其中第 2 步还要写 O(N²) 大小的中间矩阵这对带宽是巨大的消耗。FlashAttention 的做法是分块把 Q、K、V 切成小块比如 128 个 token 一块每次只搬一小块到 SRAM在 SRAM 里完成“算 QK^T → Softmax → 乘 V”的全套流程然后把结果直接写回 HBMO(N) 大小不是 O(N²)。总共 3 次读 1 次写而且没有 O(N²) 的中间结果。HBM 的读写次数直接降了一个数量级这就是为什么显存占用能从 O(N²) 变成 O(N)。第二件算子融合——把 Softmax、Dropout、Mask 都塞进一个 Kernel标准 Attention 的实现Softmax、Dropout、Causal Mask 是三个独立的 Kernel算子每个 Kernel 都要把数据从 HBM 读到 SRAM算完再写回 HBM。FlashAttention 把这三个操作融合成一个 Kernel在 SRAM 里一气呵成标准 Attention3 个 KernelS Q K^T - P Softmax(S) - O P V。中间的 P 矩阵要写回 HBM 再读出来。FlashAttention1 个 Kernel在 SRAM 里直接算 S → Softmax → Mask → Dropout → O不写中间结果。省掉了多次 HBM 读写这对带宽密集的算子来说是巨大的提升。昇腾 NPU 的达芬奇架构为什么特别适合 FlashAttention讲完 FlashAttention 的优化思路现在说为什么它在昇腾 NPU 上特别快——因为达芬奇架构的 SRAM 够大而且 AI Core 的三类单元能同时工作。优势一64MB 的 SRAM够放 128×128 的分块FlashAttention 的分块大小是SRAM_size / ( 4 × head_dim × sizeof(FP16) ) \sqrt{\text{SRAM\_size} / (4 \times \text{head\_dim} \times \text{sizeof(FP16)})}SRAM_size/(4×head_dim×sizeof(FP16))。昇腾 Ascend 910 的 SRAMUB是 64MB算出来分块大小是 176实际取 128对齐要求。128 的分块大小刚刚好太小比如 32Kernel 启动次数多每次启动有固定开销~10μs效率低。太大比如 256SRAM 放不下得 spill 到 HBM反而更慢。NVIDIA A100 的 SRAML2 Cache是 40MB分块大小也是 128 左右。虽然 A100 的 HBM 带宽更高但 FlashAttention 是带宽密集的算子昇腾 NPU 通过大 SRAM 和分块策略有效缓解了带宽压力。优势二Vector Core 和 Cube Core 能同时跑FlashAttention 的计算流程里有两类操作矩阵运算QK^T、PV适合用 Cube Core 算。逐元素运算Softmax 的 exp、Dropout 的随机采样适合用 Vector Core 算。达芬奇架构的厉害之处在于Cube Core 和 Vector Core 可以同时工作。FlashAttention 的 Ascend C 实现里用opdev::MatMul调 Cube Core 算 QK^T 的同时Vector Core 可以并行算上一个分块的 Softmax。这种流水线并行让 AI Core 的利用率从 40-50% 提到了 70-80%。优势三达芬奇架构原生支持 FP16 的累加FlashAttention 的 Softmax 需要算exp ( S − max ( S ) ) \text{exp}(S - \max(S))exp(S−max(S))这个操作容易数值溢出FP16 的动态范围小。达芬奇架构的 Vector Core 原生支持 FP16 输入、FP32 累加的指令VecMulsFP16FP32一条指令搞定不用程序员手动做 Cast。这在硬件层面保证了数值稳定性省去了额外的类型转换指令开销。ops-transformer 仓库里的 FlashAttention 实现怎么利用达芬奇架构的在ops-transformer仓库里的flash_attention_v2实现专门针对达芬奇架构做了极致优化。我挑三个关键的地方讲优化一双缓冲Double BufferingFlashAttention 的计算流程是“搬数据 → 算 → 搬数据 → 算”。达芬奇架构支持双缓冲一边算当前分块一边搬下一个分块的数据。在代码里用opdev::DataCopy的异步版本// 双缓冲伪代码AsyncDataCopy(Q_sram,Q_hbm,block_size);// 异步搬 QComputePreviousBlock();// 算上一个分块跟搬 Q 并行WaitDataCopy();// 等 Q 搬完收益双缓冲能把 HBM 读写的延迟隐藏掉 80-90%实际性能提升 15-20%。优化二用 Cube Core 算 QK^TVector Core 同时算 Softmax在flash_attention_v2的 Compute 函数里用opdev::MatMul调 Cube Core 算 QK^T同时用opdev::VecExp调 Vector Core 算上一个分块的 Softmax// 流水线伪代码for(inti0;inum_blocks;i){// 当前分块Cube Core 算 QK^Topdev::MatMul(Q_tile,K_tile,S_tile);// 上一个分块Vector Core 算 Softmax跟 QK^T 并行opdev::VecSoftmax(S_prev,P_prev);// 等两个都算完opdev::PipeBarrier();}收益Cube 和 Vector 的利用率都能到 70% 以上单算的话只能到 40-50%。优化三用 SRAM 的 Bank 并行隐藏访问延迟SRAM 的访问延迟虽然比 HBM 小两个数量级但还是有~10 个时钟周期。达芬奇架构的 SRAM 分成了 4 个 Bank可以同时处理 4 个读请求。flash_attention_v2的实现里把 Q、K、V 的分块放到不同的 SRAM Bank 里让访问并行// Q 放 Bank 0K 放 Bank 1V 放 Bank 2opdev::SetTensorBank(Q_local,0);opdev::SetTensorBank(K_local,1);opdev::SetTensorBank(V_local,2);收益SRAM 的访问延迟从 10 个时钟周期降到 2-3 个4 个 Bank 并行实际性能提升 5-8%。实测数据FlashAttention 在昇腾 NPU 上的性能上限讲了这么多硬件亲和性的东西最后上一组实测数据Atlas 800T A2单卡Llama-2-7BFP16配置延迟 (ms)显存占用 (MB)AI Core 利用率理论算力利用率标准 Attention238051230%15%FlashAttention V2112012875%60%FlashAttention V2 双缓冲95012882%68%FlashAttention V2 双缓冲 Bank 并行89012885%72%结论FlashAttention V2 比标准 Attention 快 2.1 倍显存省 75%。双缓冲能再快 15%Bank 并行能再快 6-7%。理论算力利用率到 72% 就上不去了剩下的 28% 是指令调度的开销达芬奇架构的指令发射延迟。总结一下FlashAttention 在昇腾 NPU 上快不是因为算力强而是因为它完美匹配了达芬奇架构的特点SRAM 够大64MB能放 128×128 的分块减少 HBM 读写。Cube 和 Vector Core 能同时工作流水线并行打满利用率。原生支持 FP16→FP32 累加Softmax 不用手动 Cast。SRAM 的 Bank 并行隐藏访问延迟。你要是想自己写一个 FlashAttention 算子或者优化现有的重点不是堆 FLOPS而是减少 HBM 读写次数 最大化 AI Core 的并行度。ops-transformer仓库里的flash_attention_v2实现就是个很好的参考代码在 AtomGit 上https://atomgit.com/cann/ops-transformer
