昇腾CANN ops-transformer 仓的 FlashAttention 算子:昇腾NPU上的注意力加速实现大模型推理和训练里,Self-Attention 层的计算是最大的性能瓶颈。FlashAttention 把这块的计算从 O(n²) 的显存占用降到了 O(n),靠的是分块计算——把整个注意力矩阵拆成小块,逐块在片上缓存里算完再写回 HBM。ops-transformer 仓是昇腾CANN 的 Transformer 类进阶算子库,里面就有一个昇腾NPU 原生的 FlashAttention 实现。这篇文章拆开看它怎么在昇腾达芬奇架构上做分块计算和在线 softmax,以及实际的性能表现。