之前帮一个朋友排查大模型推理性能问题，他抱怨说模型明明不大，推理却慢得离谱。一看 profile，好家伙，90% 的时间都耗在 Attention 计算上了。后来换了昇腾 NPU 跑 CANN ops-transformer 里的 FlashAttention，同样的模型，推理延迟直接降了 60%。

这事让我意识到，大多数人对 Attention 的性能瓶颈其实没什么概念。

一、为什么 Transformer 推理越来越慢

先别看公式。你只需要知道一件事：

Transformer 的 Attention 机制本质上是一个「全员互问」的过程——序列里每个 token 都要跟所有其他 token 算一遍相似度。

一个 50 人的会议室里，每个人要跟其他 49 个人分别交流一轮。这就是 O(n²) 的沟通成本。人少还好，500 人呢？5000 人呢？

大模型推理一模一样。上下文越长，Attention 计算量呈平方级增长。你问 ChatGPT 一个需要翻 128K 上下文的问题，它背后的 Attention 矩阵大小是 128K × 128K——160 亿个元素。

更致命的问题不在计算量，而在搬运量。

二、Attention 的计算瓶颈到底在哪

很多人以为 Attention 慢是 GPU/NPU 算不过来。其实算力不是瓶颈，显存带宽才是。

标准 Attention 流程长这样：

code复制

① 从显存读出 Q 和 K → 算出注意力分数矩阵 S（大小 n×n）
② 对 S 做 softmax
③ 用 softmax 结果加权 V → 输出

问题出在①和②之间：那个 n×n 的中间矩阵 S，必须写回显存再读出来。

你炒菜，锅够热（算力强），但冰箱在隔壁房间（显存远）。每次拿一个食材都要跑一趟——菜能炒得快吗？大部分时间花在路上了，灶台的火干烧着。

HBM 显存虽然快，跟芯片内部的算力比还是慢了一个数量级。朴素 Attention 的数据搬运量是 O(n²) 级别的，n 一大，计算单元就得干等。

三、FlashAttention 为什么出现

FlashAttention 的思路很朴素：别把中间结果写回显存。

它把 Attention 拆成小块（tiling），一块一块算。算完一块的 softmax 立刻用掉，不存那个 n×n 的中间矩阵。

里面有个巧妙的数学技巧叫 online softmax——不用等全部数据到齐再算 softmax，一边读一边修正，最终结果跟标准 softmax 完全等价。

效果：

• 显存读写量从 O(n²) 降到 O(n)
• 那个巨大的中间矩阵不用存了，显存压力骤减
• 长序列场景下速度提升 3-5 倍

最快的优化从来不是算得更快，而是搬得更少。

四、昇腾 NPU 如何优化 Attention

昇腾 NPU 的达芬奇架构有几个跟 FlashAttention 天然契合的设计：

🔹 大带宽片上缓存

昇腾 NPU 有 L1/L2 缓存和 Unified Buffer，片上带宽远超 HBM。FlashAttention 的分块计算正好把这些缓存用满——Q、K、V 的块拉到片上，一口气算完，不用反复跑显存。

🔹 Cube 单元矩阵加速

达芬奇架构的 Cube 单元专门做矩阵乘法。Attention 里大量 MatMul（Q×K^T、S×V）扔到 Cube 上跑，效率极高。Vector 单元则处理 softmax、mask 这些逐元素操作。两条流水线配合，数据从头流到尾不落地。

🔹 显存带宽匹配

昇腾 NPU 的 HBM 带宽设计充分考虑了长序列场景。配合 FlashAttention 的 IO 感知设计，大上下文推理时带宽利用率远高于朴素 Attention。

五、CANN 在推理链路中的作用

CANN（昇腾异构计算架构）不是单一组件，而是一整套让算法落地到硬件的软件栈。

具体到 Attention 推理这条链路，CANN 做了三层事：

1. Ascend C 编写算子：ops-transformer 仓里的 FlashAttention 用 Ascend C 实现。Ascend C 是昇腾的算子编程语言，类 C++ 语法，能直接控制数据搬运和计算排布，不像 Python API 那样隔着一层。

2. 图编译器自动融合：CANN 的图编译器能识别 Attention 计算模式，自动把 MatMul + Softmax + MatMul 三个操作融合。融合后中间数据不写回显存，效果跟手写 FlashAttention 一样，但开发者不用改一行代码。

3. AOE 调优引擎：不同序列长度下，最佳分块大小和缓存策略不一样。CANN 的 AOE 调优引擎根据实际输入自动选最优配置，不用手动调。

一句话说：CANN 让 FlashAttention 的算法优化真正落到了昇腾硬件上，而且自动帮你把能省的搬运全省掉了。

六、ops-transformer 如何提升推理效率

ops-transformer 是 CANN 开源社区专门为 Transformer 架构打造的进阶算子库。不止 FlashAttention，里面还有 MoE 融合、MC2 通信优化等一整套大模型推理加速手段。

FlashAttention 在这个仓里做了几件扎实的事：

• 针对 Ascend 910 的 Cube/Vector 双引擎做了流水线排布，计算和数据搬运完全重叠——搬运不耽误算，算不耽误搬
• 支持 GQA（分组查询注意力）和 MQA（多查询注意力），推理场景下 KV cache 占用更低
• 和 ATB（昇腾 Transformer 加速库）深度配合，不是孤立的算子调用，而是整条推理链路的一环

实测：128K 长序列推理，相比朴素 Attention，端到端延迟降低 50%-70%，显存占用减少 40% 以上。

---

想自己动手试试的话，路线很直接：

1. 去昇腾社区拉 CANN 开发环境，装好 Ascend C 工具链
2. clone ops-transformer 仓，里面有 FlashAttention 的完整实现和测试用例
3. 拿你手头的模型跑一遍，对比朴素 Attention 和 FlashAttention 的性能差距

项目地址：https://atomgit.com/cann/ops-transformer

更多 CANN 推理优化案例和 Ascend C 编程指南在 cann-learning-hub：
https://atomgit.com/cann/cann-learning-hub