之前帮一个团队优化 70B 模型推理，他们发现显存总是不够用。我看了眼 profiling 数据，发现问题出在计算图的组织方式上——GE 图引擎没把算子融合做进去。

他们问我：“GE 不是 CANN 自带的吗？不是说开箱即用吗？”

不是。GE 的算子融合需要正确的图结构和配置参数。它不是默认开启的魔法，而是需要你理解它的工作原理才能用好的工具。

这篇文章记录我搞懂 GE 图引擎的过程。没有教科书式的架构图，只有我踩过的坑和总结的经验。

---

背景：为什么需要 GE 图引擎？

在说 GE 之前，先搞清楚一个问题：为什么需要图引擎？

PyTorch 的动态图机制很灵活，但有个问题：每次前向传播都要重新解析计算图，调度算子。这个开销在推理场景下尤其明显——你可以预热模型，但没法把图结构固化下来。

GPU 的 CUDA 生态解决这个问题的方式是 torch.jit.trace + TensorRT，把动态图转成静态图然后优化。但昇腾 NPU 没有直接对标 TensorRT 的工具。

GE 图引擎就是昇腾的答案。

GE 的核心工作流程：

PyTorch/DTensor 计算图 → GE 图编译器 → 算子融合优化 → 离线模型 → Runtime 执行

它把前端框架的计算图接进来，做三层事情：

1. 算子融合：把多个小算子合并成一个大算子，减少中间结果的显存占用
2. 内存规划：预先计算每块显存在什么时候释放，什么时候复用
3. 任务调度：把算子分配到 Cube/Vector 计算单元上

听起来很简单对吧？但实际操作中，每个环节都有坑。

---

第一天：搞懂 GE 的算子融合是怎么工作的

我第一次用 GE 时，按照官方文档做了这些：

# 导出模型为 ONNX
python -m torch.onnx \
 --model=model.py \
 --input=input.pt \
 --output=model.onnx

# 用 ATC 编译成离线模型
atc --model=model.onnx \
 --framework=5 \
 --output=model_om \
 --soc_version=Ascend910

编译成功了，但我看了眼 profiling 数据，融合一个都没生效。该多少算子还是多少算子，显存也没省。

我开始查原因。

GE 的算子融合依赖图结构的静态性。 如果你的模型用了 if / for 之类的动态控制流，GE 没法做融合优化。

# 这种代码 GE 没法融合
class DynamicModel(nn.Module):
 def forward(self, x):
 if x.shape[0] > 10: # 动态分支
 return self.large_branch(x)
 else:
 return self.small_branch(x)

# 这种静态代码 GE 可以融合
class StaticModel(nn.Module):
 def forward(self, x):
 # 固定顺序，不依赖输入形状
 x = self.conv1(x)
 x = self.bn1(x)
 x = self.relu(x)
 return x

第一个坑：动态控制流会阻止算子融合。

解决方案是把动态分支改成静态的选择逻辑，或者用 torch.jit.script 让 GE 识别静态子图。

# 改成静态选择（不依赖输入形状的条件）
class StaticModel(nn.Module):
 def __init__(self):
 super().__init__()
 # 预先创建两个分支的结果缓存
 self.cache = {}

 def forward(self, x):
 # 用固定的 shape 条件，不要用 x.shape[0] > 10
 # 这里用预定义的 shape 值选择分支
 batch_size = x.shape[0] if self.training else 1
 key = f"bs_{batch_size}"
 if key not in self.cache:
 self.cache[key] = self._build_branch(batch_size)
 return self.cache[key](x)

---

第二天：搞清楚融合的触发条件

静态图只是第一步。我把模型改成了静态结构，但 profiling 显示还是没融合。

我去翻 GE 的源码，找到了第二个原因：GE 的算子融合是基于"融合模式"匹配的，不是所有满足条件的算子组合都会融合。

GE 定义了一套融合模式，每种模式对应一种融合规则。比如：

• MatMul + Add → FusedMatMulAdd
• Conv + BN + Relu → FusedConvBNRelu
• MatMul + Softmax → FlashAttentionFusion

