Star1294
2968
代码
代码
Issues158
Pull Requests525
Actions
项目讨论
Wiki
分析
Star1294
2968
已合并
dispatch将专家权重判断提至moe发送循环外，消除可选topk weight引入的性能劣化 #11394
zhong-zixin创建于 19 天前
dispatch将专家权重判断提至moe发送循环外，消除可选topk weight引入的性能劣化 #11394
已合并
zhong-zixin创建于 19 天前
从zhong-zixin/ops-transformerfix_topkw_loop合入到cann/ops-transformermaster
共 3 个文件变更+93-46
Mmc2/moe_distribute_dispatch_v2/op_kernel/arch22/moe_distribute_dispatch_v2_full_mesh.h+31-15
@@ -97,6 +97,8 @@ private:
97 __aicore__ inline void SetTilingDataAndCal(const MoeDistributeDispatchV2TilingData *tilingData);97 __aicore__ inline void SetTilingDataAndCal(const MoeDistributeDispatchV2TilingData *tilingData);
98 __aicore__ inline void SendToSharedExpert(TQue<QuePosition::VECIN, 1> inQueue, TBuf<> outBuf);98 __aicore__ inline void SendToSharedExpert(TQue<QuePosition::VECIN, 1> inQueue, TBuf<> outBuf);
99 __aicore__ inline void SendToMoeExpert(TQue<QuePosition::VECIN, 1> inQueue, TBuf<> expertMaskBuf, TBuf<> outBuf);99 __aicore__ inline void SendToMoeExpert(TQue<QuePosition::VECIN, 1> inQueue, TBuf<> expertMaskBuf, TBuf<> outBuf);
100+ __aicore__ inline void SendToMoeExpertLoop(TQue<QuePosition::VECIN, 1> inQueue, TBuf<> expertMaskBuf,
101+ bool writeExpertScale);
100 __aicore__ inline void SendToMoeExpertByBS(TQue<QuePosition::VECIN, 1> inQueue, TBuf<> outBuf);102 __aicore__ inline void SendToMoeExpertByBS(TQue<QuePosition::VECIN, 1> inQueue, TBuf<> outBuf);
101 __aicore__ inline void CalcBSTokenRange(uint32_t &startTokenId, uint32_t &endTokenId, uint32_t &myStartK,103 __aicore__ inline void CalcBSTokenRange(uint32_t &startTokenId, uint32_t &endTokenId, uint32_t &myStartK,
102 uint32_t &myEndK, bool &isMultiCorePerToken);104 uint32_t &myEndK, bool &isMultiCorePerToken);
@@ -133,13 +135,14 @@ private:
133 __aicore__ inline void SplitToCore(uint32_t curSendCnt, uint32_t curUseAivNum, uint32_t &startTokenId,135 __aicore__ inline void SplitToCore(uint32_t curSendCnt, uint32_t curUseAivNum, uint32_t &startTokenId,
134 uint32_t &endTokenId, uint32_t &sendTokenNum, bool isFront = true);136 uint32_t &endTokenId, uint32_t &sendTokenNum, bool isFront = true);
135 __aicore__ inline void SplitExpertNumToCore();137 __aicore__ inline void SplitExpertNumToCore();
136- __aicore__ inline void FillTriple(LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k);138+ __aicore__ inline void FillTriple(LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k,
139+ bool writeExpertScale);
137 __aicore__ inline void CalTokenSendExpertCnt(uint32_t dstExpertId, int32_t calCnt, int32_t &curExpertCnt);140 __aicore__ inline void CalTokenSendExpertCnt(uint32_t dstExpertId, int32_t calCnt, int32_t &curExpertCnt);
138 __aicore__ inline void TokenToExpertInQuant(GlobalTensor<XOutType> dstWinGMTensor,141 __aicore__ inline void TokenToExpertInQuant(GlobalTensor<XOutType> dstWinGMTensor,
139 TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,142 TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,
140- uint32_t toExpertId, uint32_t toExpertIndex);143+ uint32_t toExpertId, uint32_t toExpertIndex, bool writeExpertScale);
141 __aicore__ inline void TokenToExpert(GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue,144 __aicore__ inline void TokenToExpert(GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue,
142- uint32_t srcTokenIndex, uint32_t toExpertIndex);145+ uint32_t srcTokenIndex, uint32_t toExpertIndex, bool writeExpertScale);
143 __aicore__ inline void RecordRankCommDuration(LocalTensor<int32_t> &performanceInfoTensor, uint64_t startTime);146 __aicore__ inline void RecordRankCommDuration(LocalTensor<int32_t> &performanceInfoTensor, uint64_t startTime);
144 __aicore__ inline GM_ADDR GetExpertWinAddr(uint32_t dstExpertId, uint32_t tokenOffset);147 __aicore__ inline GM_ADDR GetExpertWinAddr(uint32_t dstExpertId, uint32_t tokenOffset);
145 __aicore__ inline GM_ADDR GetWindAddrByRankId(const int32_t rankId)148 __aicore__ inline GM_ADDR GetWindAddrByRankId(const int32_t rankId)
@@ -465,15 +468,15 @@ __aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFu
465 468 
466template <TemplateMC2TypeFullmeshClass>469template <TemplateMC2TypeFullmeshClass>
467__aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFunc>::FillTriple(470__aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFunc>::FillTriple(
468- LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k)471+ LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k, bool writeExpertScale)
469{472{
470 SyncFunc<AscendC::HardEvent::MTE3_S>();473 SyncFunc<AscendC::HardEvent::MTE3_S>();
471 LocalTensor<int32_t> xOutTint32 = xOutTensor.template ReinterpretCast<int32_t>();474 LocalTensor<int32_t> xOutTint32 = xOutTensor.template ReinterpretCast<int32_t>();
472 xOutTint32(tokenQuantAlign_) = epRankId_; // 0:epRankId index475 xOutTint32(tokenQuantAlign_) = epRankId_; // 0:epRankId index
473 xOutTint32(tokenQuantAlign_ + 1) = tokenIndex; // 1:token index476 xOutTint32(tokenQuantAlign_ + 1) = tokenIndex; // 1:token index
474 xOutTint32(tokenQuantAlign_ + 2) = k; // 2:topK value index477 xOutTint32(tokenQuantAlign_ + 2) = k; // 2:topK value index
475- LocalTensor<float> xOutTfloat = xOutTensor.template ReinterpretCast<float>();478+ if ((k < axisK_) && writeExpertScale) {
476- if ((k < axisK_) && hasExpertScalesFlag_) {479+ LocalTensor<float> xOutTfloat = xOutTensor.template ReinterpretCast<float>();
477 xOutTfloat(expertScaleAlign_) = expertScalesTensor_.GetValue(tokenIndex * axisK_ + k);480 xOutTfloat(expertScaleAlign_) = expertScalesTensor_.GetValue(tokenIndex * axisK_ + k);
478 }481 }
479 SyncFunc<AscendC::HardEvent::S_MTE3>();482 SyncFunc<AscendC::HardEvent::S_MTE3>();
@@ -482,7 +485,7 @@ __aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFu
482template <TemplateMC2TypeFullmeshClass>485template <TemplateMC2TypeFullmeshClass>
483__aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFunc>::TokenToExpertInQuant(486__aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFunc>::TokenToExpertInQuant(
484 GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,487 GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,
485- uint32_t fillExpertIdx, uint32_t quantExpertIdx)488+ uint32_t fillExpertIdx, uint32_t quantExpertIdx, bool writeExpertScale)
486{489{
487 DataCopyPadParams copyPadParams{true, 0U, 0U, 0U};490 DataCopyPadParams copyPadParams{true, 0U, 0U, 0U};
488 LocalTensor<XInType> xInTensor = inQueue.AllocTensor<XInType>();491 LocalTensor<XInType> xInTensor = inQueue.AllocTensor<XInType>();
@@ -493,7 +496,7 @@ __aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFu
493 if constexpr (QuantMode > UNQUANT) {496 if constexpr (QuantMode > UNQUANT) {
494 quantInst_.QuantProcess(tempTensor_, xInTensor, quantExpertIdx, scalesCount_, scalesGMTensor_);497 quantInst_.QuantProcess(tempTensor_, xInTensor, quantExpertIdx, scalesCount_, scalesGMTensor_);
495 }498 }
496- FillTriple(tempTensor_, srcTokenIndex, fillExpertIdx);499+ FillTriple(tempTensor_, srcTokenIndex, fillExpertIdx, writeExpertScale);
497 inQueue.FreeTensor<XInType>(xInTensor);500 inQueue.FreeTensor<XInType>(xInTensor);
498 SyncFunc<AscendC::HardEvent::S_V>();501 SyncFunc<AscendC::HardEvent::S_V>();
499 LocalTensor<int32_t> tempTensorInt32 = tempTensor_.template ReinterpretCast<int32_t>();502 LocalTensor<int32_t> tempTensorInt32 = tempTensor_.template ReinterpretCast<int32_t>();
@@ -514,7 +517,7 @@ __aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFu
514template <TemplateMC2TypeFullmeshClass>517template <TemplateMC2TypeFullmeshClass>
515__aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFunc>::TokenToExpert(518__aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFunc>::TokenToExpert(
516 GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,519 GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,
517- uint32_t toExpertIndex)520+ uint32_t toExpertIndex, bool writeExpertScale)
518{521{
519 DataCopyPadParams copyPadParams{false, 0U, 0U, 0U};522 DataCopyPadParams copyPadParams{false, 0U, 0U, 0U};
520 LocalTensor<XInType> xInTensor = inQueue.AllocTensor<XInType>();523 LocalTensor<XInType> xInTensor = inQueue.AllocTensor<XInType>();
@@ -523,7 +526,7 @@ __aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFu
523 }526 }
524 inQueue.EnQue(xInTensor);527 inQueue.EnQue(xInTensor);
525 xInTensor = inQueue.DeQue<XInType>();528 xInTensor = inQueue.DeQue<XInType>();
526- FillTriple(xInTensor, srcTokenIndex, toExpertIndex);529+ FillTriple(xInTensor, srcTokenIndex, toExpertIndex, writeExpertScale);
527 SyncFunc<AscendC::HardEvent::S_V>();530 SyncFunc<AscendC::HardEvent::S_V>();
528 LocalTensor<int32_t> xInTensorInt32 = xInTensor.template ReinterpretCast<int32_t>();531 LocalTensor<int32_t> xInTensorInt32 = xInTensor.template ReinterpretCast<int32_t>();
529 LocalTensor<int32_t> outTensorInt32 = outTensor_.template ReinterpretCast<int32_t>();532 LocalTensor<int32_t> outTensorInt32 = outTensor_.template ReinterpretCast<int32_t>();
@@ -601,9 +604,9 @@ __aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFu
601 if constexpr ((QuantMode > UNQUANT) || (QuantMode == UNQUANT && !Std::IsSame<ExpandXOutType, XType>::value)) {604 if constexpr ((QuantMode > UNQUANT) || (QuantMode == UNQUANT && !Std::IsSame<ExpandXOutType, XType>::value)) {
602 uint32_t fillExpertIdx = axisK_ + toSharedExpertIndex;605 uint32_t fillExpertIdx = axisK_ + toSharedExpertIndex;
603 uint32_t quantExpertIdx = toSharedExpertIndex;606 uint32_t quantExpertIdx = toSharedExpertIndex;
604- TokenToExpertInQuant(dstWinGMTensor, inQueue, srcTokenIndex, fillExpertIdx, quantExpertIdx);607+ TokenToExpertInQuant(dstWinGMTensor, inQueue, srcTokenIndex, fillExpertIdx, quantExpertIdx, false);
605 } else {608 } else {
606- TokenToExpert(dstWinGMTensor, inQueue, srcTokenIndex, axisK_ + toSharedExpertIndex);609+ TokenToExpert(dstWinGMTensor, inQueue, srcTokenIndex, axisK_ + toSharedExpertIndex, false);
607 }610 }
608 }611 }
609}612}
@@ -656,10 +659,22 @@ __aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFu
656 TQue<QuePosition::VECIN, 1> inQueue, TBuf<> expertMaskBuf, TBuf<> outBuf)659 TQue<QuePosition::VECIN, 1> inQueue, TBuf<> expertMaskBuf, TBuf<> outBuf)
657{660{
658 // 按专家分核（步进）661 // 按专家分核（步进）
659- uint32_t calExpertIdsIdx;
660 SplitExpertNumToCore();662 SplitExpertNumToCore();
661 // 计算专家发送数据量 && 发送663 // 计算专家发送数据量 && 发送
662 CalExpertSendNum(outBuf, expertMaskBuf);664 CalExpertSendNum(outBuf, expertMaskBuf);
665+ // 在循环外区分有无专家权重，避免性能劣化
666+ if (hasExpertScalesFlag_) {
667+ SendToMoeExpertLoop(inQueue, expertMaskBuf, true);
668+ } else {
669+ SendToMoeExpertLoop(inQueue, expertMaskBuf, false);
670+ }
671+}
672+ 
673+template <TemplateMC2TypeFullmeshClass>
674+__aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFunc>::SendToMoeExpertLoop(
675+ TQue<QuePosition::VECIN, 1> inQueue, TBuf<> expertMaskBuf, bool writeExpertScale)
676+{
677+ uint32_t calExpertIdsIdx;
663 uint32_t maskN64Num = Ceil(expertIdsCnt_, 64); // 64：ScalarGetSFFValue按照64长度一次计算678 uint32_t maskN64Num = Ceil(expertIdsCnt_, 64); // 64：ScalarGetSFFValue按照64长度一次计算
664 GlobalTensor<XOutType> dstWinGMTensor;679 GlobalTensor<XOutType> dstWinGMTensor;
665 LocalTensor<uint64_t> expertMaskTensorU64 = expertMaskBuf.Get<uint64_t>();680 LocalTensor<uint64_t> expertMaskTensorU64 = expertMaskBuf.Get<uint64_t>();
@@ -689,9 +704,10 @@ __aicore__ inline void MoeDistributeDispatchV2FullMesh<TemplateMC2TypeFullmeshFu
689 if constexpr ((QuantMode > UNQUANT) ||704 if constexpr ((QuantMode > UNQUANT) ||
690 (QuantMode == UNQUANT && !Std::IsSame<ExpandXOutType, XType>::value)) {705 (QuantMode == UNQUANT && !Std::IsSame<ExpandXOutType, XType>::value)) {
691 uint32_t quantExpertIdx = dstExpertId + sharedExpertNum_;706 uint32_t quantExpertIdx = dstExpertId + sharedExpertNum_;
692- TokenToExpertInQuant(dstWinGMTensor, inQueue, srcTokenIndex, topKIndex, quantExpertIdx);707+ TokenToExpertInQuant(dstWinGMTensor, inQueue, srcTokenIndex, topKIndex, quantExpertIdx,
708+ writeExpertScale);
693 } else {709 } else {
694- TokenToExpert(dstWinGMTensor, inQueue, srcTokenIndex, topKIndex);710+ TokenToExpert(dstWinGMTensor, inQueue, srcTokenIndex, topKIndex, writeExpertScale);
695 }711 }
696 dstTokenPreCnt++;712 dstTokenPreCnt++;
697 uint64_t cleanMask = ~(uint64_t(1) << curValidIdx);713 uint64_t cleanMask = ~(uint64_t(1) << curValidIdx);
Mmc2/moe_distribute_dispatch_v2/op_kernel/arch35/moe_distribute_dispatch_v2_a5_full_mesh.h+36-21
@@ -100,6 +100,8 @@ private:
100 __aicore__ inline void SetTilingDataAndCal(const MoeDistributeDispatchV2TilingData *tilingData);100 __aicore__ inline void SetTilingDataAndCal(const MoeDistributeDispatchV2TilingData *tilingData);
101 __aicore__ inline void SendToSharedExpert(TQue<QuePosition::VECIN, 1> inQueue, TBuf<> outBuf);101 __aicore__ inline void SendToSharedExpert(TQue<QuePosition::VECIN, 1> inQueue, TBuf<> outBuf);
102 __aicore__ inline void SendToMoeExpert(TQue<QuePosition::VECIN, 1> inQueue, TBuf<> outBuf);102 __aicore__ inline void SendToMoeExpert(TQue<QuePosition::VECIN, 1> inQueue, TBuf<> outBuf);
103+ __aicore__ inline void SendToMoeExpertLoop(TQue<QuePosition::VECIN, 1> inQueue, uint32_t validTokenNum,
104+ bool writeExpertScale);
103 __aicore__ inline void ExpertActiveMaskInit();105 __aicore__ inline void ExpertActiveMaskInit();
104 __aicore__ inline void ExpertActiveMaskCal();106 __aicore__ inline void ExpertActiveMaskCal();
105 __aicore__ inline void CalcSendTokenBufNum(TBuf<> &outBuf);107 __aicore__ inline void CalcSendTokenBufNum(TBuf<> &outBuf);
@@ -131,13 +133,14 @@ private:
131 __aicore__ inline void ClearLocalWindowDataFlags(TBuf<> tBuf, uint32_t validNum, uint32_t localExpertNum);133 __aicore__ inline void ClearLocalWindowDataFlags(TBuf<> tBuf, uint32_t validNum, uint32_t localExpertNum);
132 __aicore__ inline void SplitToCore(uint32_t curSendCnt, uint32_t curUseAivNum, uint32_t &startTokenId,134 __aicore__ inline void SplitToCore(uint32_t curSendCnt, uint32_t curUseAivNum, uint32_t &startTokenId,
133 uint32_t &endTokenId, uint32_t &sendTokenNum, bool isFront = true);135 uint32_t &endTokenId, uint32_t &sendTokenNum, bool isFront = true);
134- __aicore__ inline void FillTriple(LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k);136+ __aicore__ inline void FillTriple(LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k,
137+ bool writeExpertScale);
135 __aicore__ inline void CalTokenSendExpertCnt(uint32_t dstExpertId, int32_t calCnt, int32_t &curExpertCnt);138 __aicore__ inline void CalTokenSendExpertCnt(uint32_t dstExpertId, int32_t calCnt, int32_t &curExpertCnt);
136 __aicore__ inline void TokenToExpertInQuant(GlobalTensor<XOutType> dstWinGMTensor,139 __aicore__ inline void TokenToExpertInQuant(GlobalTensor<XOutType> dstWinGMTensor,
137 TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,140 TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,
138- uint32_t toExpertId, uint32_t toExpertIndex);141+ uint32_t toExpertId, uint32_t toExpertIndex, bool writeExpertScale);
139 __aicore__ inline void TokenToExpert(GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue,142 __aicore__ inline void TokenToExpert(GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue,
140- uint32_t srcTokenIndex, uint32_t toExpertIndex);143+ uint32_t srcTokenIndex, uint32_t toExpertIndex, bool writeExpertScale);
141 __aicore__ inline void RecordRankCommDuration(LocalTensor<int32_t> &performanceInfoTensor, uint64_t startTime);144 __aicore__ inline void RecordRankCommDuration(LocalTensor<int32_t> &performanceInfoTensor, uint64_t startTime);
142 __aicore__ inline GM_ADDR GetWindAddrByRankId(const int32_t rankId)145 __aicore__ inline GM_ADDR GetWindAddrByRankId(const int32_t rankId)
143 {146 {
@@ -468,14 +471,14 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy
468 471 
469template <TemplateMC2A5FullMeshTypeClass>472template <TemplateMC2A5FullMeshTypeClass>
470__aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::FillTriple(473__aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::FillTriple(
471- LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k)474+ LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k, bool writeExpertScale)
472{475{
473 LocalTensor<int32_t> xOutTint32 = xOutTensor.template ReinterpretCast<int32_t>();476 LocalTensor<int32_t> xOutTint32 = xOutTensor.template ReinterpretCast<int32_t>();
474 xOutTint32(tokenQuantAlign_) = epRankId_; // 0:epRankId index477 xOutTint32(tokenQuantAlign_) = epRankId_; // 0:epRankId index
475 xOutTint32(tokenQuantAlign_ + 1) = tokenIndex; // 1:token index478 xOutTint32(tokenQuantAlign_ + 1) = tokenIndex; // 1:token index
476 xOutTint32(tokenQuantAlign_ + 2) = k; // 2:topK value index479 xOutTint32(tokenQuantAlign_ + 2) = k; // 2:topK value index
477- LocalTensor<float> xOutTfloat = xOutTensor.template ReinterpretCast<float>();480+ if ((k < axisK_) && writeExpertScale) {
478- if ((k < axisK_) && hasExpertScalesFlag_) {481+ LocalTensor<float> xOutTfloat = xOutTensor.template ReinterpretCast<float>();
479 xOutTfloat(expertScaleAlign_) = expertScalesTensor_.GetValue(tokenIndex * axisK_ + k);482 xOutTfloat(expertScaleAlign_) = expertScalesTensor_.GetValue(tokenIndex * axisK_ + k);
480 }483 }
481}484}
@@ -483,7 +486,7 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy
483template <TemplateMC2A5FullMeshTypeClass>486template <TemplateMC2A5FullMeshTypeClass>
484__aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::TokenToExpertInQuant(487__aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::TokenToExpertInQuant(
485 GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,488 GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,
486- uint32_t fillExpertIdx, uint32_t quantExpertIdx)489+ uint32_t fillExpertIdx, uint32_t quantExpertIdx, bool writeExpertScale)
487{490{
488 DataCopyPadParams copyPadParams{true, 0U, 0U, 0U};491 DataCopyPadParams copyPadParams{true, 0U, 0U, 0U};
489 LocalTensor<XInType> xInTensor = inQueue.AllocTensor<XInType>();492 LocalTensor<XInType> xInTensor = inQueue.AllocTensor<XInType>();
@@ -498,7 +501,7 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy
498 }501 }
499 inQueue.FreeTensor<XInType>(xInTensor);502 inQueue.FreeTensor<XInType>(xInTensor);
500 SyncFunc<AscendC::HardEvent::V_S>();503 SyncFunc<AscendC::HardEvent::V_S>();
501- FillTriple(tempTensor_, srcTokenIndex, fillExpertIdx);504+ FillTriple(tempTensor_, srcTokenIndex, fillExpertIdx, writeExpertScale);
502 SyncFunc<AscendC::HardEvent::S_V>();505 SyncFunc<AscendC::HardEvent::S_V>();
503 AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(syncFlagId_ % sendTokenBufNum_);506 AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(syncFlagId_ % sendTokenBufNum_);
504 LocalTensor<int32_t> tempTensorInt32 = tempTensor_.template ReinterpretCast<int32_t>();507 LocalTensor<int32_t> tempTensorInt32 = tempTensor_.template ReinterpretCast<int32_t>();
@@ -519,7 +522,7 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy
519template <TemplateMC2A5FullMeshTypeClass>522template <TemplateMC2A5FullMeshTypeClass>
520__aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::TokenToExpert(523__aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::TokenToExpert(
521 GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,524 GlobalTensor<XOutType> dstWinGMTensor, TQue<QuePosition::VECIN, 1> inQueue, uint32_t srcTokenIndex,
522- uint32_t toExpertIndex)525+ uint32_t toExpertIndex, bool writeExpertScale)
523{526{
524 DataCopyPadParams copyPadParams{false, 0U, 0U, 0U};527 DataCopyPadParams copyPadParams{false, 0U, 0U, 0U};
525 LocalTensor<XInType> xInTensor = inQueue.AllocTensor<XInType>();528 LocalTensor<XInType> xInTensor = inQueue.AllocTensor<XInType>();
@@ -536,7 +539,7 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy
536 inQueue.EnQue(xInTensor);539 inQueue.EnQue(xInTensor);
537 xInTensor = inQueue.DeQue<XInType>();540 xInTensor = inQueue.DeQue<XInType>();
538 SyncFunc<AscendC::HardEvent::MTE2_S>();541 SyncFunc<AscendC::HardEvent::MTE2_S>();
539- FillTriple(xInTensor, srcTokenIndex, toExpertIndex);542+ FillTriple(xInTensor, srcTokenIndex, toExpertIndex, writeExpertScale);
540 SyncFunc<AscendC::HardEvent::S_V>();543 SyncFunc<AscendC::HardEvent::S_V>();
541 AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(syncFlagId_ % sendTokenBufNum_);544 AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(syncFlagId_ % sendTokenBufNum_);
542 LocalTensor<int32_t> xInTensorInt32 = xInTensor.template ReinterpretCast<int32_t>();545 LocalTensor<int32_t> xInTensorInt32 = xInTensor.template ReinterpretCast<int32_t>();
@@ -620,9 +623,9 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy
620 if constexpr ((QuantMode > UNQUANT) || (QuantMode == UNQUANT && !Std::IsSame<ExpandXOutType, XType>::value)) {623 if constexpr ((QuantMode > UNQUANT) || (QuantMode == UNQUANT && !Std::IsSame<ExpandXOutType, XType>::value)) {
621 uint32_t fillExpertIdx = axisK_ + toSharedExpertIndex;624 uint32_t fillExpertIdx = axisK_ + toSharedExpertIndex;
622 uint32_t quantExpertIdx = toSharedExpertIndex;625 uint32_t quantExpertIdx = toSharedExpertIndex;
623- TokenToExpertInQuant(dstWinGMTensor, inQueue, srcTokenIndex, fillExpertIdx, quantExpertIdx);626+ TokenToExpertInQuant(dstWinGMTensor, inQueue, srcTokenIndex, fillExpertIdx, quantExpertIdx, false);
624 } else {627 } else {
625- TokenToExpert(dstWinGMTensor, inQueue, srcTokenIndex, axisK_ + toSharedExpertIndex);628+ TokenToExpert(dstWinGMTensor, inQueue, srcTokenIndex, axisK_ + toSharedExpertIndex, false);
626 }629 }
627 }630 }
628 for (int i = 0; i < sendTokenBufNum_; i++) {631 for (int i = 0; i < sendTokenBufNum_; i++) {
@@ -638,15 +641,31 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy
638 LocalTensor<float> outTensorFp32 = outBuf.Get<float>();641 LocalTensor<float> outTensorFp32 = outBuf.Get<float>();
639 Duplicate<float>(outTensorFp32, float(1), hCommuSize_ * sendTokenBufNum_ / sizeof(float));642 Duplicate<float>(outTensorFp32, float(1), hCommuSize_ * sendTokenBufNum_ / sizeof(float));
640 uint32_t validTokenNum = isTokenMaskFlag_ ? (activeMaskBsCnt_ * axisK_) : expertIdsCnt_;643 uint32_t validTokenNum = isTokenMaskFlag_ ? (activeMaskBsCnt_ * axisK_) : expertIdsCnt_;
641- GlobalTensor<XOutType> dstWinGMTensor;
642- 
643- int32_t dstTokenIdx = 0;
644 syncFlagId_ = 0;644 syncFlagId_ = 0;
645 645 
646 for (int i = 0; i < sendTokenBufNum_; i++) {646 for (int i = 0; i < sendTokenBufNum_; i++) {
647 AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(i % sendTokenBufNum_);647 AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(i % sendTokenBufNum_);
648 }648 }
649 649 
650+ // 在循环外区分有无专家权重，避免性能劣化
651+ if (hasExpertScalesFlag_) {
652+ SendToMoeExpertLoop(inQueue, validTokenNum, true);
653+ } else {
654+ SendToMoeExpertLoop(inQueue, validTokenNum, false);
655+ }
656+ 
657+ for (int i = 0; i < sendTokenBufNum_; i++) {
658+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(i % sendTokenBufNum_);
659+ }
660+}
661+ 
662+template <TemplateMC2A5FullMeshTypeClass>
663+__aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::SendToMoeExpertLoop(
664+ TQue<QuePosition::VECIN, 1> inQueue, uint32_t validTokenNum, bool writeExpertScale)
665+{
666+ GlobalTensor<XOutType> dstWinGMTensor;
667+ int32_t dstTokenIdx = 0;
668+ 
650 for (int32_t index = aivId_; index < validTokenNum; index += moeUsedAivNum_) {669 for (int32_t index = aivId_; index < validTokenNum; index += moeUsedAivNum_) {
651 int32_t tokenId = index / axisK_;670 int32_t tokenId = index / axisK_;
652 int32_t topKId = index % axisK_;671 int32_t topKId = index % axisK_;
@@ -673,15 +692,11 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy
673 }692 }
674 if constexpr ((QuantMode > UNQUANT) || (QuantMode == UNQUANT && !Std::IsSame<ExpandXOutType, XType>::value)) {693 if constexpr ((QuantMode > UNQUANT) || (QuantMode == UNQUANT && !Std::IsSame<ExpandXOutType, XType>::value)) {
675 uint32_t quantExpertIdx = expertId + sharedExpertNum_;694 uint32_t quantExpertIdx = expertId + sharedExpertNum_;
676- TokenToExpertInQuant(dstWinGMTensor, inQueue, tokenId, topKId, quantExpertIdx);695+ TokenToExpertInQuant(dstWinGMTensor, inQueue, tokenId, topKId, quantExpertIdx, writeExpertScale);
677 } else {696 } else {
678- TokenToExpert(dstWinGMTensor, inQueue, tokenId, topKId);697+ TokenToExpert(dstWinGMTensor, inQueue, tokenId, topKId, writeExpertScale);
679 }698 }
680 }699 }
681- 
682- for (int i = 0; i < sendTokenBufNum_; i++) {
683- AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(i % sendTokenBufNum_);
684- }
685}700}
686 701 
687template <TemplateMC2A5FullMeshTypeClass>702template <TemplateMC2A5FullMeshTypeClass>
Mmc2/moe_distribute_dispatch_v2/op_kernel/moe_distribute_dispatch_v2.h+26-10
@@ -78,9 +78,10 @@ public:
78private:78private:
79 __aicore__ inline void ProcessToken(GlobalTensor<XOutType> &outTokenGT, uint32_t tokenIndex, uint32_t topKIndex,79 __aicore__ inline void ProcessToken(GlobalTensor<XOutType> &outTokenGT, uint32_t tokenIndex, uint32_t topKIndex,
80 DataCopyPadParams &padParams, DataCopyParams &scaleInParams,80 DataCopyPadParams &padParams, DataCopyParams &scaleInParams,
81- uint32_t expertIndex);81+ uint32_t expertIndex, bool writeExpertScale);
82 __aicore__ inline void SendToSharedExpert();82 __aicore__ inline void SendToSharedExpert();
83 __aicore__ inline void SendToMoeExpert();83 __aicore__ inline void SendToMoeExpert();
84+ __aicore__ inline void SendToMoeExpertLoop(uint32_t startTokenId, uint32_t endTokenId, bool writeExpertScale);
84 __aicore__ inline void AlltoAllDispatch();85 __aicore__ inline void AlltoAllDispatch();
85 __aicore__ inline void LocalWindowCopy();86 __aicore__ inline void LocalWindowCopy();
86#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)87#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)
@@ -104,7 +105,8 @@ private:
104 __aicore__ inline void GetCumSumA5(LocalTensor<int32_t> &outLocal);105 __aicore__ inline void GetCumSumA5(LocalTensor<int32_t> &outLocal);
105 __aicore__ inline void SplitToCore(uint32_t curSendCnt, uint32_t curUseAivNum, uint32_t &startTokenId,106 __aicore__ inline void SplitToCore(uint32_t curSendCnt, uint32_t curUseAivNum, uint32_t &startTokenId,
106 uint32_t &endTokenId, uint32_t &sendTokenNum, bool isFront = true);107 uint32_t &endTokenId, uint32_t &sendTokenNum, bool isFront = true);
107- __aicore__ inline void FillTriple(LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k);108+ __aicore__ inline void FillTriple(LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k,
109+ bool writeExpertScale);
108 __aicore__ inline void CalTokenSendExpertCnt(uint32_t dstExpertId, int32_t calCnt, int32_t &curExpertCnt);110 __aicore__ inline void CalTokenSendExpertCnt(uint32_t dstExpertId, int32_t calCnt, int32_t &curExpertCnt);
109 __aicore__ inline void SyncCntOnCore(LocalTensor<float> &gatherMaskOutTensor,111 __aicore__ inline void SyncCntOnCore(LocalTensor<float> &gatherMaskOutTensor,
110 LocalTensor<uint32_t> &gatherTmpTensor,112 LocalTensor<uint32_t> &gatherTmpTensor,
@@ -586,14 +588,14 @@ __aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::Spli
586 588 
587template <TemplateDispatchV2TypeClass>589template <TemplateDispatchV2TypeClass>
588__aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::FillTriple(590__aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::FillTriple(
589- LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k)591+ LocalTensor<XOutType> &xOutTensor, uint32_t tokenIndex, uint32_t k, bool writeExpertScale)
590{592{
591 LocalTensor<int32_t> xOutTint32 = xOutTensor.template ReinterpretCast<int32_t>();593 LocalTensor<int32_t> xOutTint32 = xOutTensor.template ReinterpretCast<int32_t>();
592 xOutTint32(tokenQuantAlign_) = epRankId_;594 xOutTint32(tokenQuantAlign_) = epRankId_;
593 xOutTint32(tokenQuantAlign_ + 1) = tokenIndex;595 xOutTint32(tokenQuantAlign_ + 1) = tokenIndex;
594 xOutTint32(tokenQuantAlign_ + 2) = k;596 xOutTint32(tokenQuantAlign_ + 2) = k;
595- LocalTensor<float> xOutTfloat = xOutTensor.template ReinterpretCast<float>();597+ if (writeExpertScale) {
596- if ((k < axisK_) && (hasExpertScalesFlag_)) {598+ LocalTensor<float> xOutTfloat = xOutTensor.template ReinterpretCast<float>();
597 xOutTfloat(expertScaleAlign_) = expertScalesTensor_.GetValue(tokenIndex * axisK_ + k);599 xOutTfloat(expertScaleAlign_) = expertScalesTensor_.GetValue(tokenIndex * axisK_ + k);
598 }600 }
599}601}
@@ -601,7 +603,7 @@ __aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::Fill
601template <TemplateDispatchV2TypeClass>603template <TemplateDispatchV2TypeClass>
602__aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::ProcessToken(604__aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::ProcessToken(
603 GlobalTensor<XOutType> &outTokenGT, uint32_t tokenIndex, uint32_t topKIndex, DataCopyPadParams &padParams,605 GlobalTensor<XOutType> &outTokenGT, uint32_t tokenIndex, uint32_t topKIndex, DataCopyPadParams &padParams,
604- DataCopyParams &scaleInParams, uint32_t expertIndex)606+ DataCopyParams &scaleInParams, uint32_t expertIndex, bool writeExpertScale)
605{607{
606 if constexpr ((QuantMode > UNQUANT) || (QuantMode == UNQUANT && !Std::IsSame<ExpandXOutType, XType>::value)) {608 if constexpr ((QuantMode > UNQUANT) || (QuantMode == UNQUANT && !Std::IsSame<ExpandXOutType, XType>::value)) {
607 xInTensor_ = xInQueue_.AllocTensor<XInType>();609 xInTensor_ = xInQueue_.AllocTensor<XInType>();
@@ -627,7 +629,7 @@ __aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::Proc
627 }629 }
628#endif630#endif
629 SyncFunc<AscendC::HardEvent::V_S>();631 SyncFunc<AscendC::HardEvent::V_S>();
630- FillTriple(xOutTensor_, tokenIndex, topKIndex);632+ FillTriple(xOutTensor_, tokenIndex, topKIndex, writeExpertScale);
631 xOutQueue_.EnQue(xOutTensor_);633 xOutQueue_.EnQue(xOutTensor_);
632 xInQueue_.FreeTensor<XInType>(xInTensor_);634 xInQueue_.FreeTensor<XInType>(xInTensor_);
633 xOutTensor_ = xOutQueue_.DeQue<XOutType>();635 xOutTensor_ = xOutQueue_.DeQue<XOutType>();
@@ -647,7 +649,7 @@ __aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::Proc
647 }649 }
648#endif650#endif
649 SyncFunc<AscendC::HardEvent::MTE2_S>();651 SyncFunc<AscendC::HardEvent::MTE2_S>();
650- FillTriple(xTmpTensor_, tokenIndex, topKIndex);652+ FillTriple(xTmpTensor_, tokenIndex, topKIndex, writeExpertScale);
651 xQueue_.EnQue(xTmpTensor_);653 xQueue_.EnQue(xTmpTensor_);
652 xTmpTensor_ = xQueue_.DeQue<XOutType>();654 xTmpTensor_ = xQueue_.DeQue<XOutType>();
653 DataCopyPad(outTokenGT, xTmpTensor_, hCommuCopyOutParams_);655 DataCopyPad(outTokenGT, xTmpTensor_, hCommuCopyOutParams_);
@@ -684,7 +686,7 @@ __aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::Send
684 }686 }
685 GlobalTensor<XOutType> tempTensor = dstWinGMTensor[tokenIndex * hAlignWinCnt_];687 GlobalTensor<XOutType> tempTensor = dstWinGMTensor[tokenIndex * hAlignWinCnt_];
686 ProcessToken(tempTensor, srcTokenIndex, axisK_ + toSharedExpertIndex, padParams, scaleInParams,688 ProcessToken(tempTensor, srcTokenIndex, axisK_ + toSharedExpertIndex, padParams, scaleInParams,
687- toSharedExpertIndex);689+ toSharedExpertIndex, false);
688 }690 }
689}691}
690 692 
@@ -726,6 +728,19 @@ __aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::Send
726 if (startTokenId >= sendToMoeExpTokenCnt_) {728 if (startTokenId >= sendToMoeExpTokenCnt_) {
727 return;729 return;
728 }730 }
731+ // 在循环外区分有无专家权重，避免性能劣化
732+ if (hasExpertScalesFlag_) {
733+ SendToMoeExpertLoop(startTokenId, endTokenId, true);
734+ } else {
735+ SendToMoeExpertLoop(startTokenId, endTokenId, false);
736+ }
737+}
738+ 
739+template <TemplateDispatchV2TypeClass>
740+__aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::SendToMoeExpertLoop(uint32_t startTokenId,
741+ uint32_t endTokenId,
742+ bool writeExpertScale)
743+{
729 GlobalTensor<XOutType> dstWinGMTensor;744 GlobalTensor<XOutType> dstWinGMTensor;
730 DataCopyPadParams padParams = {true, 0, 0, 0};745 DataCopyPadParams padParams = {true, 0, 0, 0};
731 DataCopyParams scaleInParams = {1U, static_cast<uint16_t>(scaleInBytes_), 0U, 0U};746 DataCopyParams scaleInParams = {1U, static_cast<uint16_t>(scaleInBytes_), 0U, 0U};
@@ -762,7 +777,8 @@ __aicore__ inline void MoeDistributeDispatchV2<TemplateDispatchV2TypeFunc>::Send
762 (epRankId_ * moeExpertNumPerRank_ + dstExpertId % moeExpertNumPerRank_)) +777 (epRankId_ * moeExpertNumPerRank_ + dstExpertId % moeExpertNumPerRank_)) +
763 hAlignWinSize_ * curExpertCnt); // 计算地址偏移778 hAlignWinSize_ * curExpertCnt); // 计算地址偏移
764 dstWinGMTensor.SetGlobalBuffer((__gm__ XOutType *)rankGM);779 dstWinGMTensor.SetGlobalBuffer((__gm__ XOutType *)rankGM);
765- ProcessToken(dstWinGMTensor, tokenIndex, topKIndex, padParams, scaleInParams, dstExpertId + sharedExpertNum_);780+ ProcessToken(dstWinGMTensor, tokenIndex, topKIndex, padParams, scaleInParams, dstExpertId + sharedExpertNum_,
781+ writeExpertScale);
766 }782 }
767}783}
768 784