已合并
[feature]新增RainFusionAttention算子的plugin #54
mazhixin00_00创建于 2025年12月11日
[feature]新增RainFusionAttention算子的plugin #54
已合并
mazhixin00_00创建于 2025年12月11日
mazhixin00_00成员
2025年12月11日

[feature]新增RainFusionAttention算子的plugin
背景:
CANN中已有稀疏算子RainFusionAttention,但是未提供python的调用方式。

方案:
MindIE-SD通过注册torch算子的方式,提供算子调用方式

m.def("rainfusionattention_mindie_sd(Tensor query, Tensor key, Tensor value, Tensor select_idx, \
        Tensor select_num_idx, int[] blockshape, Tensor? attn_mask=None, int[]? actual_seq_qlen=None, \
        int[]? actual_seq_kvlen=None, Tensor? block_table=None, str q_input_layout='TND', str kv_input_layout='TND', \
        int head_num=1, int mask_type=0, float scale=1.0, \
        int inner_precise=1, int block_size=0) -> (Tensor, Tensor)");

测试结果:
全量UT:
image.png
约束:

  • query,key和value目前只支持shape为(batch*seqlen, head, headdim)的tensor
  • q_input_layout和kv_input_layout目前只支持TND
  • 该接口依赖的CANN包暂未发布,暂不支持外部调用
likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 mazhixin00_00 的贡献)
Mmazhixin00_00成员
2025年12月11日 创建了 pull request,commit 23764c76
ascend-robot
ascend-robot成员
2025年12月11日 评论:

Thank your for your pull-request.

The full list of commands accepted by me can be found at here.

You can get sig-info at here

likedislike
ascend-robot
ascend-robot成员
2025年12月11日 评论:

以下是根据您提交的修改文件推荐的Reviewer和Committer序列,需各模块评审通过后方可合入

Module List Reviewers Committers
repo-Ascend/MindIE-SD yyansifu, wenjiali_sh, zhou-wenxue, shaopeng666, muziyuhui666 betta18, guowenna1, lanwangli, blian, mazhixin00_00
likedislike
ascend-robotascend-robot成员
2025年12月11日 添加了label:ascend-cla/yes
mazhixin00_00成员
2025年12月11日 评论:

compile

likedislike
此处折叠了79条消息 查看更多
lanwangli成员
2025年12月15日 评论:

/merge

likedislike
ascend-robotascend-robot成员
2025年12月15日 添加了label:keeper_approved
ascend-robot
ascend-robot成员
2025年12月15日 评论:

Review Guide

This Pull-Request Passes Review.
Committers who writed a comment of /approve are: betta18.
Reviewers who writed a comment of /lgtm are: guowenna1, lanwangli, changetheway, betta18.

likedislike
ascend-robotascend-robot成员
2025年12月15日 合入了pull request
Mmazhixin00_00成员
2025年12月26日 修改了pull request 的描述