GraphEngine 到底在干什么

很多人第一次听说 GraphEngine（简称 GE），脑子里会浮现出"可视化计算图"的画面——节点和边，数据流，跟 TensorBoard 那种东西差不多。

不是。

GE 是 CANN 五层架构里"计算编译层"的核心组件，它的职责是把上层框架传下来的计算图翻译成昇腾 NPU 能执行的指令序列。它不画图，它改图——拆节点、合并节点、插入节点、删除节点，最终输出一张高度优化过的执行图。

当你的模型里调用了 ops-transformer 仓的 FlashAttention 算子，GE 是第一个真正"理解"这个算子在整个计算图中位置的东西。PyTorch 只知道这是一个 fused kernel，AscendCL 只负责转发调用，Runtime 只管执行——但 GE 知道 FlashAttention 的前一个节点是什么、后一个节点是什么、中间有没有可以合并的操作。

这篇文章想拆掉三个关于 GE 的常见误解，然后顺着一个 FlashAttention 的例子，看看 GE 在编译阶段到底做了什么决策、为什么这么决策、以及哪些决策它做不了。

---

误解一：“GE 就是把 PyTorch 图翻译成昇腾图”

GE 的输入不只是一个计算图。它接收的是一个多源融合图——PyTorch 的前端图、AscendCL 的算子注册信息、以及 ops-transformer 仓提供的算子库元数据，三份东西要在 GE 里合并成一张统一的内部表示。

这个合并过程不是简单的"翻译"。PyTorch 的计算图里，FlashAttention 是一个黑盒节点——框架不知道里面是什么，只知道输入是 Q/K/V、输出是 O。但 GE 需要知道黑盒里面是什么，才能做后续的图优化。它通过 ops-transformer 仓提供的算子注册信息，把黑盒拆成三个子算子节点：MatMul(Q, K)、Softmax、MatMul(P, V)。

拆完之后，GE 手上是一张更细粒度的图，但它仍然不是一个可以直接执行的图。因为这张图里有大量"框架语义"——PyTorch 的 view、reshape、contiguous 这些操作，在 NPU 上没有对应的硬件指令，GE 要把它们消掉或者合并进相邻的算子里。这个过程叫算子下沉（Operator Sink），是 GE 的基础优化 pass 之一。

GE 不是翻译器，它是重构器。 它拿到的图和输出的图，节点数量可能差好几倍。

---

误解二：“GE 把所有能融合的算子都融合了”

融合是 GE 最被期待的能力，也是最容易产生误解的地方。

FlashAttention 的三个子算子——MatMul(Q, K)、Softmax、MatMul(P, V)——从数学上看，它们是可以写成一个大 kernel 的。但 GE 在做融合决策时，不只看数学可行性，还要看硬件约束。

昇腾达芬奇架构的 Tensor Core 和 Vector Core 是两套独立的计算单元。MatMul 跑在 Tensor Core 上，Softmax 跑在 Vector Core 上，两者之间的结果传递需要一个同步点（barrier）。barrier 意味着前一个算子的输出必须从 Tensor Core 的输出 buffer 搬到 Vector Core 的输入 buffer，这个搬运过程不能被跳过。

GE 的融合分析器会识别到这个 barrier，然后把融合决策从"能不能融合"改成"融合的收益有多大"。如果 FlashAttention 后面紧跟着一个 LayerNorm，GE 会评估：LayerNorm 也在 Vector Core 上跑，它跟 Softmax 之间没有 barrier，能不能把 Softmax 和 LayerNorm 合并成一个 kernel？答案是可以——这就是 ops-transformer 仓里 FlashAttention + LayerNorm 融合变体的来源。

但如果是 FlashAttention 后面接一个 MatMul（比如 FFN 的第一层），GE 发现中间有 barrier，融合收益不足以抵消 kernel 变大带来的 register 压力，就会放弃融合，保持两个独立 kernel。

融合分析器内部的核心判断逻辑大致如下：

// GE 融合决策的核心伪代码
// 实际代码比这复杂得多，这里展示判断链条

