Pull Request已成功合入, 合并人@ascend-robot
(感谢 linjiyuan 的贡献)变更摘要
此 PR 将 NPU Inductor 的快速启动(fast launch)功能重构为通过运行时补丁(runtime patch)机制按需启用,而非默认加载。核心思路是引入 python_wrapper_fast_launch 子包,在 INDUCTOR_ASCEND_FAST_LAUNCH 环境变量开启时,对 NPUPythonWrapperCodeGen.generate_kernel_call 和 TritonCompileResultNpu.make_launcher 两个关键路径进行 monkey-patch,将内核发射从原有的 autotuner 流程替换为预计算布局的 PlannedFastLaunch 计划式发射,并通过 C++ 层面的 FastLaunchPlan 实现参数打包与直接内核调用,绕过 Python 层的重复序列化开销。
主要改动
-
新增
python_wrapper_fast_launch子包:包含backend.py(计划构建与PlannedFastLaunch可调用对象)、bind.py(BoundFastLaunch状态机——warmup 后自动升级为计划发射,支持负缓存与 schema 补全)、codegen.py(调用点元数据构建)、launcher.py(为 launcher 附加 ABI 元数据)、patch.py(运行时 monkey-patch 入口,按环境变量INDUCTOR_ASCEND_FAST_LAUNCH控制启用)、wrapper_codegen.py(FastLaunchWrapperEmitter生成 Python wrapper 中的内核调用代码)、types.py(FastLaunchError与FastLaunchPlanUnavailable异常)等模块 -
新增 C++ 绑定
csrc/bindings.cpp:实现MakeFastLaunchPlan(预计算参数布局BuildPackedLayout、对齐、FFTS 地址、SIMT 配置等)和FastLaunchWithPlan(按布局打包参数、校验 grid 范围、通过OpCommand::SetCustomHandler提交rtKernelLaunch或rtKernelLaunchWithFlagV2),并在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 时执行,对NPUPythonWrapperCodeGen的generate_kernel_call和write_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 路径)


代码审查
I've completed a thorough review of all 21 files. Here is my closing summary:
Closing Summary
审查的 21 个变更文件:
test/_inductor/experimental/__init__.py— 空文件,无问题。test/_inductor/experimental/python_wrapper_fast_launch/__init__.py— 空文件,无问题。test/_inductor/experimental/python_wrapper_fast_launch/test_npu_fast_launch.py— 测试文件,覆盖全面,无问题。test/_inductor/experimental/python_wrapper_fast_launch/test_npu_fast_launch_patch.py— 测试文件,无问题。test/_inductor/experimental/python_wrapper_fast_launch/test_npu_fast_launch_static.py— 静态检查测试,无问题。torch_npu/_inductor/__init__.py— 新增 2 行(import 和调用),无问题。torch_npu/_inductor/docs/npu_inductor_fast_launch_zh.md— 文档,无安全隐患,无危险操作指令。torch_npu/_inductor/experimental/__init__.py— 空文件,无问题。torch_npu/_inductor/experimental/python_wrapper_fast_launch/__init__.py— 延迟导入模式,无问题。torch_npu/_inductor/experimental/python_wrapper_fast_launch/backend.py— 核心 Plan 构建和调用逻辑,参数校验完备,ABI 处理正确,无问题。torch_npu/_inductor/experimental/python_wrapper_fast_launch/bind.py— 绑定/路由逻辑,正负缓存和回退语义正确,无问题。torch_npu/_inductor/experimental/python_wrapper_fast_launch/codegen.py— 代码生成 metadata 构建,无问题。torch_npu/_inductor/experimental/python_wrapper_fast_launch/csrc/bindings.cpp— C++ backend,内存管理正确(shared_ptr/ move 语义),边界检查严格(grid维度、INT32_MAX、UINT16_MAX等),无问题。torch_npu/_inductor/experimental/python_wrapper_fast_launch/csrc/bindings.h— 头文件,无问题。torch_npu/_inductor/experimental/python_wrapper_fast_launch/env.py— 环境变量辅助,无问题。torch_npu/_inductor/experimental/python_wrapper_fast_launch/launcher.py— Launcher metadata 附加,无问题。torch_npu/_inductor/experimental/python_wrapper_fast_launch/patch.py— 发现 1 个 P3 问题(except Exception过于宽泛)。torch_npu/_inductor/experimental/python_wrapper_fast_launch/types.py— 异常类型定义,无问题。torch_npu/_inductor/experimental/python_wrapper_fast_launch/wrapper_codegen.py— 代码生成发射器,无问题。torch_npu/csrc/InitNpuBindings.cpp— 注册绑定调用,无问题。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 可能给问题诊断带来不便,但这是"最佳努力"设计意图的体现,且回退路径完整,不构成运行时风险。
⚠️ 已识别出整体风险,但无法提取行内评论,请参考整体评估。


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.




