本文基于昇腾CANN和昇腾NPU,围绕 OM Model 离线模型 技术展开。

ATC 编译产出的 .om 文件不是"一个串行算子列表"。它包含完整的执行计划:哪些阶段可以并行、每个阶段用几个 Stream、每个 Kernel 配哪个 L1 Buffer、中间 Buffer 的生命周期和复用关系。Runtime 加载 OM 后基本不需要做运行时决策——直接照着执行计划跑。


OM 文件的内部结构

OM 是一个二进制文件,由 ATC 的代码生成阶段序列化而成。内部包含四块:

OM 文件结构:

Header:
  - magic: "OM\0"
  - can_version: "8.0"
  - soc_type: "Ascend910"
  - node_count: 120 (融合后的执行单元数)

Weight Section:
  - 所有编译期已固定的权重数据(FP16)
  - 加载时一次拷进 NPU 显存

Execution Graph:
  - 拓扑排序后的执行阶段列表
  - 每个阶段的融合 Kernel 定义 + Stream 分配

Buffer Plan:
  - 每个中间 Tensor 的 Buffer ID、大小、生命周期
  - Pool 的初始分配方案

Runtime 加载 OM 的四步

acl.mdl.load_from_file("model.om") 做四件事:

第一步:头校验。 读 Magic Number 和 CANN 版本。版本不匹配直接报错——OM 文件跨 CANN 大版本不兼容。升级 CANN 后要重新编译一次。

第二步:权重搬运。 把 OM 里的 Weight Section 整个拷到 NPU 显存。一次 aclrtMemcpy 完成,不走任何前向计算。LLaMA-7B 的 14GB 权重约 0.5 秒搬完(PCIe 4.0 x16 ~32GB/s)。

第三步:Buffer Pool 初始化。 根据 OM 里的 Buffer Plan 预申请显存 Pool。Kernel执行前 Runtime 已经把中间结果的 Buffer 都分配好了——零运行时 aclrtMalloc

第四步:调度表注册。 把 Execution Graph 的阶段和依赖关系注册到 Runtime 的任务队列。后续每次 acl.mdl.execute_async 从调度表里按序取阶段执行。

Runtime 加载和执行 OM 的完整链路:

加载:
  acl.mdl.load_from_file("model.om")
    → 头校验(版本 + 芯片型号)
    → 权重搬进显存
    → 创建 Buffer Pool
    → 注册调度表
  → model_id

推理(每次调用):
  acl.mdl.execute_async(model_id, input, output, stream)
    → 取调度表,按阶段顺序
    → for each stage:
        → 分配当前阶段的临时 Buffer(从 Pool)
        → 提交 Kernel 到分配好的 Stream
    → 等所有 Stream 完成
    → 回收临时 Buffer(回 Pool)

Tensor调度和 Kernel 衔接

OM 里的 Buffer Plan 标注了每个中间 Tensor 的生命周期——从哪个 Kernel 产生、哪个 Kernel 消费、什么时候可以释放。两个不重叠的 Tensor 复用同一块 Buffer 地址。Runtime 按这个计划分配——不需要运行时分析。

融合 Kernel 之间如果还有独立算子需要传递数据,Buffer Plan 会标注这两块 Buffer 是否可以做零拷贝衔接——分配为同一块物理地址。Kernel A 的输出直接变成 Kernel B 的输入,省一次 DDR 写回和一次 DDR 读入。


跟 ONNX 的本质区别

ONNX 是"模型描述"——告诉你要算什么,但不说怎么算。OM 是"执行方案"——Kernel 选好了、Buffer 分配好了、Stream 排好了、Stage 拆好了。ONNX 在不同的推理引擎上行为不同,OM 在昇腾 Runtime 上行为确定。

这也是为什么 OM 跨 CANN 版本不兼容——不同版本的 Kernel 实现和融合策略不同,老 OM 里标注的 Kernel 可能在新的 AOL 算子库里已经不存在了。



Kernel执行和 Tensor调度的配合

OM 里的 Execution Graph 不只是"第 N 个阶段跑哪些 Kernel"——它标了每个 Kernel 的 Tensor调度策略:输入从哪个 Buffer 读、输出写到哪个 Buffer、这个 Buffer 在之后的哪个阶段还能复用。

以 FusedAttention 所在阶段为例:Q/K/V 三个输入分配了 3 个 Buffer ID(从 Pool 取),中间 Score 矩阵和 Softmax 结果各占 1 个 Buffer。OM 标了这些 Buffer 在阶段结束后全部释放——下一个阶段 FusedFFN 复用同一批 Buffer 地址。Tensor调度 全部在编译期完成,Runtime 照单执行,零运行时决策。

参考仓库

GE 图引擎

Runtime 运行时

ATC 编译工具

CANN 学习中心

Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