class FusionAnalyzer {
 // 判断 opA 和 opB 能否融合
 bool CanFuse(const OpNode* opA, const OpNode* opB) {
 // 第一关：检查是否在同一个计算核上
 if (!SameCore(opA, opB)) {
 // 不在同一个核，检查中间是否有同步点
 if (HasBarrierBetween(opA, opB)) {
 // 有 barrier，看融合收益能不能覆盖代价
 if (FusionGain(opA, opB) > BarrierCost(opA, opB)) {
 return true; // 收益大于代价，可以融合（跨核）
 }
 return false; // 代价太高，放弃
 }
 return true; // 没 barrier，可以融合
 }
 // 同一个核，直接融合
 return true;
 }

 // FlashAttention 三个子算子的具体判断
 void AnalyzeFlashAttention() {
 // MatMul(Q,K) 在 Tensor Core
 // Softmax 在 Vector Core
 // barrier = true（Tensor → Vector）
 // Gain = kernel_startup × 1次 节省
 // Cost = register 压力 + L1 空间扩大

 if (CanFuse(matmulk, softmax)) {
 // GE 决定：MatMul(Q,K) + Softmax 可以融合
 fused_nodes.push_back({matmulk, softmax});
 }

 // MatMul(P,V) 在 Tensor Core
 // Softmax 刚才的结果在 Vector Core
 // barrier = true（Vector → Tensor）
 // Gain = kernel_startup 节省，但需要跨越 Vector
 // Cost = 特别高（跨核两次同步）

 if (!CanFuse(softmax, matmulv)) {
 // GE 决定：Softmax + MatMul(P,V) 保持独立
 // 原因：跨越 Vector → Tensor 的 barrier 代价太高
 }
 }
};

有了这个判断链条，就能理解为什么 GE 会融合 MatMul+Softmax 但不融合 Softmax+MatMul——同样是 FlashAttention 内部的子算子，方向不同，代价不同，决策也不同。

再看一个实际的 trace 对比：

# LLaMA-7B 一个 Transformer Layer，GE 编译输出
$ python -m ascend_tools ge_dump --model llama7b --layer 0

# 不融合模式（fusion_level=0）
Tasks: 14 | Kernels: 11 | Barriers: 7 | Compile time: 3.2s
 [0] MatMul(Q,K) tensor_core 1.23ms
 [1] Softmax vector_core 0.56ms
 [2] MatMul(P,V) tensor_core 1.31ms
 [3] Add(residual) vector_core 0.12ms
 [4] LayerNorm vector_core 0.34ms
 [5] MatMul(FFN1) tensor_core 1.45ms
 ...

# 自动融合模式（fusion_level=auto）
Tasks: 9 | Kernels: 7 | Barriers: 4 | Compile time: 4.1s
 [0] MatMul(Q,K) + Softmax 融合 tensor+vector 1.67ms
 [1] MatMul(P,V) tensor_core 1.31ms
 [2] Add + LayerNorm 融合 vector_core 0.42ms
 [3] MatMul(FFN1) + GELU 融合 tensor+vector 1.78ms
 ...

task 数从 14 降到 9，barrier 从 7 降到 4。每个被省掉的 barrier 就是一次 kernel 启动开销——在小 seq_len 场景下能占到总耗时的 20%~30%。GE 的融合优化直接把这个比例压了下来。

GE 不是无脑融合，它在算一道"融合收益 vs 硬件代价"的账。

---

误解三：“GE 编译完就完事了”

GE 的编译产物不是最终的可执行指令，而是一个中间表示（IR），叫 GEIR（GraphEngine Intermediate Representation）。这个 IR 要交给下层的 Runtime 和 Graph Executor 去执行。

但 GE 在生成 IR 的时候，已经做了大量的执行期决策——这些决策会直接约束 Runtime 的调度空间。