【合入来源】
Fixes https://gitcode.com/Ascend/pytorch/issues/3471
【修改方案】
TORCHINDUCTOR_NPU_FAST_LAUNCH环境变量控制启用,默认关闭。patch_fast_launch(),运行时 patchNPUPythonWrapperCodeGen的初始化、Triton header 和 kernel call 生成逻辑,以及TritonCompileResultNpu.make_launcher();patch 支持重复调用幂等,并在批量属性替换失败时回滚已应用的修改。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 继续使用原生成路径。NPUCachingAutotuner.run完成 autotune、coordinate descent 和 store cubin 生命周期,再将唯一且稳定的 best launcher 提升为PlannedFastLaunch;grouped autotuner 继续使用其自身的原autotuner.run完整入口;profiler、benchmark_run、runtime kwargs、dump_fx_graph、check_accuracy、triton_interpret、dump_launch_params等普通 autotuner 完整入口场景继续使用原NPUCachingAutotuner.run。plan 已建立后如检测到活动 launch hook,仅本次调用已选中的原 launcher,hook 清空后恢复 planned path。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(Tensordata_ptr、标量及best_runtime_blocks补入参数)重新分配和填充 packed args buffer;非 SIMT 分支调用rtKernelLaunch,SIMT 分支在提交时检查并调用弱符号rtKernelLaunchWithFlagV2,均经OpCommand.SetCustomHandler(...).Run()提交。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两个内部函数。组件交互链路:
torch.compile→ NPU Triton backend 初始化 →(环境变量开启时)patch_fast_launch→ Python Wrapper codegen/launcher metadata → callsite 首次调用惰性绑定。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_plan→OpCommand.SetCustomHandler(...).Run()→rtKernelLaunch/rtKernelLaunchWithFlagV2提交。不满足 plan 条件的调用继续使用原完整入口或原 launcher。【资料变更】
涉及。
新增
torch_npu/_inductor/docs/npu_inductor_fast_launch_zh.md,说明 Planned Fast Launch 的功能边界、前置条件、环境变量启用方式、最小示例、launcher 生命周期、ABI 支持范围、plan 创建条件、grid/stream 校验、回退与负缓存行为,以及关闭和故障处理方法。【接口变更】
涉及。
TORCHINDUCTOR_NPU_FAST_LAUNCH,默认关闭;值经去空白并忽略大小写后为1、true、yes、on或y时启用。该变量须在首次加载 NPU Triton backend 前设置,推荐在导入 PyTorch/torch_npu 前设置;运行中修改不会补装或撤销 patch,需要重启 Python 进程并重新编译 Wrapper。BUILD_LIBTORCH的torch_npu._CPython 扩展中新增 3 个内部 Python 可见符号(均非用户 API):_NPUInductorFastLaunchPlan:无公开构造器,作为 plan 的内部持有类型;_npu_inductor_make_fast_launch_plan;_npu_inductor_fast_launch_with_plan。FastLaunchPlan位于 C++ 匿名命名空间,不新增客户可直接调用的 C++ plan API;现有模型代码及torch.compile(..., backend="inductor")调用方式不变。【功能验证】
已新增以下 UT 看护:
test_npu_fast_launch.py:覆盖环境变量关闭、codegen schema 完整/不完整、plan 生命周期与 binary owner、FFTS ABI、隐藏资源拒绝、warmup 后 promotion、负缓存、profiler、launch hook metadata、动态 grid、backend 异常不重放、launcher 变化和 runtime block 参数等场景。test_npu_fast_launch_patch.py:覆盖默认关闭时不解析 patch 目标、patch 只安装一次、代码生成路由、launcher metadata 附加,以及 patch 操作失败后的逆序回滚。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]