Star1662
3340
代码
代码
Issues456
Pull Requests713
Actions
项目讨论
Wiki
分析
Star1662
3340
已合并
refactor(moe_gating_top_k): hash模式拆分为独立kernel模板MoeGatingTopKHashRegbase #15062
CPZhou创建于 8 天前
refactor(moe_gating_top_k): hash模式拆分为独立kernel模板MoeGatingTopKHashRegbase #15062
已合并
CPZhou创建于 8 天前
从ZhouChuping/ops-transformermaster-topk-hash合入到cann/ops-transformermaster
共 3 个文件变更+412-341
Amoe/moe_gating_top_k/op_kernel/arch35/moe_gating_top_k_hash_regbase.h+232-0
@@ -0,0 +1,232 @@
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+/*!
12+ * \file moe_gating_top_k_hash_regbase.h
13+ * \brief Arch35 register-based kernel for MoE Gating TopK hash mode.
14+ */
15+#ifndef MOE_GATING_TOP_K_HASH_REGBASE_H
16+#define MOE_GATING_TOP_K_HASH_REGBASE_H
17+ 
18+#include "moe_gating_top_k_regbase.h"
19+ 
20+namespace MoeGatingTopK {
21+using namespace AscendC;
22+using Reg::RegTensor;
23+ 
24+template <typename T, typename U1, typename U2>
25+class MoeGatingTopKHashRegbase : public MoeGatingTopKRegbase<T> {
26+public:
27+ __aicore__ inline MoeGatingTopKHashRegbase(){};
28+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds, GM_ADDR tid2eid, GM_ADDR y,
29+ GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace,
30+ const MoeGatingTopKRegbaseTilingData* tilingData, TPipe* tPipe);
31+ __aicore__ inline void Process();
32+ 
33+private:
34+ using MoeGatingTopKRegbase<T>::CopyInBias;
35+ using MoeGatingTopKRegbase<T>::CopyInX;
36+ using MoeGatingTopKRegbase<T>::ComputeX;
37+ using MoeGatingTopKRegbase<T>::CopyOutXNorm;
38+ using MoeGatingTopKRegbase<T>::CopyOut;
39+ using MoeGatingTopKRegbase<T>::expertIdxOutQueue_;
40+ using MoeGatingTopKRegbase<T>::sortedInGroupBuf_;
41+ using MoeGatingTopKRegbase<T>::xSigmoidBuf_;
42+ using MoeGatingTopKRegbase<T>::yOutQueue_;
43+ using MoeGatingTopKRegbase<T>::k_;
44+ using MoeGatingTopKRegbase<T>::eps_;
45+ using MoeGatingTopKRegbase<T>::routedScalingFactor_;
46+ using MoeGatingTopKRegbase<T>::noNorm_;
47+ using MoeGatingTopKRegbase<T>::perGroupExpertCount_;
48+ using MoeGatingTopKRegbase<T>::perGroupExpertCountAlign_;
49+ using MoeGatingTopKRegbase<T>::curCoreRowCount_;
50+ using MoeGatingTopKRegbase<T>::tilingData_;
51+ using MoeGatingTopKRegbase<T>::blockIdx_;
52+ 
53+ __aicore__ inline void HashCompute(int64_t row);
54+ __aicore__ inline void HashGatherWithSmallKAlignE(LocalTensor<float> xSigmoidTensor,
55+ LocalTensor<int32_t> expertIdxTensor, LocalTensor<T> yTensor,
56+ uint32_t k, float eps, float routedScalingFactor);
57+ __aicore__ inline void HashGatherWithSmallKNotAlignE(LocalTensor<float> xSigmoidTensor,
58+ LocalTensor<int32_t> expertIdxTensor, LocalTensor<T> yTensor,
59+ uint32_t k, float eps, float routedScalingFactor,
60+ int32_t expertIdxPad, int32_t perGroupExpertCountAlign);
61+ 
62+private:
63+ GlobalTensor<U1> inputIdsGm_;
64+ GlobalTensor<U2> tid2eidGm_;
65+};
66+ 
67+template <typename T, typename U1, typename U2>
68+__aicore__ inline void MoeGatingTopKHashRegbase<T, U1, U2>::HashGatherWithSmallKAlignE(
69+ LocalTensor<float> xSigmoidTensor, LocalTensor<int32_t> expertIdxTensor, LocalTensor<T> yTensor, uint32_t k,
70+ float eps, float routedScalingFactor)
71+{
72+ __ubuf__ float* inputAddr = (__ubuf__ float*)xSigmoidTensor.GetPhyAddr();
73+ __ubuf__ uint32_t* expertIdxAddr = (__ubuf__ uint32_t*)expertIdxTensor.GetPhyAddr();
74+ __ubuf__ T* outputAddr = (__ubuf__ T*)yTensor.GetPhyAddr();
75+ 
76+ if (noNorm_) {
77+ __VEC_SCOPE__
78+ {
79+ RegTensor<uint32_t> vreg1;
80+ RegTensor<float> vreg2;
81+ 
82+ Reg::MaskReg preg0 = Reg::UpdateMask<uint32_t>(k);
83+ Reg::LoadAlign(vreg1, expertIdxAddr);
84+ Reg::Gather(vreg2, inputAddr, vreg1, preg0);
85+ Reg::Muls(vreg2, vreg2, routedScalingFactor, preg0);
86+ ops::StoreOneTensorForDtypeT<T>(outputAddr, vreg2, preg0, 0);
87+ }
88+ } else {
89+ __VEC_SCOPE__
90+ {
91+ RegTensor<uint32_t> vreg1;
92+ RegTensor<float> vreg2;
93+ RegTensor<float> vreg3;
94+ RegTensor<float> vreg4;
95+ 
96+ Reg::MaskReg preg0 = Reg::UpdateMask<uint32_t>(k);
97+ Reg::LoadAlign(vreg1, expertIdxAddr);
98+ Reg::Gather(vreg2, inputAddr, vreg1, preg0);
99+ Reg::Reduce<Reg::ReduceType::SUM>(vreg3, vreg2, preg0);
100+ Reg::Adds(vreg3, vreg3, eps, preg0);
101+ Reg::Duplicate(vreg4, vreg3, preg0);
102+ Reg::Div(vreg4, vreg2, vreg4, preg0);
103+ Reg::Muls(vreg4, vreg4, routedScalingFactor, preg0);
104+ ops::StoreOneTensorForDtypeT<T>(outputAddr, vreg4, preg0, 0);
105+ }
106+ }
107+}
108+ 
109+template <typename T, typename U1, typename U2>
110+__aicore__ inline void MoeGatingTopKHashRegbase<T, U1, U2>::HashGatherWithSmallKNotAlignE(
111+ LocalTensor<float> xSigmoidTensor, LocalTensor<int32_t> expertIdxTensor, LocalTensor<T> yTensor, uint32_t k,
112+ float eps, float routedScalingFactor, int32_t expertIdxPad, int32_t perGroupExpertCountAlign)
113+{
114+ __ubuf__ float* inputAddr = (__ubuf__ float*)xSigmoidTensor.GetPhyAddr();
115+ __ubuf__ uint32_t* expertIdxAddr = (__ubuf__ uint32_t*)expertIdxTensor.GetPhyAddr();
116+ __ubuf__ T* outputAddr = (__ubuf__ T*)yTensor.GetPhyAddr();
117+ 
118+ if (noNorm_) {
119+ __VEC_SCOPE__
120+ {
121+ RegTensor<uint32_t> vreg1;
122+ RegTensor<float> vreg2;
123+ RegTensor<uint32_t> vregAlign;
124+ 
125+ Reg::MaskReg preg0 = Reg::UpdateMask<float>(k);
126+ Reg::Duplicate(vregAlign, perGroupExpertCountAlign, preg0);
127+ 
128+ Reg::LoadAlign(vreg1, expertIdxAddr);
129+ Reg::Gather(vreg2, inputAddr, vreg1, preg0);
130+ Reg::Muls(vreg2, vreg2, routedScalingFactor, preg0);
131+ 
132+ Reg::Div(vregAlign, vreg1, vregAlign, preg0);
133+ Reg::Muls(vregAlign, vregAlign, expertIdxPad, preg0);
134+ Reg::Sub(vreg1, vreg1, vregAlign, preg0);
135+ 
136+ ops::StoreOneTensorForDtypeT<T>(outputAddr, vreg2, preg0, 0);
137+ }
138+ } else {
139+ __VEC_SCOPE__
140+ {
141+ RegTensor<uint32_t> vreg1;
142+ RegTensor<float> vreg2;
143+ RegTensor<float> vreg3;
144+ RegTensor<float> vreg4;
145+ RegTensor<uint32_t> vregAlign;
146+ 
147+ Reg::MaskReg preg0 = Reg::UpdateMask<float>(k);
148+ Reg::Duplicate(vregAlign, perGroupExpertCountAlign, preg0);
149+ 
150+ Reg::LoadAlign(vreg1, expertIdxAddr);
151+ Reg::Gather(vreg2, inputAddr, vreg1, preg0);
152+ Reg::Reduce<Reg::ReduceType::SUM>(vreg3, vreg2, preg0);
153+ Reg::Adds(vreg3, vreg3, eps, preg0);
154+ Reg::Duplicate(vreg4, vreg3, preg0);
155+ Reg::Div(vreg4, vreg2, vreg4, preg0);
156+ Reg::Muls(vreg4, vreg4, routedScalingFactor, preg0);
157+ 
158+ Reg::Div(vregAlign, vreg1, vregAlign, preg0);
159+ Reg::Muls(vregAlign, vregAlign, expertIdxPad, preg0);
160+ Reg::Sub(vreg1, vreg1, vregAlign, preg0);
161+ 
162+ ops::StoreOneTensorForDtypeT<T>(outputAddr, vreg4, preg0, 0);
163+ }
164+ }
165+}
166+ 
167+template <typename T, typename U1, typename U2>
168+__aicore__ inline void MoeGatingTopKHashRegbase<T, U1, U2>::HashCompute(int64_t row)
169+{
170+ LocalTensor<int32_t> expertIdxTensor = expertIdxOutQueue_.template AllocTensor<int32_t>();
171+ LocalTensor<U2> hashExpertId = sortedInGroupBuf_.template Get<U2>();
172+ LocalTensor<int32_t> hashExpertIdInt32 = hashExpertId.template ReinterpretCast<int32_t>();
173+ 
174+ U1 key = inputIdsGm_.GetValue(row);
175+ SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
176+ DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(k_ * sizeof(U2)), 0, 0, 0};
177+ DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast<U2>(0)};
178+ DataCopyPad(hashExpertId, tid2eidGm_[key * k_], dataCopyParams, dataCopyPadParams);
179+ SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
180+ if constexpr (IsSameType<U2, int32_t>::value) {
181+ DataCopy(expertIdxTensor, hashExpertId, Align(k_, sizeof(int32_t)));
182+ } else {
183+ Cast(hashExpertIdInt32, hashExpertId, RoundMode::CAST_NONE, Align(k_, sizeof(U2)));
184+ PipeBarrier<PIPE_V>();
185+ DataCopy(expertIdxTensor, hashExpertIdInt32, Align(k_, sizeof(int32_t)));
186+ }
187+ PipeBarrier<PIPE_V>();
188+ 
189+ LocalTensor<float> xSigmoidTensor = xSigmoidBuf_.template Get<float>();
190+ LocalTensor<T> yTensor = yOutQueue_.template AllocTensor<T>();
191+ int32_t expertIdxPad = perGroupExpertCountAlign_ - perGroupExpertCount_;
192+ if (expertIdxPad != 0) {
193+ HashGatherWithSmallKNotAlignE(xSigmoidTensor, expertIdxTensor, yTensor, k_, eps_, routedScalingFactor_,
194+ expertIdxPad, perGroupExpertCountAlign_);
195+ } else {
196+ HashGatherWithSmallKAlignE(xSigmoidTensor, expertIdxTensor, yTensor, k_, eps_, routedScalingFactor_);
197+ }
198+ yOutQueue_.EnQue(yTensor);
199+ expertIdxOutQueue_.template EnQue<int32_t>(expertIdxTensor);
200+}
201+ 
202+template <typename T, typename U1, typename U2>
203+__aicore__ inline void MoeGatingTopKHashRegbase<T, U1, U2>::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds,
204+ GM_ADDR tid2eid, GM_ADDR y, GM_ADDR expertIdx,
205+ GM_ADDR out, GM_ADDR workspace,
206+ const MoeGatingTopKRegbaseTilingData* tilingData,
207+ TPipe* tPipe)
208+{
209+ MoeGatingTopKRegbase<T>::Init(x, bias, y, expertIdx, out, workspace, tilingData, tPipe);
210+ inputIdsGm_.SetGlobalBuffer((__gm__ U1*)inputIds);
211+ tid2eidGm_.SetGlobalBuffer((__gm__ U2*)tid2eid);
212+}
213+ 
214+template <typename T, typename U1, typename U2>
215+__aicore__ inline void MoeGatingTopKHashRegbase<T, U1, U2>::Process()
216+{
217+ CopyInBias();
218+ CopyInX(0);
219+ for (int64_t row = 1; row < curCoreRowCount_; row++) {
220+ ComputeX();
221+ CopyOutXNorm(row - 1);
222+ CopyInX(row);
223+ HashCompute(row - 1 + tilingData_->perCoreRowCount * blockIdx_);
224+ CopyOut(row - 1);
225+ }
226+ ComputeX();
227+ CopyOutXNorm(curCoreRowCount_ - 1);
228+ HashCompute(curCoreRowCount_ - 1 + tilingData_->perCoreRowCount * blockIdx_);
229+ CopyOut(curCoreRowCount_ - 1);
230+}
231+} // namespace MoeGatingTopK
232+#endif // MOE_GATING_TOP_K_HASH_REGBASE_H
Mmoe/moe_gating_top_k/op_kernel/moe_gating_top_k_apt.cpp+9-4
@@ -14,6 +14,7 @@
14 */14 */
15 15 
16#include "arch35/moe_gating_top_k_regbase.h"16#include "arch35/moe_gating_top_k_regbase.h"
17+#include "arch35/moe_gating_top_k_hash_regbase.h"
17#include "arch35/moe_gating_top_k_without_group_regbase.h"18#include "arch35/moe_gating_top_k_without_group_regbase.h"
18#include "arch35/moe_gating_top_k_e_k_fullload_regbase.h"19#include "arch35/moe_gating_top_k_e_k_fullload_regbase.h"
19using namespace AscendC;20using namespace AscendC;
@@ -57,20 +58,24 @@ extern "C" __global__ __aicore__ void moe_gating_top_k(GM_ADDR x, GM_ADDR bias,
57 MoeGatingTopKWithoutGroupRegbase<DTYPE_X> op;58 MoeGatingTopKWithoutGroupRegbase<DTYPE_X> op;
58 op.Init(x, bias, y, expertIdx, out, userWS, tilingData, &tPipe);59 op.Init(x, bias, y, expertIdx, out, userWS, tilingData, &tPipe);
59 op.Process();60 op.Process();
60- } else if (TILING_KEY_IS(TILING_KEY_REGBASE) || TILING_KEY_IS(TILING_KEY_HASH_INT32_INT32)) {61+ } else if (TILING_KEY_IS(TILING_KEY_REGBASE)) {
61 MoeGatingTopKRegbase<DTYPE_X> op;62 MoeGatingTopKRegbase<DTYPE_X> op;
63+ op.Init(x, bias, y, expertIdx, out, userWS, tilingData, &tPipe);
64+ op.Process();
65+ } else if (TILING_KEY_IS(TILING_KEY_HASH_INT32_INT32)) {
66+ MoeGatingTopKHashRegbase<DTYPE_X, int32_t, int32_t> op;
62 op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);67 op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);
63 op.Process();68 op.Process();
64 } else if (TILING_KEY_IS(TILING_KEY_HASH_INT32_INT64)) {69 } else if (TILING_KEY_IS(TILING_KEY_HASH_INT32_INT64)) {
65- MoeGatingTopKRegbase<DTYPE_X, int32_t, int64_t> op;70+ MoeGatingTopKHashRegbase<DTYPE_X, int32_t, int64_t> op;
66 op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);71 op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);
67 op.Process();72 op.Process();
68 } else if (TILING_KEY_IS(TILING_KEY_HASH_INT64_INT64)) {73 } else if (TILING_KEY_IS(TILING_KEY_HASH_INT64_INT64)) {
69- MoeGatingTopKRegbase<DTYPE_X, int64_t, int64_t> op;74+ MoeGatingTopKHashRegbase<DTYPE_X, int64_t, int64_t> op;
70 op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);75 op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);
71 op.Process();76 op.Process();
72 } else if (TILING_KEY_IS(TILING_KEY_HASH_INT64_INT32)) {77 } else if (TILING_KEY_IS(TILING_KEY_HASH_INT64_INT32)) {
73- MoeGatingTopKRegbase<DTYPE_X, int64_t, int32_t> op;78+ MoeGatingTopKHashRegbase<DTYPE_X, int64_t, int32_t> op;
74 op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);79 op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);
75 op.Process();80 op.Process();
76 }81 }