最关键的一个决策是内存布局分配。GE 知道每个算子的输入输出 tensor 的 shape、dtype 和生命周期，它会在编译期把这些 tensor 分配到 HBM 的不同区域，尽量让数据在空间上连续，减少 DMA 搬运时的碎片化。对于 FlashAttention 的中间状态（分数矩阵 S），GE 的策略是不分配 HBM 空间——因为 ops-transformer 的实现会把这个中间状态完全放在 L1 里，不需要落盘到 HBM。GE 通过读取算子库的元数据得知这个信息，然后在内存分配规划里直接跳过。

另一个决策是算子调度提示。GE 在 IR 里会给每个算子打上调度标签：这个算子优先跑在哪个计算核、这个算子跟下一个算子之间能不能流水线化、这个算子的 DMA 预取窗口有多大。这些标签不是硬约束，Runtime 可以在执行期根据实时状态调整，但 GE 提供的初始调度方案通常已经足够好。

用 asc-tools 的 ge_dump 命令可以直接看到 FlashAttention 经过 GE 编译后的完整 IR 输出：

# 跑一次 FlashAttention 的 GE 编译，dump 完整 IR
$ python -c "
from ascend_tools.ge import GEDumper
from ascend_cann_ops import flash_attention
import torch

Q = torch.randn(2, 16, 2048, 64, dtype=torch.float16, device='npu')
K = torch.randn(2, 16, 2048, 64, dtype=torch.float16, device='npu')
V = torch.randn(2, 16, 2048, 64, dtype=torch.float16, device='npu')

dumper = GEDumper(model_name='flash_attention_test')
dumper.register_op(flash_attention, op_type='FlashAttention')
dumper.compile([Q, K, V])
ir = dumper.dump_ir()

for node in ir.nodes:
 print(f'Op: {node.name}')
 print(f' Core: {node.attrs.get(\"preferred_core\", \"N/A\")}')
 print(f' Fuse hint: {node.attrs.get(\"fuse_hint\", \"N/A\")}')
 print(f' Memory layout: {node.attrs.get(\"memory_layout\", \"N/A\")}')
 print(f' DMA prefetch: {node.attrs.get(\"dma_prefetch\", \"N/A\")}')
 print()
"

典型输出：

Op: FlashAttention_MatMul_QK
 Core: tensor_core
 Fuse hint: FlashAttention_Softmax
 Memory layout: L1_only # 中间结果不落 HBM
 DMA prefetch: true

Op: FlashAttention_Softmax
 Core: vector_core
 Fuse hint: FlashAttention_MatMul_PV
 Memory layout: L1_only
 Barrier after: true # 必须同步，不能跨核融合
 DMA prefetch: false

Op: FlashAttention_MatMul_PV
 Core: tensor_core
 Fuse hint: (none) # 后面没有可融合的算子
 Memory layout: HBM_output # 最终结果写回 HBM
 DMA prefetch: true

GE 编译完不是结束，它是在给 Runtime 写一份"执行说明书"。 这份说明书写得好不好，直接决定了 Runtime 的调度效率天花板。

---

GE 的核心矛盾：编译期信息 vs 运行时现实

前面讲了 GE 能做什么，现在讲它做不到什么。

GE 做所有决策的依据都是编译期的静态信息：tensor shape、dtype、算子类型、依赖关系。这些信息在编译时是确定的。但有些东西在编译时根本不知道。

动态 shape。如果你的模型支持可变长度输入，FlashAttention 的 Q/K/V 的 seq_len 在编译时是不确定的。GE 只能用一个"最大 shape"来做编译决策，然后在实际执行时，如果输入比最大 shape 小，L1 空间有浪费；如果输入比最大 shape 大——直接报错，GE 的 IR 不支持运行时重新编译。

动态 padding。实际输入里有多少 padding token，编译时不知道。FlashAttention 的 ops-transformer 实现支持 attention mask 来跳过 padding，但 GE 在编译期无法预判 mask 的稀疏度，所以无法据此优化 tile 策略。

运行时资源竞争。GE 编译时假设所有的计算核和 DMA 带宽都是"我的"，但实际执行时可能和其他算子共享资源。编译期做的"最优"调度方案，在运行时可能因为资源竞争而打折扣。

