[feature]新增RainFusionAttention算子的plugin-MindIE-SD-AtomGit

已合并

[feature]新增RainFusionAttention算子的plugin #54

mazhixin00_00创建于 2025年12月11日从mazhixin00_00/MindIE-SD-public： dev 合入到 Ascend/MindIE-SD： dev

[feature]新增RainFusionAttention算子的plugin #54

已合并

mazhixin00_00创建于 2025年12月11日

从mazhixin00_00/MindIE-SD-public dev合入到Ascend/MindIE-SD dev

共 4 个文件变更+111-1

Mcsrc/CMakeLists.txt+2-1

	@@ -52,7 +52,8 @@ add_library(PTAExtensionOPS SHARED
52	./plugin/batchmatmulv3duo.cpp	52	./plugin/batchmatmulv3duo.cpp
53	./plugin/adalayernorm.cpp	53	./plugin/adalayernorm.cpp
54	./plugin/find_op_path.cpp	54	./plugin/find_op_path.cpp
55	- ./plugin/la_preprocess.cpp)	55	+ ./plugin/la_preprocess.cpp
		56	+ ./plugin/rainfusionattention.cpp)
56		57
57	target_compile_features(PTAExtensionOPS PRIVATE cxx_std_17)	58	target_compile_features(PTAExtensionOPS PRIVATE cxx_std_17)
58	if(DEFINED ENV{USER_ABI_VERSION})	59	if(DEFINED ENV{USER_ABI_VERSION})

Acsrc/plugin/rainfusionattention.cpp+65-0

@@ -0,0 +1,65 @@
	1	+/**
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
	3	+ * MindIE is licensed under Mulan PSL v2.
	4	+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
	5	+ * You may obtain a copy of Mulan PSL v2 at:
	6	+ * http://license.coscl.org.cn/MulanPSL2
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
	8	+ * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
	9	+ * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
	10	+ * See the Mulan PSL v2 for more details.
	11	+ */
	12	+
	13	+#include <string_view>
	14	+#include <torch/library.h>
	15	+#include "torch_npu/csrc/framework/utils/OpAdapter.h"
	16	+#include "torch_npu/csrc/core/npu/NPUFormat.h"
	17	+#include "pytorch_npu_helper.h"
	18	+
	19	+#include "rainfusionattention.h"
	20	+
	21	+using namespace at;
	22	+using npu_preparation = at_npu::native::OpPreparation;
	23	+namespace {
	24	+constexpr int EXPECTED_TENSOR_DIMENSION = 3;
	25	+constexpr std::string_view RAINFUSIONATTENTIONOP_NAME = "aclnnRainFusionAttention";
	26	+}
	27	+std::tuple<at::Tensor, at::Tensor> rainfusionattention_mindie_sd_impl_npu(
	28	+ const at::Tensor &query,
	29	+ const at::Tensor &key,
	30	+ const at::Tensor &value,
	31	+ const at::Tensor &select_idx,
	32	+ const at::Tensor &select_num_idx,
	33	+ at::IntArrayRef blockshape,
	34	+ const c10::optional<at::Tensor> &attn_mask,
	35	+ c10::OptionalIntArrayRef actual_seq_qlen,
	36	+ c10::OptionalIntArrayRef actual_seq_kvlen,
	37	+ const c10::optional<at::Tensor> &block_table,
	38	+ std::string q_input_layout,
	39	+ std::string kv_input_layout,
	40	+ int64_t head_num, int64_t mask_type, double scale,
	41	+ int64_t inner_precise, int64_t block_size)
	42	+{
	43	+ TORCH_CHECK(query.dim() == EXPECTED_TENSOR_DIMENSION, "Query must be 3D tensor.");
	44	+ TORCH_CHECK(key.dim() == EXPECTED_TENSOR_DIMENSION, "Key must be 3D tensor.");
	45	+ TORCH_CHECK(value.dim() == EXPECTED_TENSOR_DIMENSION, "Value must be 3D tensor.");
	46	+ TORCH_CHECK(q_input_layout == "TND", "q_input_layout must be 'TND'.");
	47	+ TORCH_CHECK(kv_input_layout == "TND", "kv_input_layout must be 'TND'.");
	48	+ const at::Tensor& attenMask = c10::value_or_else(attn_mask, [] {return at::Tensor();});
	49	+ auto actualSeqLengths = actual_seq_qlen.value_or(at::IntArrayRef{});
	50	+ auto actualSeqLengthsKv = actual_seq_kvlen.value_or(at::IntArrayRef{});
	51	+ const at::Tensor& blockTable = c10::value_or_else(block_table, [] {return at::Tensor();});
	52	+
	53	+ const char* qlayoutPtr = q_input_layout.data();
	54	+ const char* kvlayoutPtr = kv_input_layout.data();
	55	+
	56	+ at::Tensor attentionOut = at_npu::native::empty_with_format(query.sizes(), query.options(),
	57	+ at_npu::native::get_npu_format(query));
	58	+ at::Tensor softmaxLse = at_npu::native::empty_with_format({query.sizes()[0], query.sizes()[1], query.sizes()[2]},
	59	+ query.options(), at_npu::native::get_npu_format(query));
	60	+
	61	+ EXEC_NPU_CMD<RAINFUSIONATTENTIONOP_NAME>(query, key, value, select_idx, select_num_idx, blockshape,
	62	+ attenMask, actualSeqLengths, actualSeqLengthsKv, blockTable, qlayoutPtr, kvlayoutPtr,
	63	+ head_num, mask_type, scale, inner_precise, block_size, attentionOut, softmaxLse);
	64	+ return std::tuple<at::Tensor, at::Tensor>(attentionOut, softmaxLse);
	65	+}

