如果说 Runtime 是昇腾 NPU 的执行官,Graph Engine(GE)就是作战参谋部。模型在被送进 Runtime 之前,GE 已经完成了一整套编排工作:把模型的计算图拆解、融合、优化,最终产出一份 Runtime 可以直接"照单执行"的调度方案。

很多人把 GE 当成一个普通的计算图执行框架,但在 CANN 体系里 GE 的角色比这更重——它是连接前端框架和底层硬件的中枢,也是 CANN 能在 Transformer 这类复杂模型上压出性能的关键。

GE 在 CANN 中的位置

CANN 的五层架构里,GE 站在计算编译层和计算服务层的交界处。

应用层(AscendCL)
    ↓
图引擎(GE)← 接收 OM 模型,执行图优化
    ↓
计算编译层(Graph Compiler / ATC)
    ↓
计算执行层(Runtime)
    ↓
硬件层(昇腾达芬奇架构)

通常的推理流程是:OM 模型加载到 GE → GE 解析模型的计算图 → 执行图优化(算子融合、内存复用、Tensor 重排)→ 生成优化后的执行计划 → 交给 Runtime 逐算子分发到 NPU。

GE 不直接参与某个算子的计算,它管的是"宏观调度"——决定算子按什么顺序跑、哪些可以合并、中间 Tensor 放哪里。

为什么 AI 推理需要图执行

推理和训练的最大区别在于:推理的计算图是固定的。模型训练结束后,网络结构就锁死了,每次推理只是用不同的输入数据跑同一张图。

固定图这件事给了 GE 巨大的优化空间。它可以在模型加载阶段花时间做全局优化——分析整张图的数据流、识别可以合并的算子模式、消除冗余的内存搬运——这些工作在训练场景中很难做(因为图在变),但在推理场景中只需要做一次,后续推理全部受益。

不用图执行的方案是逐算子下发:应用层依次调用每个算子的 Launch 函数,每个算子独立申请内存、独立调度。这种方式的问题是算子间的依赖信息被丢失了,GE 无法做跨算子的优化——比如把两个算子的中间 Tensor 直接在片上传递,省掉一次 DDR 读写。

图融合为什么能减少推理耗时

算子融合是 GE 最重要的优化手段。以 Transformer 中的 MHA(Multi-Head Attention)→ Add → LayerNorm → FFN 这个子图为例:

融合前,GE 会拆成 7-8 个独立算子,每个算子:

  1. 从 DDR 读取输入 Tensor
  2. 把计算结果写回 DDR
  3. 下一个算子再读回来

融合后,GE 把 QKV 投影 + Split + Attention Score + Softmax + 输出投影 + Add + LayerNorm 合并成一个 Fusion 算子。中间 Tensor 在片上 AI Core 内部流转,不落 DDR。

数据搬运的耗时通常占单算子执行时间的 60%-80%。融合一次就省掉两级 DDR 读写——融合 4 个算子消除了 6 次搬运。Transformer 推理中,算子融合带来的端到端加速通常在 1.5x-2.5x 之间,具体取决于模型规模和融合策略。

Graph Engine 如何优化 Transformer

Transformer 模型的结构高度规整——重复的 Decoder Block、确定的 Attention 模式、固定的 FFN 结构。GE 针对这类模型做了专门优化。

Pattern Match 融合。 GE 内置了 Transfomer 常见子图的匹配模板。加载 OM 图时,GE 自动扫描图中的算子序列,一旦匹配到 BMM + Softmax + BMM 这类 Attention 模式,自动替换成等效的融合算子 FlashAttention。不需要应用层做任何改动。

内存生命周期分析。 Transformer 推理中每个 Block 产生大量中间 Tensor:Q、K、V、Score、加权和、残差输出。GE 做全局内存分析,判断每个 Tensor 的最短生命周期,让不同 Block 的中间 Tensor 复用同一块显存。以 LLaMA-7B 为例,GE 的内存复用可以将峰值显存占用降低 40% 以上。

Tensor 调度重排。 GE 分析算子计算量和对带宽的需求,把计算密集型的 MatMul 算子和带宽瓶颈型的 Softmax 算子交替排列,让 AI Core 的计算单元和 Memory 搬运单元同时工作。单流执行时计算单元空闲等待数据搬运,重排后两路并行,整体吞吐可提升 20-30%。

昇腾图执行链路解析

完整的 GE 执行链路走六步:

第一步:图解析。 GE 读取 OM 模型中的序列化计算图,重建算子节点和数据边。每个算子节点记录了算子类型、输入输出 Tensor 的形状和数据类型、以及属性参数。

第二步:图分析。 GE 遍历全图做数据流分析。这一步产生两个关键信息:每个 Tensor 的引用计数(哪些算子读它、哪个算子最后写它)和算子之间的依赖关系(哪些算子可以并行执行)。

输入 Tensor A → [算子1] → Tensor B → [算子2] → Tensor C
                                    ↘ [算子3] → Tensor D

分析结果是 GE 知道算子 2 和算子 3 在等待 Tensor B 就绪后可以同时提交。

第三步:图优化。 执行算子融合、常量折叠、无用节点消除等优化 pass。融合策略由内置模式库和 AOE(自动调优引擎)共同决定。

第四步:内存分配。 基于生命周期分析结果,为每个 Tensor 分配物理显存地址。目标是最大化复用、最小化碎片。

第五步:任务生成。 将优化后的图拆解成 Runtime 可以执行的 Task 序列。Task 包含了算子的输入输出地址、执行配置和 Stream 分配信息。

第六步:下发执行。 把 Task 提交到 Runtime 的 Stream 队列。GE 在这一步可以选择逐算子提交或 Batch Launch(一次提交一组算子以减少 Runtime 调度开销)。

大模型推理中的 GE 瓶颈

GE 在模型加载阶段的优化跑得再好,推理阶段的瓶颈依然集中在几个点上。

编译时间。 GE 的图优化 pass 数量多,大型模型(70B+)的 OM 图可能有数万个算子节点。图分析和融合优化在模型加载时需要几秒到几十秒,这个延迟对于需要快速扩缩容的推理服务来说是不能接受的。当前的做法是缓存优化后的执行计划,下次加载直接复用。

动态 Shape 的处理代价。 GE 在推理阶段遇到动态 Shape 输入时,部分优化(如内存分配策略)需要重新计算。每次 recompile 都会打断推理的连续性。GE 的应对是预分配最大 Shape 的缓冲区,在小 Shape 输入时复用,用空间换时间。

结语

Graph Engine 是 CANN 推理链条中最容易被低估的一层。它不直接产出性能数字,但没有 GE 的算子融合和 Tensor 调度优化,Runtime 拿到的就是一张未经优化的算子列表——每个算子在硬件上裸跑,数据在 DDR 和片上存储之间反复搬运,性能可能只有优化后的三分之一。理解了 GE 的执行链路,才算理解了"一张计算图怎么跑上昇腾 NPU"这个问题的全部答案。

CANN Graph Engine 仓库

Graph AutoFusion 优化框架

Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