Star1205
2811
代码
代码
Issues225
Pull Requests452
Actions
项目讨论
Wiki
分析
Star1205
2811
已合并
add fused operator finalRouting #5941
slj创建于 5月26日
add fused operator finalRouting #5941
已合并
slj创建于 5月26日
从slj02/ops-transformerfeat/FinalRouting合入到cann/ops-transformermaster
共 4 个文件变更+428-0
Aexperimental/moe/final_routing/CMakeLists.txt+18-0
@@ -0,0 +1,18 @@
1+message(STATUS "BUILD_TORCH_OPS ON in final_routing")
2+# FINAL_ROUTING operation sources
3+file(GLOB FINAL_ROUTING_NPU_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")
4+ 
5+set(FINAL_ROUTING_SOURCES ${FINAL_ROUTING_NPU_SOURCES})
6+# Mark .cpp files with special properties
7+set_source_files_properties(
8+ ${FINAL_ROUTING_NPU_SOURCES} PROPERTIES
9+ LANGUAGE CXX
10+ COMPILE_FLAGS "--cce-soc-version=Ascend910B1 --cce-soc-core-type=VecCore --cce-auto-sync -xcce"
11+)
12+ 
13+# Create object library
14+add_library(final_routing_objects OBJECT ${FINAL_ROUTING_SOURCES})
15+ 
16+target_compile_options(final_routing_objects PRIVATE ${COMMON_COMPILE_OPTIONS})
17+target_include_directories(final_routing_objects PRIVATE ${COMMON_INCLUDE_DIRS})
18+return()
Aexperimental/moe/final_routing/README.md+85-0
@@ -0,0 +1,85 @@
1+# 算子名称：FinalRouting
2+ 
3+## 产品支持情况
4+ 
5+| 产品 | 是否支持 |
6+| :----------------------------------------------------------- | :------: |
7+| Atlas A2 训练系列产品 | 是 |
8+ 
9+## 功能说明
10+ 
11+- 算子功能：用于MoE模型的最终路由阶段，将各个Expert计算的结果按照评分加权合并，得到每个Token的最终输出，完成MoE的Combine阶段。
12+ 
13+- 计算公式：
14+对于每个Token `t`：
15+ 
16+$$
17+out[t] = \sum_{e} \big(in[token\_table[t, e]] \cdot score\_table[t, e]\big)
18+$$
19+ 
20+其中，仅当 `token_table[t, e] >= 0` 时参与计算。
21+ 
22+## 参数说明
23+ 
24+<table style="undefined;table-layout: fixed; width: 820px"><colgroup>
25+ <col style="width: 100px">
26+ <col style="width: 150px">
27+ <col style="width: 190px">
28+ <col style="width: 260px">
29+ <col style="width: 120px">
30+ </colgroup>
31+ <thead>
32+ <tr>
33+ <th>参数名</th>
34+ <th>输入/输出/属性</th>
35+ <th>描述</th>
36+ <th>数据类型</th>
37+ <th>数据格式</th>
38+ </tr></thead>
39+ <tbody>
40+ <tr>
41+ <td>blockDim</td>
42+ <td>输入</td>
43+ <td>AI CORE的数量，比如：Ascend910B是40。</td>
44+ <td>int64_t</td>
45+ <td>-</td>
46+ </tr>
47+ <tr>
48+ <td>in</td>
49+ <td>输入</td>
50+ <td>expert输出张量， shape为(expert_num*token_num, hidden_size)</td>
51+ <td>BFLOAT16</td>
52+ <td>ND</td>
53+ </tr>
54+ <tr>
55+ <td>token_table</td>
56+ <td>输入</td>
57+ <td>token到expert的映射表， shape为(token_num, expert_num)</td>
58+ <td>int32_t</td>
59+ <td>ND</td>
60+ </tr>
61+ <tr>
62+ <td>score_table</td>
63+ <td>输入</td>
64+ <td>每个token在每个expert的评分，shape为(token_num, expert_num)</td>
65+ <td>BFLOAT16</td>
66+ <td>ND</td>
67+ </tr>
68+ <tr>
69+ <td>out</td>
70+ <td>输出</td>
71+ <td>加权合并后的输出张量，shape为(token_num, hidden_size)</td>
72+ <td>BFLOAT16</td>
73+ <td>ND</td>
74+ </tr>
75+ </tbody></table>
76+ 
77+## 约束说明
78+ 
79+- token_table中小于0的值表示该expert对此token无效
80+ 
81+## 调用说明
82+ 
83+```
84+torch.ops.npu_ops_transformer_ext.final_routing(block_dim, input, token_table,score_table, output)
85+```
Aexperimental/moe/final_routing/final_routing.cpp+279-0
@@ -0,0 +1,279 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <ATen/Operators.h>
12+#include <torch/all.h>
13+#include <torch/library.h>
14+#include "acl/acl.h"
15+#include "torch_npu/csrc/core/npu/NPUStream.h"
16+#include "torch_npu/csrc/core/npu/DeviceUtils.h"
17+#include "torch_npu/csrc/framework/OpCommand.h"
18+ 
19+namespace npu_ops_transformer_ext {
20+namespace FinalRouting {
21+#include <iostream>
22+#include <stdio.h>
23+#include "kernel_operator.h"
24+ 
25+using namespace AscendC;
26+ 
27+constexpr int64_t UB_MAX_BYTES = 184*1024;
28+constexpr int64_t BUFFER_NUM = 1;
29+ 
30+class final_routing {
31+public:
32+ __aicore__ inline final_routing() {}
33+ 
34+ __aicore__ inline void Init(GM_ADDR in, GM_ADDR token_table, GM_ADDR score_table, GM_ADDR out,
35+ const int64_t expert_num, const int64_t token_num, const int64_t copy_byte)
36+ {
37+ blockNum_ = GetBlockNum();
38+ blockIdx_ = GetBlockIdx();
39+ 
40+ gbExpertNum_ = expert_num;
41+ gbTokenNum_ = token_num;
42+ gbCopyByte_ = copy_byte;
43+ gbCopyNum_ = gbCopyByte_ / sizeof(bfloat16_t);
44+ 
45+ gbTokenNumAlign_ = AlignUp(gbTokenNum_, 64 / sizeof(int32_t));
46+ gbExpertNumAlign_ = AlignUp(gbExpertNum_, 64 / sizeof(int32_t));
47+ 
48+ bkExpertNum_ = 1;
49+ bkTokenNum_ = 1;
50+ bkCopyByte_ = gbCopyByte_;
51+ bkCopyNum_ = bkCopyByte_ / sizeof(bfloat16_t);
52+ bkCopyByteAlign_ = AlignUp(bkCopyByte_, 64);
53+ bkCopyNumAlign_ = bkCopyByteAlign_ / sizeof(bfloat16_t);
54+
55+ gmIn_.SetGlobalBuffer((__gm__ bfloat16_t*)(in), gbExpertNum_ * gbTokenNum_ * gbCopyNum_);
56+ gmTokenTable_.SetGlobalBuffer((__gm__ int32_t*)(token_table), gbTokenNum_ * gbExpertNum_);
57+ gmScoreTable_.SetGlobalBuffer((__gm__ bfloat16_t*)(score_table), gbTokenNum_ * gbExpertNum_);
58+ gmOut_.SetGlobalBuffer((__gm__ bfloat16_t*)(out), gbTokenNum_ * gbCopyNum_);
59+ 
60+ pipe_.InitBuffer(inQueIn_, BUFFER_NUM, bkCopyByteAlign_ * 2);
61+ pipe_.InitBuffer(inQueTokenTable_, BUFFER_NUM, gbExpertNumAlign_ * sizeof(int32_t));
62+ pipe_.InitBuffer(inQueScoreTable_, BUFFER_NUM, gbExpertNumAlign_ * sizeof(bfloat16_t));
63+ pipe_.InitBuffer(outQueOut_, BUFFER_NUM, bkCopyByteAlign_ * 2);
64+ }
65+ 
66+ __aicore__ inline void Process()
67+ {
68+ for (int64_t tis = 0; tis < gbTokenNum_; tis += blockNum_) {
69+ int64_t token_idx = tis + blockIdx_;
70+ 
71+ if (token_idx < gbTokenNum_) {
72+ LocalTensor<int32_t> local_token_table = inQueTokenTable_.AllocTensor<int32_t>();
73+ LocalTensor<bfloat16_t> local_score_table = inQueScoreTable_.AllocTensor<bfloat16_t>();
74+
75+ DataCopyParams copy_pas{1, (uint16_t)(gbExpertNum_ * sizeof(int32_t)), 0, 0};
76+ DataCopyPadParams pad_pas;
77+ 
78+ DataCopyPad(local_token_table, gmTokenTable_[token_idx * gbExpertNum_], copy_pas, pad_pas);
79+ copy_pas.blockLen = (uint16_t)(gbExpertNum_ * sizeof(bfloat16_t));
80+ DataCopyPad(local_score_table, gmScoreTable_[token_idx * gbExpertNum_], copy_pas, pad_pas);
81+
82+ inQueTokenTable_.EnQue(local_token_table);
83+ inQueScoreTable_.EnQue(local_score_table);
84+
85+ lmTokenTable_ = inQueTokenTable_.DeQue<int32_t>();
86+ lmScoreTable_ = inQueScoreTable_.DeQue<bfloat16_t>();
87+
88+ lmOut_ = outQueOut_.AllocTensor<bfloat16_t>();
89+ lmOutFloat_ = lmOut_.ReinterpretCast<float>();
90+ Duplicate(lmOutFloat_, (float)0.0, bkCopyNumAlign_);
91+
92+ for (int64_t expert_idx = 0; expert_idx < gbExpertNum_; expert_idx++) {
93+ int64_t routing_token_idx = lmTokenTable_.GetValue(expert_idx);
94+ if (routing_token_idx >= 0) {
95+ CopyIn(routing_token_idx);
96+ Compute(AscendC::ToFloat(lmScoreTable_.GetValue(expert_idx)));
97+ }
98+ }
99+
100+ Cast(lmOut_, lmOutFloat_, RoundMode::CAST_ROUND, bkCopyNum_);
101+ outQueOut_.EnQue(lmOut_);
102+ outQueOut_.DeQue<bfloat16_t>();
103+ CopyOut(token_idx);
104+
105+ inQueTokenTable_.FreeTensor(local_token_table);
106+ inQueScoreTable_.FreeTensor(local_score_table);
107+ outQueOut_.FreeTensor(lmOut_);
108+ }
109+ }
110+ }
111+ 
112+private:
113+ __aicore__ inline void CopyIn(int64_t routing_token_idx)
114+ {
115+ LocalTensor<bfloat16_t> local_in = inQueIn_.AllocTensor<bfloat16_t>();
116+
117+ DataCopyExtParams copy_pas{1, (uint32_t)(bkCopyByte_), 0, 0, 0};
118+ DataCopyPadExtParams<bfloat16_t> pad_pas;
119+
120+ int64_t offset = routing_token_idx * gbCopyNum_;
121+ DataCopyPad(local_in[bkCopyNumAlign_], gmIn_[offset], copy_pas, pad_pas);
122+ inQueIn_.EnQue(local_in);
123+ }
124+ 
125+ __aicore__ inline void Compute(float score)
126+ {
127+ LocalTensor<bfloat16_t> local_in = inQueIn_.DeQue<bfloat16_t>();
128+ LocalTensor<float> local_in_float = local_in.ReinterpretCast<float>();
129+
130+ Cast(local_in_float, local_in[bkCopyNumAlign_], RoundMode::CAST_NONE, bkCopyNum_);
131+ 
132+ Muls(local_in_float, local_in_float, score, bkCopyNum_);
133+ Add(lmOutFloat_, lmOutFloat_, local_in_float, bkCopyNum_);
134+ 
135+ inQueIn_.FreeTensor(local_in);
136+ }
137+ 
138+ __aicore__ inline void CopyOut(int64_t token_idx)
139+ {
140+ DataCopyExtParams copy_pas{1, (uint32_t)(bkCopyByte_), 0, 0, 0};
141+ DataCopyPad(gmOut_[token_idx * bkCopyNum_], lmOut_, copy_pas);
142+ }
143+ 
144+private:
145+ TPipe pipe_;
146+ TQue<QuePosition::VECIN, BUFFER_NUM> inQueIn_, inQueTokenTable_, inQueScoreTable_;
147+ TQue<QuePosition::VECOUT, BUFFER_NUM> outQueOut_;
148+ 
149+ GlobalTensor<bfloat16_t> gmIn_, gmOut_;
150+ GlobalTensor<int32_t> gmTokenTable_;
151+ GlobalTensor<bfloat16_t> gmScoreTable_;
152+
153+ LocalTensor<int32_t> lmTokenTable_;
154+ LocalTensor<bfloat16_t> lmScoreTable_;
155+ LocalTensor<bfloat16_t> lmOut_;
156+ LocalTensor<float> lmOutFloat_;
157+ 
158+ int64_t blockNum_, blockIdx_;
159+
160+ int64_t gbExpertNum_, gbTokenNum_, gbCopyByte_, gbCopyNum_;
161+ int64_t gbExpertNumAlign_, gbTokenNumAlign_;
162+ 
163+ int64_t bkExpertNum_, bkTokenNum_, bkCopyByte_, bkCopyNum_;
164+ int64_t bkCopyByteAlign_, bkCopyNumAlign_;
165+};
166+ 
167+extern "C" __global__ __aicore__ void compute_final_routing(
168+ GM_ADDR in, GM_ADDR token_table, GM_ADDR score_table, GM_ADDR out,
169+ const int64_t expert_num, const int64_t token_num, const int64_t copy_byte)
170+{
171+ final_routing op;
172+ op.Init(in, token_table, score_table, out, expert_num, token_num, copy_byte);
173+ op.Process();
174+}
175+ 
176+void final_routing_kernel_lanuch(int64_t block_dim, void* stream,
177+ uint8_t* in, uint8_t* token_table, uint8_t* score_table, uint8_t* out,
178+ const int64_t expert_num, const int64_t token_num, const int64_t copy_byte)
179+{
180+ compute_final_routing<<<block_dim, nullptr, stream>>>(in, token_table, score_table, out,
181+ expert_num, token_num, copy_byte);
182+}
183+ 
184+inline int64_t align_up(const int64_t number, const int64_t alignSize)
185+{
186+ if (number % alignSize == 0) {
187+ return number;
188+ }
189+
190+ return ((number / alignSize + 1) * alignSize);
191+}
192+ 
193+ 
194+int judge_final_routing_lanuch(const int64_t expert_num, const int64_t token_num, const int64_t copy_byte,
195+ const int64_t ubSize, const int64_t vCores)
196+{
197+ (void)(vCores);
198+ 
199+ constexpr int64_t BUFFER_NUM = 1;
200+
201+ int64_t gbExpertNum_;
202+ int64_t gbTokenNum_;
203+ int64_t gbCopyByte_;
204+ int64_t gbExpertNumAlign_;
205+ int64_t gbTokenNumAlign_;
206+ 
207+ int64_t bkCopyByte_;
208+ int64_t bkCopyByteAlign_;
209+ 
210+ gbExpertNum_ = expert_num;
211+ gbTokenNum_ = token_num;
212+ gbCopyByte_ = copy_byte;
213+ 
214+ gbTokenNumAlign_ = align_up(gbTokenNum_, 64 / sizeof(int32_t));
215+ gbExpertNumAlign_ = align_up(gbExpertNum_, 64 / sizeof(int32_t));
216+ 
217+ bkCopyByte_ = gbCopyByte_;
218+ bkCopyByteAlign_ = align_up(bkCopyByte_, 64);
219+
220+ float use_byte = BUFFER_NUM * bkCopyByteAlign_ * 2;
221+ use_byte += gbTokenNumAlign_ * sizeof(int32_t);
222+ use_byte += gbExpertNumAlign_ * sizeof(int32_t);
223+ use_byte += gbExpertNumAlign_ * sizeof(int16_t);
224+ 
225+ if (use_byte > ubSize) {
226+ std::cout << __FUNCTION__ << ": " << use_byte/1024 << " KB" << std::endl;
227+ return 1;
228+ }
229+ return 0;
230+}
231+ 
232+int final_routing_lanuch(int64_t block_dim, void* stream,
233+ uint8_t* in, uint8_t* token_table, uint8_t* score_table, uint8_t* out,
234+ const int64_t expert_num, const int64_t token_num, const int64_t copy_byte)
235+{
236+ int64_t ubSize = 184 * 1024;
237+ int64_t vCores = 40;
238+
239+ int ret = judge_final_routing_lanuch(expert_num, token_num, copy_byte, ubSize, vCores);
240+ if (ret == 0) {
241+ final_routing_kernel_lanuch(block_dim, stream, in, token_table, score_table,
242+ out, expert_num, token_num, copy_byte);
243+ return 0;
244+ }
245+
246+ std::cout << __FUNCTION__ << ": " << "UB size is limited, please check!" << std::endl;
247+ return 1;
248+}
249+ 
250+int64_t final_routing_npu(int64_t block_dim, torch::Tensor &in, torch::Tensor &token_table,
251+ torch::Tensor &score_table, torch::Tensor &out)
252+{
253+ TORCH_CHECK(torch_npu::utils::is_npu(in), "input tensor must be on NPU device");
254+ TORCH_CHECK(torch_npu::utils::is_npu(token_table), "token table tensor must be on NPU device");
255+ TORCH_CHECK(torch_npu::utils::is_npu(score_table), "score table tensor must be on NPU device");
256+ TORCH_CHECK(torch_npu::utils::is_npu(out), "output tensor must be on NPU device");
257+ 
258+ const int64_t token_num = token_table.size(0);
259+ const int64_t expert_num = token_table.size(1);
260+ const int64_t copy_byte = in.element_size() * in.size(1);
261+
262+ auto stream = c10_npu::getCurrentNPUStream().stream(false);
263+ int launchStatus = 0;
264+ auto acl_call = [=, &launchStatus]() -> int {
265+ launchStatus = final_routing_lanuch(block_dim, stream, (uint8_t *)in.data_ptr(),
266+ (uint8_t *)token_table.data_ptr(), (uint8_t *)score_table.data_ptr(), (uint8_t *)out.data_ptr(),
267+ expert_num, token_num, copy_byte);
268+ return 0;
269+ };
270+ at_npu::native::OpCommand::RunOpApi("FinalRouting", acl_call);
271+ return launchStatus;
272+}
273+ 
274+TORCH_LIBRARY_IMPL(npu_ops_transformer_ext, PrivateUse1, m)
275+{
276+ m.impl("final_routing", final_routing_npu);
277+}
278+}
279+}
Aexperimental/moe/final_routing/test/test_final_routing.py+46-0
@@ -0,0 +1,46 @@
1+#!/usr/bin/env python3
2+# -*- coding: UTF-8 -*-
3+# -----------------------------------------------------------------------------------------------------------
4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# -----------------------------------------------------------------------------------------------------------
12+ 
13+import logging
14+import torch
15+import torch_npu
16+import npu_ops_transformer_ext
17+ 
18+logging.basicConfig(level=logging.INFO)
19+ 
20+TOKEN_NUM = 2
21+EXPERT_NUM = 4
22+HIDDEN_SIZE = 8
23+BLOCK_DIM = 40
24+ 
25+input_tensor = torch.randint(1, 10, (TOKEN_NUM * EXPERT_NUM, HIDDEN_SIZE)).to(torch.bfloat16)
26+token_table = torch.arange(TOKEN_NUM * EXPERT_NUM, dtype=torch.int32).reshape(TOKEN_NUM, EXPERT_NUM)
27+score_table = torch.randint(0, 10, (TOKEN_NUM, EXPERT_NUM)).to(torch.bfloat16)
28+output = torch.empty(TOKEN_NUM, HIDDEN_SIZE, dtype=torch.bfloat16)
29+output_cpu = torch.empty_like(output).float()
30+ 
31+input_32 = input_tensor.float()
32+for t in range(TOKEN_NUM):
33+ for e in range(EXPERT_NUM):
34+ input_idx = token_table[t, e].item()
35+ if input_idx < 0:
36+ continue
37+ score = score_table[t, e].item()
38+ feat = input_32[input_idx]
39+ output_cpu[t] += feat * score
40+ 
41+output_cpu = output_cpu.to(torch.bfloat16)
42+output_npu = output.npu()
43+torch.ops.npu_ops_transformer_ext.final_routing(BLOCK_DIM, input_tensor.npu(), token_table.npu(),
44+ score_table.npu(), output_npu)
45+ 
46+logging.info(f"cpu result vs npu result: {torch.equal(output_cpu, output_npu.cpu())}")