Acsrc/plugin/rainfusionattention.h+37-0

@@ -0,0 +1,37 @@
	1	+/**
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
	3	+ * MindIE is licensed under Mulan PSL v2.
	4	+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
	5	+ * You may obtain a copy of Mulan PSL v2 at:
	6	+ * http://license.coscl.org.cn/MulanPSL2
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
	8	+ * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
	9	+ * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
	10	+ * See the Mulan PSL v2 for more details.
	11	+ */
	12	+
	13	+#ifndef RAINFUSIONATTENTION_MINDIE_SD_IMPL_H
	14	+#define RAINFUSIONATTENTION_MINDIE_SD_IMPL_H
	15	+
	16	+#include <ATen/Tensor.h>
	17	+#include <c10/util/Optional.h>
	18	+#include <string>
	19	+#include <tuple>
	20	+
	21	+std::tuple<at::Tensor, at::Tensor> rainfusionattention_mindie_sd_impl_npu(
	22	+ const at::Tensor &query,
	23	+ const at::Tensor &key,
	24	+ const at::Tensor &value,
	25	+ const at::Tensor &select_idx,
	26	+ const at::Tensor &select_num_idx,
	27	+ at::IntArrayRef blockshape,
	28	+ const c10::optional<at::Tensor> &attn_mask,
	29	+ c10::OptionalIntArrayRef actual_seq_qlen,
	30	+ c10::OptionalIntArrayRef actual_seq_kvlen,
	31	+ const c10::optional<at::Tensor> &block_table,
	32	+ std::string q_input_layout,
	33	+ std::string kv_input_layout,
	34	+ int64_t head_num, int64_t mask_type, double scale,
	35	+ int64_t inner_precise, int64_t block_size);
	36	+
	37	+#endif // RAINFUSIONATTENTION_MINDIE_SD_IMPL_H

Mcsrc/plugin/register_ops.cpp+7-0

	@@ -20,6 +20,7 @@
20	#include "batchmatmulv2.h"	20	#include "batchmatmulv2.h"
21	#include "adalayernorm.h"	21	#include "adalayernorm.h"
22	#include "la_preprocess.h"	22	#include "la_preprocess.h"
		23	+#include "rainfusionattention.h"
23		24
24		25
25	TORCH_LIBRARY(mindie, m)	26	TORCH_LIBRARY(mindie, m)
	@@ -47,6 +48,11 @@ TORCH_LIBRARY(mindie, m)
47	-> Tensor");	48	-> Tensor");
48	m.def("la_preprocess_mindie_sd(Tensor query, Tensor key, Tensor value, int align_len=256) \	49	m.def("la_preprocess_mindie_sd(Tensor query, Tensor key, Tensor value, int align_len=256) \
49	-> (Tensor, Tensor, Tensor)");	50	-> (Tensor, Tensor, Tensor)");
		51	+ m.def("rainfusionattention_mindie_sd(Tensor query, Tensor key, Tensor value, Tensor select_idx, \
		52	+ Tensor select_num_idx, int[] blockshape, Tensor? attn_mask=None, int[]? actual_seq_qlen=None, \
		53	+ int[]? actual_seq_kvlen=None, Tensor? block_table=None, str q_input_layout='TND', str kv_input_layout='TND', \
		54	+ int head_num=1, int mask_type=0, float scale=1.0, \
		55	+ int inner_precise=1, int block_size=0) -> (Tensor, Tensor)");
50	}	56	}
51		57
52		58
	@@ -60,4 +66,5 @@ TORCH_LIBRARY_IMPL(mindie, PrivateUse1, m)
60	m.impl("batchmatmulv2_mindie_sd", &batchmatmulv2_mindie_sd_impl_npu);	66	m.impl("batchmatmulv2_mindie_sd", &batchmatmulv2_mindie_sd_impl_npu);
61	m.impl("adaln_mindie_sd", &adaln_mindie_sd_impl_npu);	67	m.impl("adaln_mindie_sd", &adaln_mindie_sd_impl_npu);
62	m.impl("la_preprocess_mindie_sd", &la_preprocess_mindie_sd_impl_npu);	68	m.impl("la_preprocess_mindie_sd", &la_preprocess_mindie_sd_impl_npu);
		69	+ m.impl("rainfusionattention_mindie_sd", &rainfusionattention_mindie_sd_impl_npu);
63	}	70	}