已合并
feat(inductor): add experimental planned fast launch for NPU Python Wrapper #42634
feat(inductor): add experimental planned fast launch for NPU Python Wrapper #42634
已合并
linjiyuan创建于 7月24日
linjiyuan
7月24日

【合入来源】

如有社区issue,请关联issue链接
请勿携带内部流程信息(需求链接、问题单、内部issue等)

Fixes https://gitcode.com/Ascend/pytorch/issues/3471

【修改方案】

请描述修改内容的具体实现,涉及哪些组件之间进行交互,可以用1、2、3、...进行罗列
如果是需求或者重构类的PR,需要补充详细设计文档(说明上下游组件关系、时序图、类图、DFX能力等内容)

  1. 新增 NPU Inductor Python Wrapper Planned Fast Launch 实验模块,并通过 TORCHINDUCTOR_NPU_FAST_LAUNCH 环境变量控制启用,默认关闭。
  2. 在 NPU Triton backend 初始化阶段调用 patch_fast_launch(),运行时 patch NPUPythonWrapperCodeGen 的初始化、Triton header 和 kernel call 生成逻辑,以及 TritonCompileResultNpu.make_launcher();patch 支持重复调用幂等,并在批量属性替换失败时回滚已应用的修改。
  3. triton=True、非 C++ Wrapper、非 graph partition 子图且未启用 compile-time autotune 的 callsite 生成 metadata 和单槽调用缓存;首次调用惰性绑定,普通 autotuner 创建 BoundFastLaunch,grouped autotuner 直接缓存原 autotuner.run。C++ Wrapper、graph partition 子图、非 Triton kernel 和 compile-time autotune 继续使用原生成路径。
  4. 在保留原 Triton launcher 的基础上附加 kernel name、kernel stub 及 owner、最终 ABI 参数类型、动态 grid resolver、SIMT/shared memory、FFTS 和隐藏资源需求等元数据;元数据提取失败时保留原 launcher 行为。
  5. 新增 launcher 稳定性判断和运行时路由:普通 autotuner 必要时先通过原 NPUCachingAutotuner.run 完成 autotune、coordinate descent 和 store cubin 生命周期,再将唯一且稳定的 best launcher 提升为 PlannedFastLaunch;grouped autotuner 继续使用其自身的原 autotuner.run 完整入口;profiler、benchmark_run、runtime kwargs、dump_fx_graphcheck_accuracytriton_interpretdump_launch_params 等普通 autotuner 完整入口场景继续使用原 NPUCachingAutotuner.run。plan 已建立后如检测到活动 launch hook,仅本次调用已选中的原 launcher,hook 清空后恢复 planned path。
  6. 新增 schema/ABI、参数类别、stream、grid、workspace、sync lock、device print、FFTS 和 backend 可用性校验;不支持的稳定 launcher 使用负缓存并复用原 launcher,动态 grid 等单次可恢复错误回退原路径,进入 C++ backend 后的异常不重放 kernel,避免重复提交。
  7. 新增 C++ FastLaunchPlan 及 pybind11 内部接口:创建 plan 时保存 kernel name、kernel stub 的 Python 引用和裸指针、归一化 ABI 参数类别、SIMT/shared-memory、force-SIMT 与 FFTS 状态/地址,并预计算隐藏字段、参数字段和 grid 字段的 offset 及 packed buffer 总长度;Python 侧通过 plan._owner 额外持有产生 stub 的 binary owner。每次调用不复用参数缓冲区,而是新建 PackedLaunch:当前 stream 单独保存为 rtStream_t,并依据当前 grid 和 canonical args(Tensor data_ptr、标量及 best_runtime_blocks 补入参数)重新分配和填充 packed args buffer;非 SIMT 分支调用 rtKernelLaunch,SIMT 分支在提交时检查并调用弱符号 rtKernelLaunchWithFlagV2,均经 OpCommand.SetCustomHandler(...).Run() 提交。
  8. 通过 torch_npu/csrc/inductor/CMakeLists.txt 将 fast launch C++ 源文件纳入 INDUCTOR_SRCS;该实现仅在非 BUILD_LIBTORCH 的 CPython 扩展构建中编译生效,并在 torch_npu._C 初始化时注册内部类型 _NPUInductorFastLaunchPlan_npu_inductor_make_fast_launch_plan_npu_inductor_fast_launch_with_plan 两个内部函数。
  9. 新增 Python 单元测试、patch/回滚测试和静态结构测试,覆盖默认关闭、代码生成路由、launcher promotion、schema 冲突、负缓存、动态 grid、profiler、FFTS、隐藏资源、backend 异常及 C++ packed layout 等场景。