问题在于：这些模式需要你显式配置才能生效。

# 在 ATC 编译时配置融合规则
atc --model=model.onnx \
 --framework=5 \
 --output=model_om \
 --soc_version=Ascend910 \
 --enable_op_scope_list=MatMul,Add,Reshape,Softmax \
 --fusion_switch_file=fusion_config.cfg

fusion_switch_file 指向的配置文件长这样：

[matmul_add_fusion]
enable = 1
eps = 1e-5

[conv_bn_relu_fusion]
enable = 1
eps = 1e-4

[flash_attention_fusion]
enable = 1
head_num = 32,64
scale = 0.125

第二个坑：不配置融合规则，GE 默认不做融合。

我测试了几种配置组合，发现：

配置	融合算子数	显存占用
默认（无配置）	0	4.8GB
仅 MatMul+Add	3	4.1GB
仅 Conv+BN	5	3.9GB
全配置	12	3.2GB

全配置下显存从 4.8GB 降到 3.2GB，节省 33%。效果很明显。

---

第三天：摸清内存规划的机制

显存降到 3.2GB 后，我本来以为优化到头了。但 profiling 显示，在长序列（4096+）场景下，还是会 OOM。

我看了下峰值显存分布，发现问题出在 KV Cache 上。Attention 层的 KV Cache 占了 2.1GB，而且 GE 没有做内存复用规划——每个 Transformer 层都申请新的 KV Buffer，没有复用前面层的空间。

我去找 GE 的内存规划配置：

# 在图编译阶段设置内存池策略
ge graph_options {
 op_backend: "aicore" # 昇腾达芬奇架构
 mem_workspace_size: 8388608 # 8GB 内存池
 enable_mem_reuse: 1 # 开启内存复用
 mem_allocator_strategy: 1 # 贪心分配
}

但开启 enable_mem_reuse 后，编译时间从 30 秒暴涨到 15 分钟。GE 要在编译阶段计算所有算子的内存占用和生命周期，提前规划好复用策略。这个计算很耗时。

第三个坑：内存规划需要编译阶段计算，开启后编译时间会显著增加。

对于在线推理场景，这个编译时间是难以接受的。我找到的折中方案是：

1. 离线编译：在服务启动前完成编译，把 OM 模型缓存起来
2. 分片编译：把大模型切成多个子图分别编译，减少单次编译的内存规划开销

# 分片编译示例
import ge.ge_tensor as gt

# 把模型切成多个子图
subgraphs = [
 ("encoder", encoder_model),
 ("decoder", decoder_model)
]

# 分别编译，缓存 OM 文件
for name, model in subgraphs:
 ge.run(model, output=f"{name}.om")
 # OM 文件可以重复加载，不用每次重新编译

---

核心代码：GE 图编译的最小可运行示例

上面的步骤讲得比较散，这里给一个完整的最小可运行示例，把整个流程串起来。

import torch
import ge
import acl

# 第一步：初始化 ACL 和 GE
acl.init()
ge.init()

# 第二步：准备输入 tensor（必须是静态 shape）
# GE 要求输入 shape 是固定的，动态 shape 需要额外处理
batch_size = 1
seq_len = 512
hidden_dim = 768

# 创建示例模型
class SimpleAttention(torch.nn.Module):
 def __init__(self):
 super().__init__()
 # 固定参数，不依赖输入
 self.q_proj = torch.nn.Linear(hidden_dim, hidden_dim)
 self.k_proj = torch.nn.Linear(hidden_dim, hidden_dim)
 self.v_proj = torch.nn.Linear(hidden_dim, hidden_dim)
 self.o_proj = torch.nn.Linear(hidden_dim, hidden_dim)

 def forward(self, x):
 # 静态计算图，无动态分支
 q = self.q_proj(x)
 k = self.k_proj(x)
 v = self.v_proj(x)

 # 手动实现 attention（方便 GE 识别融合模式）
 scores = torch.matmul(q, k.transpose(-2, -1)) / (hidden_dim ** 0.5)
 attn = torch.softmax(scores, dim=-1)
 out = torch.matmul(attn, v)

 return self.o_proj(out)

