已合并
fix(triton_experimental): adapt launcher to vendored triton core 3.5.0 #43760
fix(triton_experimental): adapt launcher to vendored triton core 3.5.0 #43760
已合并
AllenGuan创建于 8月4日
AllenGuan成员
8月4日

【合入来源】

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

【修改方案】

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

  1. 新增统一版本开关 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 分支(安全方向)。
  2. launcher 启动适配(npu_triton_heuristics.py make_launcher
    • launch hooks 读取改为 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,类属性已不存在);
    • attrs_dict 分支(core 3.5.0 上 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 语义。
  3. codegen 侧 signature 适配(codegen/triton.py add_constexpr_arg:3.6 分支额外将 BLOCK constexpr 以 ConstexprArg 纳入 signature(上游 3.5.0 要求 constexpr 出现在 signature 中,3.2.x 不需要)。
  4. rsplit-outer combine kernel attrs 适配(codegen/triton.py _npu_emit_rsplit_combine:手写 combine kernel 的 triton_meta["configs"] 构造改为双分支——3.6 产出 plain dict {(arg_idx,): [["tt.divisibility", 16]]}(core 3.5.0 ASTSource attrs / ASTFunction.deserialize 的消费格式,与 torch AttrsDescriptorWrapper 的 dict 分支一致);legacy 分支保留原 AttrsDescriptor.from_dict payload 逐字不变,且 import 移入分支内使 3.6 永不执行。
  5. 关键取舍:所有版本差异点统一收敛到 IS_TRITON_36_PLUS 单一判据 + 使用点显式 if IS_TRITON_36_PLUS / else 双 block(而非逐点散落 knobs is None / torch 谓词、或复用 torch AttrsDescriptorWrapper 将分支逻辑藏进上游内部)——目的是版本差异集中可 grep、legacy 分支与旧代码逐字等价(老版本行为零变更的可审查证据)、后续整体迁移 3.6 时可按 else 臂逐块删除。

详细设计/根因归因文档:docs/未归档/triton_experimental_triton版本兼容终极方案.mddocs/未归档/triton_experimental兼容triton3.6根因分析与3.2.2兼容性评估.md(含三段对照实验的证据链、双版本 API 锚点实测表、遗留风险登记)。

【资料变更】

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

不涉及

【接口变更】

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

不涉及

【功能验证】

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

验证环境: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 缓存目录后重跑,确保不命中陈旧编译产物):

  1. 端到端冒烟(既有用例看护)test/_inductor/test_triton_experimental_enable.py(三入口 options/config/env × 多 shape/dtype + default/experimental 同进程隔离,共 13 例):
    • 修改前(基线):13/13 全部失败——AttributeError: type object 'CompiledKernel' has no attribute 'launch_enter_hook'
    • 修改后:13 passed(清缓存重跑,两轮复现一致)。
    • 运行方式:清缓存后以未设 TORCHINDUCTOR_NPU_BACKEND 的环境运行(隔离用例的裸 compile 须落 default 后端)。
  2. rsplit combine 路径(症状修复前后对照,动态 shape OUTER 归约 e2e):动态首维 + 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。
  3. launcher 参数表(症状二):修复前该最小图在 hooks 修复后显形为 launcher() missing 1 required positional argument: 'XBLOCK'(外层表现为 RuntimeError: All tiling for [...] are not runnable.);修复后随 1 一并通过。
  4. legacy 分支等价性:全部 else 臂与修改前代码逐字一致(仅判据、分支顺序与 import 位置变化);3.2.2 环境的实机回归待后续版本切换验证。

【CheckList】

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

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 AllenGuan 的贡献)
AAllenGuan成员
8月4日 创建了 pull request,commit fcc4fff1
atomgit-bot
atomgit-bot
8月4日 评论:

变更摘要

此 PR 针对 triton_experimental 模块进行兼容性适配,使其能正常对接 vendored triton core 3.5.0(triton-ascend)。主要解决两个问题:一是 launch hooks 在 triton 3.5.0 中已从 CompiledKernel 类属性迁移至 knobs.runtime.*(HookChain),通过引入 knobs 兼容层实现新旧版本共存;二是修复 launcher 签名参数过滤逻辑,将 XBLOCKYBLOCKnum_warps 等所有 constexpr 隐式常量统一从 def_argscall_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_warpsnum_stages 的逻辑扩展为 implicit_constants 集合(包含 num_warpsnum_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 二进制,无需由调用方传入。

likedislike
不准确?
AAllenGuan成员
8月4日 关联了看板:FrameworkPTAdapter 版本issue看板
atomgit-bot
atomgit-bot
8月4日 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
此处折叠了94条消息 查看更多
htchu成员
20 天前 评论:

/approve

likedislike
ascend-robotascend-robot成员
20 天前 添加了label:approvedlgtm
ascend-robotascend-robot成员
20 天前 关闭了关联的issue
ascend-robotascend-robot成员
20 天前 合入了pull request
ascend-robot
ascend-robot成员
20 天前 评论:
流水线 pytorch_gitcode_PR_multiVersion#14483 [ commitID:fad5a366 ] 已完成
likedislike