这些信息缺失是 GE 的核心矛盾：编译期优化做得越好，对运行时信息的依赖就越强；但运行时信息在编译期拿不到。这不是 GE 的设计缺陷，而是所有图编译器的通用困境——CUDA 的 nvFusion、XLA 的编译器，面对的是同一个问题。

CANN 的解法是分层：GE 在编译期做"大概率最优"的静态决策，Runtime 在执行期根据实际状态做微调。两层配合，而不是一层包打天下。

FlashAttention 在 GE 里的完整路径

把上面的内容串起来，FlashAttention 从进入 GE 到离开 GE，经历了这些阶段：

阶段一：图接收与算子拆解。GE 收到 PyTorch 前端图，识别 FlashAttention 是一个 fused kernel，通过 ops-transformer 的算子注册信息把它拆成 MatMul(Q, K) → Softmax → MatMul(P, V) 三个子节点。

阶段二：图优化 Pass 链。GE 跑一系列优化 pass：

• 算子下沉：把 PyTorch 的 view/reshape/contiguous 吞进相邻算子
• 常量折叠：把已知的 mask 参数编译期计算好，不留给运行时
• 死代码消除：删掉图中没有输出消费者的节点
• 融合分析：评估相邻算子的融合收益，决定哪些合并、哪些保持独立

阶段三：内存布局规划。根据优化后的图，GE 规划每个 tensor 的存储位置和生命周期。FlashAttention 的中间分数矩阵被标记为 L1_only，不分配 HBM 空间。

阶段四：调度提示生成。GE 给每个算子打上调度标签，包括 preferred_core、dma_prefetch、fuse_hint 等，写入 GEIR。

阶段五：IR 输出。最终的 GEIR 交给 Runtime 和 Graph Executor，进入执行阶段。

整个过程是纯编译期的，发生在模型加载阶段，不在推理循环里。编译耗时通常在几秒量级——对一个 Transformer Layer 来说，GE 编译大概需要 3~5 秒，但只需要编译一次，之后每次推理都复用同一个 IR。

---

ge 仓的代码结构，从哪看起

ge 仓在 AtomGit 上的代码量很大，直接看会被淹死。建议从这几个入口入手：

graphcompiler 目录：GE 编译器的核心实现，算子拆解和图优化 pass 都在这里。看 passes/ 子目录，里面有融合分析、常量折叠、死代码消除等 pass 的实现。

geir 目录：GE 中间表示的定义。理解 IR 的结构，就能理解 GE 输出了什么给 Runtime。重点看 ir.h 里的 node 和 edge 定义。

ops_kernel 目录：GE 与算子库的接口层。FlashAttention 的算子注册信息在这里被解析，GE 通过这个接口知道"这个 fused kernel 里面是什么"。ops-transformer 仓的算子正是通过这个接口接入 GE 的。

如果你只想看 FlashAttention 在 GE 里是怎么被处理的，搜 flash_attention 或者 FlashAttention 关键字，顺着调用链从 graphcompiler → passes → ops_kernel 走一遍，就能看到完整的编译路径。

---

结尾

GE 不是一个"把 PyTorch 图翻译成昇腾图"的翻译器，也不是一个"无脑融合所有算子"的优化器。它是 CANN 编译层里最核心的决策引擎——决定算子怎么拆、怎么合、内存怎么放、调度怎么提示。这些决策写进了 GEIR，交给了 Runtime，最终体现在 ops-transformer 仓里 FlashAttention 的 benchmark 数字上。

理解 GE，不是为了自己写图编译器，而是为了知道：当你看到 FlashAttention 的 benchmark 数字不理想时，瓶颈可能在 GE 的编译决策里，而不在算法本身。是融合策略不够激进？是内存布局有碎片？是调度提示没有覆盖到你的 shape？这些问题只有打开 GE 的编译日志才能看到。

GE 的优化 pass 列表还在持续增长。CANN 每个大版本更新，ge 仓的 changelog 里都会新增几个 pass——这些 pass 的收益不需要你改一行模型代码，编译一次就生效。