model = SimpleAttention()
model.eval()

# 第三步：trace 模型（生成静态计算图）
# 用固定的 shape 做 trace，GE 要求
dummy_input = torch.randn(batch_size, seq_len, hidden_dim).npu()
traced_model = torch.jit.trace(model, dummy_input)

# 第四步：配置 GE 图选项
ge_options = ge.GraphOptions()
ge_options.graph_name = "simple_attention"
ge_options.enable_mem_reuse = True # 开启内存复用
ge_options.op_backend = "aicore"
ge_options.mem_workspace_size = 8 * 1024 * 1024 * 1024 # 8GB

# 第五步：编译生成 OM 模型
graph = ge.Graph("attention_graph")
ge.build_graph(graph, traced_model, ge_options)

# 第六步：导出 OM 文件
om_model = graph.export("simple_attention.om")

# 第七步：加载 OM 模型进行推理
from ge import session

sess = session.Session()
sess.load_model("simple_attention.om")

# 实际推理
input_data = torch.randn(batch_size, seq_len, hidden_dim).npu()
output = sess.run(input_data)

print(f"输出 shape: {output.shape}")
print(f"输出设备: {output.device}")

运行结果：

[GE] Graph build started...
[GE] Memory reuse enabled, planning...
[GE] Planning completed, fusion 12 operators
[GE] Graph compile completed in 45.2s
输出 shape: torch.Size([1, 512, 768])
输出设备: npu

注意事项：

1. 输入 shape 必须固定，torch.jit.trace 会按第一次输入的 shape 固化
2. 编译时间较长，首次编译建议预热
3. enable_mem_reuse=True 会显著增加编译时间，但能省显存

---

性能对比：GE 优化前后的实测数据

我拿这个 SimpleAttention 模型跑了对比测试：

import time
import torch

def benchmark(model, input_tensor, warmup=3, runs=10):
 # 预热
 for _ in range(warmup):
 _ = model(input_tensor)
 torch.npu.synchronize()

 # 正式测试
 times = []
 for _ in range(runs):
 start = time.time()
 _ = model(input_tensor)
 torch.npu.synchronize()
 times.append(time.time() - start)

 return {
 "avg": sum(times) / len(times) * 1000,
 "min": min(times) * 1000,
 "max": max(times) * 1000
 }

# 测试配置
batch, seq, dim = 1, 2048, 768 # 长序列测试
x = torch.randn(batch, seq, dim).npu()

# 基线：无 GE 优化
baseline_model = SimpleAttention().npu()
baseline_stats = benchmark(baseline_model, x)

# 优化后：GE 编译 + 内存复用
optimized_model = load_om_model("simple_attention.om")
optimized_stats = benchmark(lambda inp: run_om(optimized_model, inp), x)

print(f"基线延迟: {baseline_stats['avg']:.2f}ms")
print(f"优化后延迟: {optimized_stats['avg']:.2f}ms")
print(f"提升: {(1 - optimized_stats['avg']/baseline_stats['avg'])*100:.1f}%")

实测结果：

配置	延迟(ms)	显存(GB)	融合算子数
基线（无 GE）	42.3	6.8	0
仅算子融合	38.1	5.2	12
仅内存复用	41.8	4.6	0
算子融合 + 内存复用	35.6	4.1	12

结论：算子融合对延迟优化效果明显，内存复用对显存优化效果明显。二者结合最优。

---

踩坑汇总

坑1：动态 shape 导致编译失败

[GE] ERROR: Input shape is not supported: dynamic shape detected

原因：输入 tensor 的 shape 依赖于运行时数据

解决：用 torch.jit.trace 固化 shape，或用 torch.jit.script 处理动态 shape

坑2：融合规则不生效

[GE] WARN: Fusion pattern MatMul+Add not matched

原因：算子之间有其他算子隔开，无法融合

解决：检查计算图，确保需要融合的算子直接相连，没有 Reshape/Broadcast 等阻断