组件交互链路:

torch.compile → NPU Triton backend 初始化 →(环境变量开启时)patch_fast_launch → Python Wrapper codegen/launcher metadata → callsite 首次调用惰性绑定。

  • grouped autotuner 继续执行其自身的原 autotuner.run;普通 autotuner 的 full-entry 场景继续执行原 NPUCachingAutotuner.run
  • 普通场景进入 BoundFastLaunch,必要时先执行原 NPUCachingAutotuner.run 完成 autotune/coordinate-descent/store-cubin 生命周期,再由 build_planned_fast_launch 调用 torch_npu._C._npu_inductor_make_fast_launch_plan 创建 C++ plan,并在 Python 侧封装为 PlannedFastLaunch;稳态调用通过 torch_npu._C._npu_inductor_fast_launch_with_planOpCommand.SetCustomHandler(...).Run()rtKernelLaunch/rtKernelLaunchWithFlagV2 提交。不满足 plan 条件的调用继续使用原完整入口或原 launcher。

【资料变更】

请确认是否涉及资料变更。如涉及,需要在PR中体现,并简要说明修改内容。如不涉及,需填写“不涉及”

涉及。

新增 torch_npu/_inductor/docs/npu_inductor_fast_launch_zh.md,说明 Planned Fast Launch 的功能边界、前置条件、环境变量启用方式、最小示例、launcher 生命周期、ABI 支持范围、plan 创建条件、grid/stream 校验、回退与负缓存行为,以及关闭和故障处理方法。

【接口变更】

请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及,需要详细说明接口以及对应的变更内容,同时需要在资料中体现。如不涉及,需填写“不涉及”

涉及。

  1. 新增用户可见环境变量 TORCHINDUCTOR_NPU_FAST_LAUNCH,默认关闭;值经去空白并忽略大小写后为 1trueyesony 时启用。该变量须在首次加载 NPU Triton backend 前设置,推荐在导入 PyTorch/torch_npu 前设置;运行中修改不会补装或撤销 patch,需要重启 Python 进程并重新编译 Wrapper。
  2. 在非 BUILD_LIBTORCHtorch_npu._C Python 扩展中新增 3 个内部 Python 可见符号(均非用户 API):
    • _NPUInductorFastLaunchPlan:无公开构造器,作为 plan 的内部持有类型;
    • _npu_inductor_make_fast_launch_plan
    • _npu_inductor_fast_launch_with_plan
  3. FastLaunchPlan 位于 C++ 匿名命名空间,不新增客户可直接调用的 C++ plan API;现有模型代码及 torch.compile(..., backend="inductor") 调用方式不变。

【功能验证】

说明测试场景,测试方法。如果本次测试方式与常规单元测试不同,请详细说明您的测试步骤
新增/变更内容是否已新增/适配UT测试用例看护,并补充测试自验证截图

已新增以下 UT 看护:

  1. test_npu_fast_launch.py:覆盖环境变量关闭、codegen schema 完整/不完整、plan 生命周期与 binary owner、FFTS ABI、隐藏资源拒绝、warmup 后 promotion、负缓存、profiler、launch hook metadata、动态 grid、backend 异常不重放、launcher 变化和 runtime block 参数等场景。
  2. test_npu_fast_launch_patch.py:覆盖默认关闭时不解析 patch 目标、patch 只安装一次、代码生成路由、launcher metadata 附加,以及 patch 操作失败后的逆序回滚。
  3. test_npu_fast_launch_static.py:覆盖 fast launch 仅通过 runtime patch 接入、package lazy import、C++ backend 保留 OpCommand 和运行时校验、仅暴露 planned launch 接口,以及 packed argument layout 预计算。

本地执行:

python -m unittest discover -s test/_inductor/experimental/python_wrapper_fast_launch -p "test_*.py" -v

