已合并
fix(triton_experimental): propagate NaN in min/max reductions #44773
伦创建于 17 天前
fix(triton_experimental): propagate NaN in min/max reductions #44773
已合并
Pull Request已成功合入, 合并人@ascend-robot
（感谢 伦 的贡献）
atomgit-bot
17 天前 评论：
17 天前 评论：
变更摘要
该 PR 旨在修复 Triton 实验性代码生成（triton_experimental）中 min/max 归约操作对 NaN 的传播问题，确保归约结果在输入含 NaN 时能正确返回 NaN。核心改动包括：在 NPUTritonKernelOverrides.minimum/maximum 生成的 tl.minimum/tl.maximum 指令中显式传入 tl.PropagateNan.ALL；同时在 npu_triton_helpers.py 中新增带 propagate_nan 语义的归约辅助函数，并调整 min2 的实现方式。
主要改动
triton.py中minimum/maximum的 NaN 传播:NPUTritonKernelOverrides.minimum与maximum生成的代码由tl.minimum({a}, {b})/tl.maximum({a}, {b})改为追加tl.PropagateNan.ALL参数，使元素级 min/max 运算在任一输入为 NaN 时返回 NaN。- 新增
_min_prop_nan归约内核: 在npu_triton_helpers.py中新增@triton.jit函数_min_prop_nan，内部通过core.minimum(a, b, propagate_nan=core.PropagateNan.ALL)实现 NaN 传播的最小值比较。 min2改用tl.reduce并引入core模块:min2由原来的tl.min(a, dim, propagate_nan=True)改为tl.reduce(a, dim, _min_prop_nan)，并新增from triton.language import core导入以支撑新实现。
ascend-robot
17 天前 评论：
17 天前 评论：
atomgit-bot
17 天前 评论：
17 天前 评论：
17 天前 添加了label：ascend-cla/yes
此处折叠了183条消息 查看更多
3 天前 添加了label：lgtm
3 天前 关闭了关联的issue
3 天前 合入了pull request
ascend-robot
3 天前 评论：
3 天前 评论：
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
AtlasAccount
3 天前 评论：
3 天前 评论：
流水线 pytorch_gitcode_PR_multiVersion#14800 [ commitID：dadfddd7 ] 运行失败
【合入来源】
【修改方案】
修复 triton_experimental 后端 min/max 归约的 NaN 传播语义错误。
codegen/triton.py：元素级 minimum/maximum 从裸tl.minimum(a, b)改为tl.minimum(a, b, tl.PropagateNan.ALL)（maximum 同理）。npu_triton_helpers.py：min2/max2归约从tl.min改为tl.reduce(a, dim, _min_prop_nan)（NaN 传播合体形式，对齐 cuda 后端）。根因：裸
tl.minimum/maximum落arith.minnumf/maxnumf（NaN-dropping），NaN 会被丢弃；PropagateNan.ALL落arith.minimumf/maximumf（NaN-propagating）。后者既保证 NaN语义正确，又避开 minnumf 在向量管道约 5× 的指令开销，实测 1.6–2.8× 加速。【资料变更】
不涉及
【接口变更】
不涉及
【功能验证】
【CheckList】