有次调一个长文本推理服务，128K 上下文，模型刚跑起来显存就爆了。切到昇腾 NPU，跑 CANN ops-transformer 的 FlashAttention，同样的输入，显存占用直接少了一半。

这才开始认真拆 Attention 到底把显存吃在了哪里。

为什么 Transformer 推理越来越吃显存

先说一个直观的数字。

128K 上下文的推理请求，朴素 Attention 需要在显存里维护一个 128000 × 128000 的矩阵。每个元素 2 字节（FP16），光这一个矩阵就占 32GB 显存。

32GB。

这还只是中间结果。输入 token、模型权重、KV cache 都还没往里加。

模型参数在涨，上下文在拉长，显存却涨不了那么快。你迟早得面对一个问题：能不能不存这个矩阵？

普通 Attention 到底慢在哪里

慢分两种：算得慢，和搬得慢。

普通 Attention 的致命伤是后者。

标准流程拆开看：

Step 1：从 HBM 读 Q、K → 算 S = Q×K^T → 把 S 写回 HBM
Step 2：从 HBM 读 S → 做 softmax → 把 P 写回 HBM
Step 3：从 HBM 读 P、V → 算 O = P×V → 把 O 写回 HBM

一个 128000×128000 的矩阵，来回读写三趟。

这就像一个快递分拣中心，包裹从仓库搬到分拣台，拣完又搬回仓库，下一步又要搬出来。真正在分拣台上的时间不到 30%。

计算单元大半时间在等数据。

FlashAttention 为什么出现

FlashAttention 的解决方式粗暴但有效：

不存。

Q、K、V 切成小块，每次只加载一小块到片上缓存。在片上把 softmax 和加权求和一口气算完，输出直接写回。那个 32GB 的中间矩阵不需要了。

核心技术点叫 online softmax——不用等所有分数到齐再做归一化。每算完一块就地修正统计量，最终结果跟标准 softmax 数学等价。

三句话总结收益：

• 显存读写量 O(n²) → O(n)
• 128K 上下文的中间矩阵从 32GB → 基本为零
• 速度不降反升——计算单元终于不用干等了

省显存和省时间，在这里是同一件事。

昇腾 NPU 如何减少数据搬运

昇腾 NPU 的达芬奇架构在做这件事上有几个天然顺手的地方。

🔹 片上缓存大。 Unified Buffer 和 L1 缓存提供了一个充裕的「工作台」。Q、K、V 的小块拉进来，算完出去，全程不碰 HBM。128K 上下文也一样，每块只有几千个元素。

🔹 Cube + Vector 双引擎并行。 Cube 单元专攻矩阵乘法——Q×K^T、S×V 扔上去全速跑。Vector 单元同时处理 softmax、mask、scale。两条线不是串行等，是同时干。

🔹 DMA 预取。 Cube 算当前块的时候，DMA 引擎已经把下一块数据从 HBM 搬到了缓存。算完当前块，下一块已经在等了。

算和搬完全重叠。

CANN 在执行层中的作用

算子写好了，怎么让它跑起来？CANN 在执行层管三件事：

调度。 FlashAttention 涉及 Cube、Vector、DMA 三个引擎的协同。CANN Runtime 把任务分给正确的引擎，排好时序，不打架。

融合。 图编译器发现你的模型里有 MatMul → Softmax → MatMul 这个模式，自动识别为 Attention，换上融合版 FlashAttention。你一行代码不改。

调优。 AOE 引擎第一次跑的时候自动搜最优的分块大小和缓存策略。短序列和长序列的最优配置不一样，AOE 帮你挑了。

CANN 在这里的价值不是写了什么新算法，而是让已有的算法真正跑在硬件的最优路径上。

ops-transformer 的融合优化思路

ops-transformer 仓里的 FlashAttention 不是把论文算法直译成 Ascend C 就完事了。

它做了三层融合：

算子内融合。 Q×K^T、softmax、×V 三个步骤合并成一个算子，中间数据不写回显存——这是 FlashAttention 本身的思路。ops-transformer 在此基础上针对 Ascend 910 的 Cube/Vector 排布重新做了流水线，计算延迟和搬运延迟完全重叠。

算子间融合。 如果 Q 和 K 来自同一个输入（self-attention 场景），把两个矩阵乘法合并成一次调用。scale 因子和 attention mask 直接在 softmax 之前内联计算，不单独挂算子。每少一个独立算子，就少一次 kernel launch 开销。

场景适配融合。 推理场景用 GQA 和 MQA——KV head 数远少于 Q head 数。ops-transformer 的 FlashAttention 内部做了广播优化，KV 读一次，多个 Q head 共用。配合 ATB（昇腾 Transformer 加速库），不是孤立的算子调一调，而是整条推理链路一起考虑。

最终效果不是 Attention 单独快了，是整个 Transformer block 的端到端延迟降了 50%-70%。

---

想继续深入的话，两条路：

短期上手： clone ops-transformer 仓，跑 FlashAttention 的自带测试用例。换不同序列长度，看显存占用曲线——会比看十篇文章更有感觉。

建立体系： 下一步看 ATB（昇腾 Transformer 加速库）。它把 FlashAttention、MoE 融合、KV cache 管理打包了，是大模型推理的完整加速方案。然后读 CANN 推理优化文档，搞清楚图编译和 AOE 调优的完整链路——你会对整个技术栈有系统理解。

ops-transformer：https://atomgit.com/cann/ops-transformer
ATB：https://atomgit.com/cann/ascend-transformer-boost
CANN 学习中心：https://atomgit.com/cann/cann-learning-hub