执行结果:共 30 个测试,全部通过。
当前自验证为隔离单元测试和静态结构测试,未执行 NPU/CANN 实机功能及性能测试。
当前 HEAD 已通过文档门禁(markdownlint、链接有效性、资源存在性和标签闭合检查);本次推送后尚未确认完整 PR CI。

【CheckList】
PR提交人对以下CheckList自检项进行全量自检,自检通过或不涉及,均修改 [ ] 为 [x]

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 linjiyuan 的贡献)
Llinjiyuan
7月24日 创建了 pull request,commit bf0ca844
atomgit-bot
atomgit-bot
7月24日 评论:

变更摘要

此 PR 将 NPU Inductor 的快速启动(fast launch)功能重构为通过运行时补丁(runtime patch)机制按需启用,而非默认加载。核心思路是引入 python_wrapper_fast_launch 子包,在 INDUCTOR_ASCEND_FAST_LAUNCH 环境变量开启时,对 NPUPythonWrapperCodeGen.generate_kernel_callTritonCompileResultNpu.make_launcher 两个关键路径进行 monkey-patch,将内核发射从原有的 autotuner 流程替换为预计算布局的 PlannedFastLaunch 计划式发射,并通过 C++ 层面的 FastLaunchPlan 实现参数打包与直接内核调用,绕过 Python 层的重复序列化开销。

主要改动

  • 新增 python_wrapper_fast_launch 子包:包含 backend.py(计划构建与 PlannedFastLaunch 可调用对象)、bind.pyBoundFastLaunch 状态机——warmup 后自动升级为计划发射,支持负缓存与 schema 补全)、codegen.py(调用点元数据构建)、launcher.py(为 launcher 附加 ABI 元数据)、patch.py(运行时 monkey-patch 入口,按环境变量 INDUCTOR_ASCEND_FAST_LAUNCH 控制启用)、wrapper_codegen.pyFastLaunchWrapperEmitter 生成 Python wrapper 中的内核调用代码)、types.pyFastLaunchErrorFastLaunchPlanUnavailable 异常)等模块

  • 新增 C++ 绑定 csrc/bindings.cpp:实现 MakeFastLaunchPlan(预计算参数布局 BuildPackedLayout、对齐、FFTS 地址、SIMT 配置等)和 FastLaunchWithPlan(按布局打包参数、校验 grid 范围、通过 OpCommand::SetCustomHandler 提交 rtKernelLaunchrtKernelLaunchWithFlagV2),并在 InitNpuBindings.cpp 中注册 _npu_inductor_make_fast_launch_plan_npu_inductor_fast_launch_with_plan 两个 Python 绑定

  • 补丁式集成 patch_fast_launch():在 torch_npu/_inductor/__init__.py_load_triton_backend() 末尾调用 patch_fast_launch();patch 仅在环境变量 INDUCTOR_ASCEND_FAST_LAUNCH 为 truthy 时执行,对 NPUPythonWrapperCodeGengenerate_kernel_callwrite_triton_header_once 进行包装以输出 bind_python_wrapper_kernel_fast 调用代码,同时对 TritonCompileResultNpu.make_launcher 进行包装以附加网格解析器与 ABI 元数据

  • BoundFastLaunch 调用状态机:首次调用执行原始 autotuner 运行(warmup),之后在 launcher 稳定(单一 launcher、无 coordinate descent、无 store-cubin)且 schema 兼容时自动提升为 PlannedFastLaunch 计划发射;负缓存机制对不兼容的 launcher 跳过重复检测;profiler 开启或 dump/accuracy 模式下强制回退全路径

  • 测试覆盖:新增 test_npu_fast_launch.py(计划构建、升格、负缓存、schema 补全/冲突、grid 校验、backend 错误处理等单元测试)、test_npu_fast_launch_patch.py(patch 启用/关闭、代码生成路由、失败回滚等测试)、test_npu_fast_launch_static.py(静态代码检查——确保 fast launch 代码仅在 patch 激活时进入 wrapper 和 triton_heuristics 路径)

likedislike
atomgit-bot
atomgit-bot
7月24日 评论:

代码审查

I've completed a thorough review of all 21 files. Here is my closing summary:


