FlashAttention 在昇腾NPU上的实现：从内存墙到IO感知

之前帮一个团队排查大模型训练显存溢出的问题，拿到profiling数据一看，Attention 层的 HBM 访存占了整条流水线 60% 以上的带宽。这不是算力不够------是数据搬得太频繁了。

FlashAttention 的核心思路并不复杂：别把整个注意力矩阵写到 HBM 里再读回来。 标准的 Attention 实现会先算一个完整的 QK^T 矩阵存到显存，再算 Softmax，再乘 V------每一步都涉及 HBM 的读写。FlashAttention 做的事情是把这套流程拆成小块（tile），在片上高速存储（SRAM/UB）里完成全部计算，中间结果只在必要时才落回 HBM。

ops-transformer 里的 FlashAttention：不是简单移植

昇腾CANN 的 ops-transformer 仓库是 Transformer 类大模型进阶算子的集中地，FlashAttention 正是其中最核心的算子之一。但昇腾NPU 的硬件结构和 GPU 有本质区别------没有传统的 shared memory，取而代之的是统一缓冲区（Unified Buffer, UB），容量和访问模式都不同。

这意味着不能直接把 GPU 上的 FlashAttention 算法搬过来跑。 ops-transformer 中的 FlashAttention 实现需要专门针对 UB 的容量做 tile 分块策略的调优，针对达芬奇架构的 Cube 和 Vector 计算单元做算子调度。Ascend C 编程语言提供了对 Cube（矩阵乘）和 Vector（向量计算）指令的直接调用能力，FlashAttention 的 QK^T 矩阵乘和 Softmax 计算分别对应这两个计算单元。

IO-aware：算力不是瓶颈，搬运才是

理解 FlashAttention 的关键不是数学公式，是一个朴素的事实：矩阵乘的算力密度远高于逐元素操作。

一次 QK^T 的矩阵乘，每个元素参与一次乘加；而 Softmax 要对每一行做指数运算、求和、再除------计算量不比矩阵乘小，但访存模式完全不同。传统实现里，矩阵乘的结果先写回 HBM，Softmax 再从 HBM 读出来算。这就好比炒菜时把每道菜都端到客厅再端回厨房------来回跑的路比炒菜本身还费时间。

FlashAttention 的做法是在 UB 里完成一个 tile 的 QK^T → Softmax → 乘 V 的完整流程，中间结果始终留在片上。只有最终的结果需要写回 HBM。ops-transformer 的实现进一步利用了昇腾NPU 的异步搬运能力，在当前 tile 计算的同时预加载下一个 tile 的数据，把数据搬运和计算重叠起来。

在 CANN 五层架构中的位置

ops-transformer 属于 CANN 第二层（计算服务层）的 AOL 算子库，是昇腾异构计算架构中面向大模型的关键算子集。它的上游依赖 opbase（算子基础组件库）和 ops-nn（基础神经网络算子），下游被 ascend-transformer-boost（ATB）Transformer 加速库直接调用。

当你用 PyTorch 跑一个基于 LLaMA 架构的模型时，经过 Framework Adaptor 的图转换和 AOE 调优引擎的自动优化后，Attention 层最终会落到 ops-transformer 中的 FlashAttention 算子上执行。整个过程用户无感知，但性能差距可能达到数倍。

causal mask 和 dropout 的融合处理

实际大模型训练中，FlashAttention 不可能只做"纯粹的"注意力计算。因果语言模型的 causal mask、训练时的 dropout、以及 head 维度的并行，都需要在算子层面融合处理。

ops-transformer 的 FlashAttention 实现把 causal mask 融合进了 QK^T 之后的 softmax 归一化步骤：被 mask 掉的位置在 softmax 之前被设为负无穷大，softmax 之后自然趋近于零。这种方式避免了单独做 mask 矩阵乘法带来的额外访存开销。dropout 同样在 softmax 之后、乘 V 之前在 UB 内完成，不增加任何 HBM 读写。

融合算子的价值就在这里------省一次搬运，就省一次带宽，就多出一点算力给真正的计算。

写给想动手试试的人

如果你在昇腾NPU 上做大模型训练，想确认 FlashAttention 是否真正生效，可以用 CANN 的 Profiler 工具抓一次 trace，看 Attention 算子的 HBM 访存量。如果看到大量的中间矩阵读写，说明还没走到融合算子路径------检查一下框架适配层和 ATB 的配置。

ops-transformer 仓库的代码和文档已经全面开源，可以直接去看 FlashAttention 的 Ascend C 实现细节：

https://atomgit.com/cann/ops-transformer

仓库里还有 MoE、MC2 等其他大模型算子的实现，都值得翻一翻。如果你对某个具体算子的实现有疑问，社区 Discussions 区是提问的好地方。