Pull Request已成功合入, 合并人@ascend-robot
(感谢 AllenGuan 的贡献)变更摘要
此 PR 针对 triton_experimental 模块进行兼容性适配,使其能正常对接 vendored triton core 3.5.0(triton-ascend)。主要解决两个问题:一是 launch hooks 在 triton 3.5.0 中已从 CompiledKernel 类属性迁移至 knobs.runtime.*(HookChain),通过引入 knobs 兼容层实现新旧版本共存;二是修复 launcher 签名参数过滤逻辑,将 XBLOCK、YBLOCK、num_warps 等所有 constexpr 隐式常量统一从 def_args 和 call_args 中排除,避免调用方因缺失这些非运行时参数而报错。
主要改动
-
引入
knobs兼容层适配 launch hooks:从torch._inductor.runtime.triton_compat新增导入knobs,在构建 kernel scope 时判断knobs是否为None来选择使用旧式binary.__class__.launch_enter_hook/launch_exit_hook或新式knobs.runtime.launch_enter_hook/knobs.runtime.launch_exit_hook,确保同时兼容 triton 3.5.0 和旧版本。 -
重构隐式常量过滤逻辑:将原来仅过滤
num_warps和num_stages的逻辑扩展为implicit_constants集合(包含num_warps、num_stages以及known_constants中的 XBLOCK/YBLOCK/ZBLOCK/R0_BLOCK 等 constexpr),并从def_args(launcher Python 签名)和call_args(传递给 ascend C runner 的参数)中统一移除,与上游triton_heuristics._get_arg_lists行为对齐。 -
移除旧的字符串替换方式:原逻辑将
num_warps/num_stages替换为compile_meta["constants"]中的字符串值再拼接进call_args,现改为直接过滤掉所有隐式常量,因为这些 constexpr 已编译进 kernel 二进制,无需由调用方传入。




【合入来源】
【修改方案】
torch_npu/_inductor/triton_experimental/compat.py:模块级常量IS_TRITON_36_PLUS = _has_knobs and not _has_attrs_descriptor,以能力探测而非版本字符串解析判别 triton-ascend >= 3.6(vendored core >= 3.5.0)的 API 面——探测锚点为from triton import knobs是否可用、以及triton.backends.compiler/triton.compiler.compiler两处是否仍存在AttrsDescriptor。锚点取值:3.2.2 →(无 knobs、有 AttrsDescriptor)→ False;3.6.0 →(有 knobs、无 AttrsDescriptor)→ True;混合形态判 False 落 legacy 分支(安全方向)。make_launcher):if IS_TRITON_36_PLUS: knobs.runtime.launch_enter_hook/launch_exit_hook / else: binary.__class__.launch_enter_hook/launch_exit_hook(core 3.5.0 将 hooks 迁至triton.knobs.runtime.*的 HookChain,类属性已不存在);triton_version_uses_attrs_dict()为 True)构建implicit_constants = {num_warps, num_stages} ∪ fn.constexprs并同时从def_args/call_args剔除(Config 烘入的 constexpr 不应出现在 launcher 签名与传给 C runner 的实参中),对齐上游triton_heuristics._get_arg_lists的 implicit_constants 语义。add_constexpr_arg):3.6 分支额外将 BLOCK constexpr 以ConstexprArg纳入 signature(上游 3.5.0 要求 constexpr 出现在 signature 中,3.2.x 不需要)。_npu_emit_rsplit_combine):手写 combine kernel 的triton_meta["configs"]构造改为双分支——3.6 产出 plain dict{(arg_idx,): [["tt.divisibility", 16]]}(core 3.5.0ASTSourceattrs /ASTFunction.deserialize的消费格式,与 torchAttrsDescriptorWrapper的 dict 分支一致);legacy 分支保留原AttrsDescriptor.from_dictpayload 逐字不变,且 import 移入分支内使 3.6 永不执行。IS_TRITON_36_PLUS单一判据 + 使用点显式if IS_TRITON_36_PLUS / else双 block(而非逐点散落knobs is None/ torch 谓词、或复用 torchAttrsDescriptorWrapper将分支逻辑藏进上游内部)——目的是版本差异集中可 grep、legacy 分支与旧代码逐字等价(老版本行为零变更的可审查证据)、后续整体迁移 3.6 时可按 else 臂逐块删除。详细设计/根因归因文档:
docs/未归档/triton_experimental_triton版本兼容终极方案.md、docs/未归档/triton_experimental兼容triton3.6根因分析与3.2.2兼容性评估.md(含三段对照实验的证据链、双版本 API 锚点实测表、遗留风险登记)。【资料变更】
不涉及
【接口变更】
不涉及
【功能验证】
验证环境:openEuler 24.03 LTS SP3 (aarch64) / 910B2 / CANN 9.1.0-beta.1 / torch 2.13.0 / triton-ascend 3.6.0(vendored core 3.5.0)。
测试方式与结果(与常规单测的差异:每轮均先清空 Inductor/Triton 缓存目录后重跑,确保不命中陈旧编译产物):
test/_inductor/test_triton_experimental_enable.py(三入口 options/config/env × 多 shape/dtype + default/experimental 同进程隔离,共 13 例):AttributeError: type object 'CompiledKernel' has no attribute 'launch_enter_hook';TORCHINDUCTOR_NPU_BACKEND的环境运行(隔离用例的裸 compile 须落 default 后端)。reshape(...).float().sum(0)类最小图,修改前ImportError: cannot import name 'AttrsDescriptor' from 'triton.compiler.compiler'(整个 torch.compile 失败);修改后 3 个代表用例(view_sum_dyn/view_sum_bcast_dyn/rms_norm_bw_dyn,动态 batch 迭代多轮、compiled 与 eager 逐轮数值断言)全部 PASS。launcher() missing 1 required positional argument: 'XBLOCK'(外层表现为RuntimeError: All tiling for [...] are not runnable.);修复后随 1 一并通过。【CheckList】