Ascend for PyTorch 训练营学习笔记 续二

茶思屋专题：昇腾（Ascend）for PyTorch 训练营(https://www.chaspark.com/#/s/AscendforPyTorch2026?multi=zh) 第一期

算子适配——自动生成适配代码：高效、解耦地适配 PyTorch 海量算子

• 版本解耦：无论 PyTorch 是 2.6、2.7 还是 2.x，torch_npu 各版本都通过统一的 Op-plugin 接口对接底层 CANN aclnn。算子实现不感知 PyTorch 版本变化，大幅降低维护成本。
• 自动生成机制：
  • 常规方式：通过 yaml 配置 自动生成算子注册和适配代码，开发者只需手写对接 aclnn 的 C++ 逻辑。
  • 结构化适配：部分简单算子连 C++ 都不用写，纯 yaml 配置即可完成全部适配。
  • 未来规划：用 AI Agent 自动生成 yaml，进一步降低人工工作量。

通过"统一插件接口 + yaml 自动生成"实现算子适配的版本解耦和效率提升

接着讲分布式并行，咱略过

Eager 模式调度：三大运行时优化手段

① 内存池：通过虚拟内存管理减少碎片、内存快照、自定义分配器，降低显存管理开销。

② 两级流水线下发算子：kernel launch 由专用线程异步执行（dequeue 流水线），实现 CPU 与 NPU 并行；同时支持用户自定义算子接入该流水线。

③ 绑核：自动将业务进程/线程绑定到特定 CPU 核心，保证 NUMA 亲和性，关键线程独占 CPU，减少调度抖动。

NPUGraph 解决"算子 launch 太多"的问题，Eager 优化则从内存、异步流水线和 CPU 绑核三方面解决"单次 launch 效率"的问题。

NPUGraph 使能下沉执行：对标 CUDA Graph 的 NPUGraph 机制，减少 CPU 反复 launch 算子的开销：

• 传统 Eager 模式下，每个算子都要经过一次 CPU 下发（Launch A/B/C/D/E），累积大量 launch latency；NPUGraph 先将计算图 Capture（构建） 好，之后只需一次 Launch Graph 即可批量执行，显著节省 CPU 时间和端到端耗时。
• 三层 API 对齐 CUDA：
  • ① torch.npu.NPUGraph：对标 torch.cuda.CUDAGraph，Eager 模式下手动控制 capture + replay。
  • ② torch.compile(...)：图模式下自动、无感地复用 NPUGraph 能力。
  • ③ torch.npu.make_graphed_callables：对标 make_graphed_callables，支持自动求导场景下的图捕获。
• 底层：基于 CANN 的 aclgraph runtime API 实现。