已合并
sdpa use component impl when the last dim of input is greater than 256 #5578
XianglongZeng创建于 24 天前
sdpa use component impl when the last dim of input is greater than 256 #5578
已合并
Pull Request已成功合入, 合并人@ascend-robot
（感谢 XianglongZeng 的贡献）24 天前 创建了 pull request，commit 38483749
atomgit-bot
24 天前 评论：
24 天前 评论：
变更摘要
此 PR 主要针对 scaled_dot_product_attention 算子，在输入张量最后一维（尾轴）大于 256 的场景下，当用户通过 torch_npu.npu.use_compatible_impl(True) 启用兼容实现模式时，强制走 math 路径（force_math = true），以调用与 torch_gpu 一致的小算子实现，从而保证高精度；若不启用该模式则保持原有融合大算子实现以优先性能。
主要改动
ScaledDotProductAttentionKernelNpuOpApi.cpp新增尾维大于 256 时的force_math判断: 在两处force_math赋值逻辑中，新增对query.size(-1) > 256 || key.size(-1) > 256 || value.size(-1) > 256的检查，满足条件时强制将force_math设为true，使算子退避到 component/math 实现路径。- 新增测试用例
test_sdpa_last_dim_greater_than_256: 在test_scaled_dot_product_attention.py中增加了尾维为 257 的测试场景，先调用use_compatible_impl(True)启用兼容模式，然后对比 CPU（fp32）与 NPU（fp16）输出精度，验证大尾维场景下兼容实现的正确性。
ascend-robot
24 天前 评论：
24 天前 评论：
atomgit-bot
24 天前 评论：
24 天前 评论：
24 天前 修改了pull request 的描述
此处折叠了55条消息 查看更多
chengpeng25
21 天前 评论：
21 天前 评论：
/lgtm
21 天前 添加了label：lgtm
21 天前 关闭了关联的issue
21 天前 合入了pull request
ascend-robot
21 天前 评论：
21 天前 评论：
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
【合入来源】
【修改方案】
在尾轴大于256的情况下，若用户对精度要求高，export TORCH_NPU_USE_COMPATIBLE_IMPL=1，牺牲一定性能和显存，内部实现和torch_gpu走一致实现，调小算子。若用户对性能要求高，则保持现状，调融合大算子。
【资料变更】
不涉及
【接口变更】
不涉及
【功能验证】
【CheckList】