Closing Summary

审查的 21 个变更文件:

  1. test/_inductor/experimental/__init__.py — 空文件,无问题。
  2. test/_inductor/experimental/python_wrapper_fast_launch/__init__.py — 空文件,无问题。
  3. test/_inductor/experimental/python_wrapper_fast_launch/test_npu_fast_launch.py — 测试文件,覆盖全面,无问题。
  4. test/_inductor/experimental/python_wrapper_fast_launch/test_npu_fast_launch_patch.py — 测试文件,无问题。
  5. test/_inductor/experimental/python_wrapper_fast_launch/test_npu_fast_launch_static.py — 静态检查测试,无问题。
  6. torch_npu/_inductor/__init__.py — 新增 2 行(import 和调用),无问题。
  7. torch_npu/_inductor/docs/npu_inductor_fast_launch_zh.md — 文档,无安全隐患,无危险操作指令。
  8. torch_npu/_inductor/experimental/__init__.py — 空文件,无问题。
  9. torch_npu/_inductor/experimental/python_wrapper_fast_launch/__init__.py — 延迟导入模式,无问题。
  10. torch_npu/_inductor/experimental/python_wrapper_fast_launch/backend.py — 核心 Plan 构建和调用逻辑,参数校验完备,ABI 处理正确,无问题。
  11. torch_npu/_inductor/experimental/python_wrapper_fast_launch/bind.py — 绑定/路由逻辑,正负缓存和回退语义正确,无问题。
  12. torch_npu/_inductor/experimental/python_wrapper_fast_launch/codegen.py — 代码生成 metadata 构建,无问题。
  13. torch_npu/_inductor/experimental/python_wrapper_fast_launch/csrc/bindings.cpp — C++ backend,内存管理正确(shared_ptr / move 语义),边界检查严格(grid 维度、INT32_MAXUINT16_MAX 等),无问题。
  14. torch_npu/_inductor/experimental/python_wrapper_fast_launch/csrc/bindings.h — 头文件,无问题。
  15. torch_npu/_inductor/experimental/python_wrapper_fast_launch/env.py — 环境变量辅助,无问题。
  16. torch_npu/_inductor/experimental/python_wrapper_fast_launch/launcher.py — Launcher metadata 附加,无问题。
  17. torch_npu/_inductor/experimental/python_wrapper_fast_launch/patch.py — 发现 1 个 P3 问题(except Exception 过于宽泛)。
  18. torch_npu/_inductor/experimental/python_wrapper_fast_launch/types.py — 异常类型定义,无问题。
  19. torch_npu/_inductor/experimental/python_wrapper_fast_launch/wrapper_codegen.py — 代码生成发射器,无问题。
  20. torch_npu/csrc/InitNpuBindings.cpp — 注册绑定调用,无问题。
  21. torch_npu/csrc/inductor/CMakeLists.txt — 新增 GLOB 条目,遵循现有模式,无问题。

发现汇总:

优先级 数量
P0 0
P1 0
P2 0
P3 1

整体风险评估:低风险

该变更是对 fast launch 功能的完整实现重构,代码质量较高:C++ 端使用 shared_ptr 管理 Plan 生命周期、通过 static_assert 校验 ABI 结构体布局、对 grid/参数进行多层边界检查;Python 端的正负缓存、回退和事务性 Patch 回滚机制设计清晰。唯一的发现是 patched_make_launcher 中宽泛的 except Exception 可能给问题诊断带来不便,但这是"最佳努力"设计意图的体现,且回退路径完整,不构成运行时风险。

⚠️ 已识别出整体风险,但无法提取行内评论,请参考整体评估。

likedislike
ascend-robotascend-robot成员
7月24日 添加了label:stat/needs-squash
此处折叠了159条消息 查看更多
ascend-robotascend-robot成员
27 天前 添加了label:approved
ascend-robotascend-robot成员
27 天前 关闭了关联的issue
ascend-robotascend-robot成员
27 天前 合入了pull request
ascend-robot
ascend-robot成员
27 天前 评论:

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike
ascend-robot
ascend-robot成员
27 天前 评论:
流水线 pytorch_gitcode_PR_multiVersion#13371 [ commitID:2701f7b3 ] 已完成
likedislike