Merged
fullmeshv2 flag轮询与搬出优化 #15481
zhong-zixinAdd on 9 days ago
fullmeshv2 flag轮询与搬出优化 #15481
Merged
Total: 1 files changed+156-98
| @@ -141,10 +141,14 @@ private: | |||
| 141 | __aicore__ inline int32_t GetRecvCntPrefix(uint32_t statusId); | 141 | __aicore__ inline int32_t GetRecvCntPrefix(uint32_t statusId); |
| 142 | __aicore__ inline void SetValidExpertInfo(uint32_t expInfoSize, uint32_t& validNum); | 142 | __aicore__ inline void SetValidExpertInfo(uint32_t expInfoSize, uint32_t& validNum); |
| 143 | __aicore__ inline void SetValidExpertInfoTokenNums(); | 143 | __aicore__ inline void SetValidExpertInfoTokenNums(); |
| 144 | - __aicore__ inline uint32_t CheckDataArriveWithFlag(uint32_t srcExpDataIdx, int32_t beginIdx, int32_t copyCnt); | 144 | + __aicore__ inline uint32_t LoadDataFlags(uint32_t beginIndex, uint32_t validNum, uint32_t maxCopyTokenCnt, |
| 145 | - __aicore__ inline void CopyInAndOut(LocalTensor<XOutType> xTmpTensor, GM_ADDR wAddr, uint32_t index, | 145 | + uint32_t localExpertNum); |
| 146 | - uint32_t dstPosition, uint32_t arriveCount, uint32_t dataBufferId, | 146 | + __aicore__ inline bool CheckDataArriveWithFlag(uint32_t beginToken, uint32_t copyCnt); |
| 147 | - bool isDataBufferReused); | 147 | + __aicore__ inline uint64_t LoadReadyData(LocalTensor<XOutType> xTmpTensor, uint32_t beginIndex, uint32_t endIndex, |
| 148 | + uint32_t maxCopyTokenCnt, uint32_t localExpertNum, uint32_t dataBufferId, | ||
| 149 | + bool isDataBufferReused); | ||
| 150 | + __aicore__ inline void CopyOut(LocalTensor<XOutType> xTmpTensor, uint32_t index, uint32_t dstPosition, | ||
| 151 | + uint32_t copyCnt, uint32_t dataBufferId); | ||
| 148 | __aicore__ inline void WaitAndFormatOutput(TBuf<> tBuf, uint32_t validNum); | 152 | __aicore__ inline void WaitAndFormatOutput(TBuf<> tBuf, uint32_t validNum); |
| 149 | __aicore__ inline uint32_t GetLocalWindowSrcDataBlockIdx(uint32_t srcExpertId, uint32_t localExpertNum); | 153 | __aicore__ inline uint32_t GetLocalWindowSrcDataBlockIdx(uint32_t srcExpertId, uint32_t localExpertNum); |
| 150 | __aicore__ inline void ClearLocalWindowDataFlags(TBuf<> tBuf, uint32_t validNum, uint32_t localExpertNum); | 154 | __aicore__ inline void ClearLocalWindowDataFlags(TBuf<> tBuf, uint32_t validNum, uint32_t localExpertNum); |
| @@ -227,7 +231,7 @@ private: | |||
| 227 | LocalTensor<uint32_t> expertLeftNumTensor_; | 231 | LocalTensor<uint32_t> expertLeftNumTensor_; |
| 228 | LocalTensor<uint32_t> dataStateLocalTensor_; | 232 | LocalTensor<uint32_t> dataStateLocalTensor_; |
| 229 | 233 | ||
| 230 | - LocalTensor<float> flagReduceWorkTensor_; | 234 | + LocalTensor<uint64_t> flagReadyTensor_; |
| 231 | LocalTensor<float> flagRecvTensor_; | 235 | LocalTensor<float> flagRecvTensor_; |
| 232 | 236 | ||
| 233 | TBuf<> statusBuf_; | 237 | TBuf<> statusBuf_; |
| @@ -2081,65 +2085,131 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy | |||
| 2081 | } | 2085 | } |
| 2082 | 2086 | ||
| 2083 | template <TemplateMC2A5FullMeshTypeClass> | 2087 | template <TemplateMC2A5FullMeshTypeClass> |
| 2084 | -__aicore__ inline uint32_t MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::CheckDataArriveWithFlag( | 2088 | +__aicore__ inline uint32_t MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::LoadDataFlags( |
| 2085 | - uint32_t srcExpDataIdx, int32_t beginIdx, int32_t copyCnt) | 2089 | + uint32_t beginIndex, uint32_t validNum, uint32_t maxCopyTokenCnt, uint32_t localExpertNum) |
| 2086 | { | 2090 | { |
| 2087 | - uint32_t flagNum = blockCntPerToken_ * uint32_t(copyCnt); | 2091 | + uint32_t endIndex = beginIndex; |
| 2088 | - DataCopyParams expFlagCopyParams{static_cast<uint16_t>(flagNum), 1U, | 2092 | + uint32_t tokenNum = 0U; |
| 2089 | - static_cast<uint16_t>(SPLIT_BLOCK_DATA_SIZE / UB_ALIGN), 0U}; | 2093 | + // 每个片段占一位，单批片段数不超过就绪掩码位数 |
| 2090 | - GlobalTensor<float> dataFlagGlobal; | 2094 | + constexpr uint32_t readyMaskBits = sizeof(uint64_t) * BITS_PER_BYTE; |
| 2091 | - GM_ADDR wAddr = (__gm__ uint8_t*)(windowGM_) + srcExpDataIdx * dataSectionSize_ + beginIdx * hCommuSize_ + | 2095 | + while (endIndex < validNum && tokenNum < maxCopyTokenCnt && endIndex - beginIndex < readyMaskBits) { |
| 2092 | - SPLIT_BLOCK_DATA_SIZE; | 2096 | + uint32_t copyCnt = expertLeftNumTensor_(endIndex); |
| 2093 | - dataFlagGlobal.SetGlobalBuffer((__gm__ float*)(wAddr)); | 2097 | + copyCnt = copyCnt > maxCopyTokenCnt - tokenNum ? maxCopyTokenCnt - tokenNum : copyCnt; |
| 2094 | - DataCopy(flagRecvTensor_, dataFlagGlobal, expFlagCopyParams); | 2098 | + if (copyCnt != 0U) { |
| 2099 | + uint32_t source = GetLocalWindowSrcDataBlockIdx(expertMapTensor_(endIndex), localExpertNum); | ||
| 2100 | + GlobalTensor<float> dataFlagGlobal; | ||
| 2101 | + dataFlagGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float*>( | ||
| 2102 | + windowGM_ + source * dataSectionSize_ + expertFinishNumTensor_(endIndex) * hCommuSize_ + | ||
| 2103 | + SPLIT_BLOCK_DATA_SIZE)); | ||
| 2104 | + DataCopyParams copyParams{static_cast<uint16_t>(copyCnt * blockCntPerToken_), 1U, | ||
| 2105 | + static_cast<uint16_t>(SPLIT_BLOCK_DATA_SIZE / UB_ALIGN), 0U}; | ||
| 2106 | + DataCopy(flagRecvTensor_[tokenNum * blockCntPerToken_ * UB_ALIGN_DATA_COUNT], dataFlagGlobal, copyParams); | ||
| 2107 | + tokenNum += copyCnt; | ||
| 2108 | + } | ||
| 2109 | + ++endIndex; | ||
| 2110 | + } | ||
| 2111 | + if (tokenNum == 0U) { | ||
| 2112 | + return endIndex; | ||
| 2113 | + } | ||
| 2114 | + | ||
| 2115 | + // 按实际标志数比较，尾部不读取额外标志块 | ||
| 2095 | SyncFunc<AscendC::HardEvent::MTE2_V>(); | 2116 | SyncFunc<AscendC::HardEvent::MTE2_V>(); |
| 2096 | - LocalTensor<float> flagSumOutTensor = scalarBuf_.Get<float>(); | 2117 | + SetMaskCount(); |
| 2097 | - ReduceSum(flagSumOutTensor, flagRecvTensor_, flagReduceWorkTensor_, 1U, flagNum, 1U); | 2118 | + CompareScalar(flagReadyTensor_.ReinterpretCast<uint8_t>(), flagRecvTensor_, 1.0F, CMPMODE::EQ, |
| 2119 | + static_cast<uint64_t>(tokenNum) * blockCntPerToken_ * UB_ALIGN_DATA_COUNT, 1, | ||
| 2120 | + UnaryRepeatParams{1, 1, 8, 8}); | ||
| 2121 | + SetMaskNorm(); | ||
| 2098 | SyncFunc<AscendC::HardEvent::V_S>(); | 2122 | SyncFunc<AscendC::HardEvent::V_S>(); |
| 2099 | - // aicore禁止uint->float的显式转换，借助浮点乘法隐式转换得到期望和 | 2123 | + return endIndex; |
| 2100 | - float expectFlagSum = 1.0F * flagNum; | ||
| 2101 | - return flagSumOutTensor.GetValue(0) == expectFlagSum ? static_cast<uint32_t>(copyCnt) : 0U; | ||
| 2102 | } | 2124 | } |
| 2103 | 2125 | ||
| 2104 | template <TemplateMC2A5FullMeshTypeClass> | 2126 | template <TemplateMC2A5FullMeshTypeClass> |
| 2105 | -__aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::CopyInAndOut( | 2127 | +__aicore__ inline bool MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::CheckDataArriveWithFlag( |
| 2106 | - LocalTensor<XOutType> xTmpTensor, GM_ADDR wAddr, uint32_t index, uint32_t dstPosition, uint32_t arriveCount, | 2128 | + uint32_t beginToken, uint32_t copyCnt) |
| 2107 | - uint32_t dataBufferId, bool isDataBufferReused) | 2129 | +{ |
| 2130 | + // 每个32B标志块对应一字节，只取首元素的比较位 | ||
| 2131 | + constexpr uint64_t readyMask = 0x0101010101010101ULL; | ||
| 2132 | + constexpr uint32_t flagsPerWord = sizeof(uint64_t); | ||
| 2133 | + uint32_t beginFlag = beginToken * blockCntPerToken_; | ||
| 2134 | + uint32_t endFlag = beginFlag + copyCnt * blockCntPerToken_; | ||
| 2135 | + while (beginFlag < endFlag) { | ||
| 2136 | + uint32_t offset = beginFlag % flagsPerWord; | ||
| 2137 | + uint32_t flagNum = endFlag - beginFlag; | ||
| 2138 | + flagNum = flagNum > flagsPerWord - offset ? flagsPerWord - offset : flagNum; | ||
| 2139 | + uint64_t expected = readyMask >> ((flagsPerWord - flagNum) * BITS_PER_BYTE); | ||
| 2140 | + uint64_t flags = flagReadyTensor_.GetValue(beginFlag / flagsPerWord) >> (offset * BITS_PER_BYTE); | ||
C | |||
| 2141 | + if ((flags & expected) != expected) { | ||
| 2142 | + return false; | ||
| 2143 | + } | ||
| 2144 | + beginFlag += flagNum; | ||
| 2145 | + } | ||
| 2146 | + return true; | ||
| 2147 | +} | ||
| 2148 | + | ||
| 2149 | +template <TemplateMC2A5FullMeshTypeClass> | ||
| 2150 | +__aicore__ inline uint64_t MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::LoadReadyData( | ||
| 2151 | + LocalTensor<XOutType> xTmpTensor, uint32_t beginIndex, uint32_t endIndex, uint32_t maxCopyTokenCnt, | ||
| 2152 | + uint32_t localExpertNum, uint32_t dataBufferId, bool isDataBufferReused) | ||
| 2153 | +{ | ||
| 2154 | + uint32_t tokenNum = 0U; | ||
| 2155 | + uint64_t readyMask = 0U; | ||
| 2156 | + const uint32_t packedTokenElems = blockCntPerToken_ * SPLIT_BLOCK_DATA_SIZE / sizeof(XOutType); | ||
| 2157 | + for (uint32_t index = beginIndex; index < endIndex; ++index) { | ||
| 2158 | + uint32_t copyCnt = expertLeftNumTensor_(index); | ||
| 2159 | + copyCnt = copyCnt > maxCopyTokenCnt - tokenNum ? maxCopyTokenCnt - tokenNum : copyCnt; | ||
| 2160 | + if (copyCnt == 0U) { | ||
| 2161 | + continue; | ||
| 2162 | + } | ||
| 2163 | + if (CheckDataArriveWithFlag(tokenNum, copyCnt)) { | ||
| 2164 | + if (readyMask == 0U && isDataBufferReused) { | ||
| 2165 | + AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(dataBufferId); | ||
| 2166 | + AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(dataBufferId); | ||
| 2167 | + AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(dataBufferId); | ||
| 2168 | + } | ||
| 2169 | + uint32_t source = GetLocalWindowSrcDataBlockIdx(expertMapTensor_(index), localExpertNum); | ||
| 2170 | + GlobalTensor<XOutType> dataGlobal; | ||
| 2171 | + dataGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ XOutType*>(windowGM_ + source * dataSectionSize_ + | ||
| 2172 | + expertFinishNumTensor_(index) * hCommuSize_)); | ||
| 2173 | + DataCopyParams copyParams{static_cast<uint16_t>(blockCntPerToken_ * copyCnt), | ||
| 2174 | + static_cast<uint16_t>(SPLIT_BLOCK_DATA_SIZE), static_cast<uint16_t>(UB_ALIGN), | ||
| 2175 | + 0U}; | ||
| 2176 | + DataCopyPadParams padParams{false, 0U, 0U, 0U}; | ||
| 2177 | + DataCopyPad(xTmpTensor[tokenNum * packedTokenElems], dataGlobal, copyParams, padParams); | ||
| 2178 | + readyMask |= 1ULL << (index - beginIndex); | ||
| 2179 | + } | ||
| 2180 | + tokenNum += copyCnt; | ||
| 2181 | + } | ||
| 2182 | + return readyMask; | ||
| 2183 | +} | ||
| 2184 | + | ||
| 2185 | +template <TemplateMC2A5FullMeshTypeClass> | ||
| 2186 | +__aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTypeFunc>::CopyOut( | ||
| 2187 | + LocalTensor<XOutType> xTmpTensor, uint32_t index, uint32_t dstPosition, uint32_t copyCnt, uint32_t dataBufferId) | ||
| 2108 | { | 2188 | { |
| 2109 | uint32_t hOutElemCount = hOutSize_ / sizeof(XOutType); // expandXOutGlobal申请每个token的GM Buffer空间大小 | 2189 | uint32_t hOutElemCount = hOutSize_ / sizeof(XOutType); // expandXOutGlobal申请每个token的GM Buffer空间大小 |
| 2110 | - GlobalTensor<XOutType> dataFlagGlobal, expandXOutGlobal; | 2190 | + GlobalTensor<XOutType> expandXOutGlobal; |
| 2111 | GlobalTensor<float> expertScaleInGlobal, expandScalesOutGlobal; | 2191 | GlobalTensor<float> expertScaleInGlobal, expandScalesOutGlobal; |
| 2112 | - dataFlagGlobal.SetGlobalBuffer((__gm__ XOutType*)(wAddr)); | ||
| 2113 | expandXOutGlobal.SetGlobalBuffer((__gm__ XOutType*)(expandXOutGM_) + | 2192 | expandXOutGlobal.SetGlobalBuffer((__gm__ XOutType*)(expandXOutGM_) + |
| 2114 | static_cast<uint64_t>(dstPosition) * hOutElemCount); | 2193 | static_cast<uint64_t>(dstPosition) * hOutElemCount); |
| 2115 | - DataCopyParams srcTokenCopyParams{static_cast<uint16_t>(blockCntPerToken_ * arriveCount), | ||
| 2116 | - static_cast<uint16_t>(SPLIT_BLOCK_DATA_SIZE), static_cast<uint16_t>(UB_ALIGN), 0}; | ||
| 2117 | DataCopyExtParams scalesCopyParams{ | 2194 | DataCopyExtParams scalesCopyParams{ |
| 2118 | - uint16_t(arriveCount), static_cast<uint32_t>(scaleOutBytes_), | 2195 | + uint16_t(copyCnt), static_cast<uint32_t>(scaleOutBytes_), |
| 2119 | static_cast<uint32_t>((blockCntPerToken_ * SPLIT_BLOCK_DATA_SIZE - scaleOutBytes_) / UB_ALIGN), 0U, 0U}; | 2196 | static_cast<uint32_t>((blockCntPerToken_ * SPLIT_BLOCK_DATA_SIZE - scaleOutBytes_) / UB_ALIGN), 0U, 0U}; |
| 2120 | DataCopyExtParams tokenCopyParams{ | 2197 | DataCopyExtParams tokenCopyParams{ |
| 2121 | - uint16_t(arriveCount), hOutSize_, | 2198 | + uint16_t(copyCnt), hOutSize_, |
| 2122 | static_cast<uint32_t>((blockCntPerToken_ * SPLIT_BLOCK_DATA_SIZE - hOutSize_) / UB_ALIGN), 0U, 0U}; | 2199 | static_cast<uint32_t>((blockCntPerToken_ * SPLIT_BLOCK_DATA_SIZE - hOutSize_) / UB_ALIGN), 0U, 0U}; |
| 2123 | DataCopyExtParams expandIdxCopyParams{ | 2200 | DataCopyExtParams expandIdxCopyParams{ |
| 2124 | - uint16_t(arriveCount), EXPAND_IDX_INFO * sizeof(int32_t), | 2201 | + uint16_t(copyCnt), EXPAND_IDX_INFO * sizeof(int32_t), |
| 2125 | static_cast<uint32_t>((blockCntPerToken_ * SPLIT_BLOCK_DATA_SIZE) / UB_ALIGN - 1), 0U, 0U}; | 2202 | static_cast<uint32_t>((blockCntPerToken_ * SPLIT_BLOCK_DATA_SIZE) / UB_ALIGN - 1), 0U, 0U}; |
| 2126 | - DataCopyPadParams srcTokenPadParams{false, 0U, 0U, 0U}; | ||
| 2127 | LocalTensor<int32_t> xOutInt32Tensor = xTmpTensor.template ReinterpretCast<int32_t>(); | 2203 | LocalTensor<int32_t> xOutInt32Tensor = xTmpTensor.template ReinterpretCast<int32_t>(); |
| 2128 | 2204 | ||
| 2129 | - if (isDataBufferReused) { | ||
| 2130 | - AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(dataBufferId); | ||
| 2131 | - AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(dataBufferId); | ||
| 2132 | - AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(dataBufferId); | ||
| 2133 | - } | ||
| 2134 | - DataCopyPad(xTmpTensor, dataFlagGlobal[expertFinishNumTensor_(index) * hCommuSize_ / sizeof(XOutType)], | ||
| 2135 | - srcTokenCopyParams, srcTokenPadParams); | ||
| 2136 | - AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE3>(dataBufferId); | ||
| 2137 | - AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE3>(dataBufferId); | ||
| 2138 | quantInst_.CopyScalesToOut(dstPosition, scaleOutBytes_, xTmpTensor, scalesCopyParams); | 2205 | quantInst_.CopyScalesToOut(dstPosition, scaleOutBytes_, xTmpTensor, scalesCopyParams); |
| 2139 | DataCopyPad(expandXOutGlobal, xTmpTensor, tokenCopyParams); | 2206 | DataCopyPad(expandXOutGlobal, xTmpTensor, tokenCopyParams); |
| 2140 | DataCopyPad(expandIdxGMTensor_[dstPosition * EXPAND_IDX_INFO], xOutInt32Tensor[tokenQuantAlign_], | 2207 | DataCopyPad(expandIdxGMTensor_[dstPosition * EXPAND_IDX_INFO], xOutInt32Tensor[tokenQuantAlign_], |
| 2141 | expandIdxCopyParams); | 2208 | expandIdxCopyParams); |
| 2142 | if (hasExpertScalesFlag_) { | 2209 | if (hasExpertScalesFlag_) { |
| 2210 | + uint32_t localExpertNum = isShareExpertRankFlag_ ? 1U : moeExpertNumPerRank_; | ||
| 2211 | + uint32_t source = GetLocalWindowSrcDataBlockIdx(expertMapTensor_(index), localExpertNum); | ||
| 2212 | + GM_ADDR wAddr = windowGM_ + source * dataSectionSize_; | ||
| 2143 | LocalTensor<float> xOutFloatTensor = xTmpTensor.template ReinterpretCast<float>(); | 2213 | LocalTensor<float> xOutFloatTensor = xTmpTensor.template ReinterpretCast<float>(); |
| 2144 | uint32_t expertScaleOffset = expertScaleAlign_ * sizeof(float); | 2214 | uint32_t expertScaleOffset = expertScaleAlign_ * sizeof(float); |
| 2145 | uint32_t expertScaleWinOffset = | 2215 | uint32_t expertScaleWinOffset = |
| @@ -2147,11 +2217,10 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy | |||
| 2147 | expertScaleInGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float*>( | 2217 | expertScaleInGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float*>( |
| 2148 | wAddr + expertFinishNumTensor_(index) * hCommuSize_ + expertScaleWinOffset)); | 2218 | wAddr + expertFinishNumTensor_(index) * hCommuSize_ + expertScaleWinOffset)); |
| 2149 | expandScalesOutGlobal.SetGlobalBuffer((__gm__ float*)(expandScalesOutGM_) + dstPosition); | 2219 | expandScalesOutGlobal.SetGlobalBuffer((__gm__ float*)(expandScalesOutGM_) + dstPosition); |
| 2150 | - DataCopyExtParams expertScaleCopyInParams{uint16_t(arriveCount), static_cast<uint32_t>(sizeof(float)), | 2220 | + DataCopyExtParams expertScaleCopyInParams{uint16_t(copyCnt), static_cast<uint32_t>(sizeof(float)), |
| 2151 | static_cast<uint32_t>(hCommuSize_ - sizeof(float)), 0U, 0U}; | 2221 | static_cast<uint32_t>(hCommuSize_ - sizeof(float)), 0U, 0U}; |
| 2152 | DataCopyPadExtParams<float> expertScalePadParams{false, 0U, 0U, 0.0f}; | 2222 | DataCopyPadExtParams<float> expertScalePadParams{false, 0U, 0U, 0.0f}; |
| 2153 | - DataCopyExtParams expertScaleCopyOutParams{uint16_t(arriveCount), static_cast<uint32_t>(sizeof(float)), 0U, 0U, | 2223 | + DataCopyExtParams expertScaleCopyOutParams{uint16_t(copyCnt), static_cast<uint32_t>(sizeof(float)), 0U, 0U, 0U}; |
| 2154 | - 0U}; | ||
| 2155 | AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(dataBufferId); | 2224 | AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(dataBufferId); |
| 2156 | AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(dataBufferId); | 2225 | AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(dataBufferId); |
| 2157 | DataCopyPad(xOutFloatTensor, expertScaleInGlobal, expertScaleCopyInParams, expertScalePadParams); | 2226 | DataCopyPad(xOutFloatTensor, expertScaleInGlobal, expertScaleCopyInParams, expertScalePadParams); |
| @@ -2159,7 +2228,6 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy | |||
| 2159 | AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE3>(dataBufferId); | 2228 | AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE3>(dataBufferId); |
| 2160 | DataCopyPad(expandScalesOutGlobal, xOutFloatTensor, expertScaleCopyOutParams); | 2229 | DataCopyPad(expandScalesOutGlobal, xOutFloatTensor, expertScaleCopyOutParams); |
| 2161 | } | 2230 | } |
| 2162 | - AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(dataBufferId); | ||
| 2163 | } | 2231 | } |
| 2164 | 2232 | ||
| 2165 | template <TemplateMC2A5FullMeshTypeClass> | 2233 | template <TemplateMC2A5FullMeshTypeClass> |
| @@ -2223,72 +2291,63 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy | |||
| 2223 | uint32_t index = 0; | 2291 | uint32_t index = 0; |
| 2224 | uint32_t finishNum = 0; | 2292 | uint32_t finishNum = 0; |
| 2225 | uint32_t localExpertNum = isShareExpertRankFlag_ ? 1 : moeExpertNumPerRank_; | 2293 | uint32_t localExpertNum = isShareExpertRankFlag_ ? 1 : moeExpertNumPerRank_; |
| 2226 | - uint32_t srcExpertId, dstPosition, arriveCount, copyCnt, srcDataBlockIdx; | 2294 | + uint32_t srcExpertId, dstPosition, copyCnt; |
| 2227 | uint64_t packedTokenSize = static_cast<uint64_t>(blockCntPerToken_) * SPLIT_BLOCK_DATA_SIZE; | 2295 | uint64_t packedTokenSize = static_cast<uint64_t>(blockCntPerToken_) * SPLIT_BLOCK_DATA_SIZE; |
| 2228 | - uint64_t minUbSizePerToken = | 2296 | + uint64_t minUbSizePerToken = LOCAL_COPY_BUFFER_NUM * packedTokenSize + |
| 2229 | - LOCAL_COPY_BUFFER_NUM * packedTokenSize + static_cast<uint64_t>(blockCntPerToken_) * (UB_ALIGN + sizeof(float)); | 2297 | + static_cast<uint64_t>(blockCntPerToken_) * (UB_ALIGN + sizeof(uint8_t)); |
| 2298 | + // 总预算按32B对齐，位图取整不改变可容纳的token数 | ||
| 2230 | uint32_t maxCopyTokenCnt = static_cast<uint32_t>(tBufRealSize_ / minUbSizePerToken); | 2299 | uint32_t maxCopyTokenCnt = static_cast<uint32_t>(tBufRealSize_ / minUbSizePerToken); |
| 2231 | - uint32_t dataBufferSize = 0; | 2300 | + uint32_t dataBufferSize = static_cast<uint32_t>(packedTokenSize * maxCopyTokenCnt); |
| 2232 | - uint32_t reduceWorkSize = 0; | 2301 | + uint32_t flagReadySize = Ceil(blockCntPerToken_ * maxCopyTokenCnt, UB_ALIGN) * UB_ALIGN; |
C 【review】 问题分类：安全 位置：mc2/moe_distribute_dispatch_v2/op_kernel/arch35/moe_distribute_dispatch_v2_a5_full_mesh.h:2301，代码：uint32_t flagReadySize = Ceil(blockCntPerToken_ * maxCopyTokenCnt, UB_ALIGN) * UB_ALIGN; 问题：maxCopyTokenCnt 的预算在上一行只按每个 token 增加 sizeof(uint8_t)，但 flagReadySize 实际会把所有 token 的位图一次性向 UB_ALIGN 向上取整；同时旧实现用于回算 totalBufferSize 并递减 maxCopyTokenCnt 的容量校验已删除。 可能导致的结果：当 blockCntPerToken_ * maxCopyTokenCnt 不是 UB_ALIGN 的整数倍且 TBuf 剩余空间小于补齐量时，flagReadyTensor_ 会与 flagRecvTensor_ 或后续 UB 区域重叠，CompareScalar 写入越过 TBuf 分配范围，触发非法内存访问并破坏同一 kernel 的其他临时数据。 建议的修改方法：按 flagReadySize 的实际向上取整结果计算总占用，并在确定 maxCopyTokenCnt 后验证 2 * dataBufferSize + flagReadySize + flagRecvSize <= tBufRealSize_；若不满足则递减批大小后再初始化各 LocalTensor。   zhong-zixin 6 days ago Comments： 6 days ago Comments： | |||
| 2233 | - uint32_t flagRecvSize = 0; | 2302 | + uint32_t flagRecvSize = blockCntPerToken_ * maxCopyTokenCnt * UB_ALIGN; |
| 2234 | - while (maxCopyTokenCnt > 0) { | ||
| 2235 | - dataBufferSize = static_cast<uint32_t>(packedTokenSize * maxCopyTokenCnt); | ||
| 2236 | - reduceWorkSize = Ceil(blockCntPerToken_ * maxCopyTokenCnt * sizeof(float), SIZE_ALIGN_256) * SIZE_ALIGN_256; | ||
| 2237 | - flagRecvSize = blockCntPerToken_ * maxCopyTokenCnt * UB_ALIGN; | ||
| 2238 | - uint64_t totalBufferSize = | ||
| 2239 | - static_cast<uint64_t>(LOCAL_COPY_BUFFER_NUM) * dataBufferSize + reduceWorkSize + flagRecvSize; | ||
| 2240 | - if (totalBufferSize <= tBufRealSize_) { | ||
| 2241 | - break; | ||
| 2242 | - } | ||
| 2243 | - --maxCopyTokenCnt; | ||
| 2244 | - } | ||
| 2245 | - if (maxCopyTokenCnt == 0) { | ||
| 2246 | - return; | ||
| 2247 | - } | ||
| 2248 | 2303 | ||
| 2249 | LocalTensor<XOutType> xTmpPingTensor = tBuf.GetWithOffset<XOutType>(dataBufferSize / sizeof(XOutType), 0); | 2304 | LocalTensor<XOutType> xTmpPingTensor = tBuf.GetWithOffset<XOutType>(dataBufferSize / sizeof(XOutType), 0); |
| 2250 | LocalTensor<XOutType> xTmpPongTensor = | 2305 | LocalTensor<XOutType> xTmpPongTensor = |
| 2251 | tBuf.GetWithOffset<XOutType>(dataBufferSize / sizeof(XOutType), dataBufferSize); | 2306 | tBuf.GetWithOffset<XOutType>(dataBufferSize / sizeof(XOutType), dataBufferSize); |
| 2252 | uint32_t flagWorkOffset = LOCAL_COPY_BUFFER_NUM * dataBufferSize; | 2307 | uint32_t flagWorkOffset = LOCAL_COPY_BUFFER_NUM * dataBufferSize; |
| 2253 | - flagReduceWorkTensor_ = tBuf.GetWithOffset<float>(reduceWorkSize / sizeof(float), flagWorkOffset); | 2308 | + flagReadyTensor_ = tBuf.GetWithOffset<uint64_t>(flagReadySize / sizeof(uint64_t), flagWorkOffset); |
| 2254 | - flagRecvTensor_ = tBuf.GetWithOffset<float>(flagRecvSize / sizeof(float), flagWorkOffset + reduceWorkSize); | 2309 | + flagRecvTensor_ = tBuf.GetWithOffset<float>(flagRecvSize / sizeof(float), flagWorkOffset + flagReadySize); |
| 2255 | 2310 | ||
| 2256 | uint32_t copyBatchId = 0; | 2311 | uint32_t copyBatchId = 0; |
| 2257 | - while (true) { | 2312 | + while (validNum != finishNum) { |
| 2258 | - if (expertLeftNumTensor_(index) == 0) { // 当前核负责的不需要收集 | 2313 | + uint32_t beginIndex = index; |
| 2259 | - index = (index + 1) % validNum; // 轮询查询每个有效的index | 2314 | + uint32_t endIndex = LoadDataFlags(index, validNum, maxCopyTokenCnt, localExpertNum); |
| 2315 | + uint32_t dataBufferId = copyBatchId % LOCAL_COPY_BUFFER_NUM; | ||
| 2316 | + LocalTensor<XOutType> xTmpTensor = dataBufferId == 0U ? xTmpPingTensor : xTmpPongTensor; | ||
| 2317 | + uint64_t readyMask = LoadReadyData(xTmpTensor, index, endIndex, maxCopyTokenCnt, localExpertNum, dataBufferId, | ||
| 2318 | + copyBatchId >= LOCAL_COPY_BUFFER_NUM); | ||
| 2319 | + if (readyMask == 0U) { | ||
| 2320 | + index = endIndex == validNum ? 0U : endIndex; | ||
| 2260 | continue; | 2321 | continue; |
| 2261 | } | 2322 | } |
| 2262 | - srcExpertId = expertMapTensor_(index); | 2323 | + // 已就绪片段先集中搬入，再统一等待搬入完成 |
| 2263 | - copyCnt = expertLeftNumTensor_(index) > maxCopyTokenCnt ? | 2324 | + AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE3>(dataBufferId); |
| 2264 | - maxCopyTokenCnt : | 2325 | + AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE3>(dataBufferId); |
| 2265 | - expertLeftNumTensor_(index); // 按照ub大小一次搬入多个token | 2326 | + uint32_t tokenNum = 0U; |
| 2266 | - srcDataBlockIdx = GetLocalWindowSrcDataBlockIdx(srcExpertId, localExpertNum); | 2327 | + for (; index < endIndex; ++index) { |
| 2267 | - arriveCount = CheckDataArriveWithFlag(srcDataBlockIdx, expertFinishNumTensor_(index), copyCnt); | 2328 | + if (expertLeftNumTensor_(index) == 0) { // 当前核负责的不需要收集 |
| 2268 | - if (arriveCount == copyCnt) { | 2329 | + continue; |
| 2269 | - dstPosition = srcExpertId != 0 ? GetRecvCntPrefix(srcExpertId - 1U) : 0; | ||
| 2270 | - dstPosition += expertFinishNumTensor_(index); | ||
| 2271 | - GM_ADDR wAddr = (__gm__ uint8_t*)(windowGM_) + srcDataBlockIdx * dataSectionSize_; | ||
| 2272 | - uint32_t dataBufferId = copyBatchId % LOCAL_COPY_BUFFER_NUM; | ||
| 2273 | - LocalTensor<XOutType> xTmpTensor = xTmpPingTensor; | ||
| 2274 | - if (dataBufferId != 0) { | ||
| 2275 | - xTmpTensor = xTmpPongTensor; | ||
| 2276 | } | 2330 | } |
| 2277 | - CopyInAndOut(xTmpTensor, wAddr, index, dstPosition, arriveCount, dataBufferId, | 2331 | + copyCnt = expertLeftNumTensor_(index); |
| 2278 | - copyBatchId >= LOCAL_COPY_BUFFER_NUM); | 2332 | + copyCnt = copyCnt > maxCopyTokenCnt - tokenNum ? maxCopyTokenCnt - tokenNum : copyCnt; |
| 2279 | - ++copyBatchId; | 2333 | + if ((readyMask & (1ULL << (index - beginIndex))) != 0U) { |
| 2280 | - // Update progress here; data flags are cleared in batches after all outputs are issued. | 2334 | + srcExpertId = expertMapTensor_(index); |
| 2281 | - expertFinishNumTensor_(index) += arriveCount; | 2335 | + dstPosition = srcExpertId != 0 ? GetRecvCntPrefix(srcExpertId - 1U) : 0; |
| 2282 | - expertLeftNumTensor_(index) -= arriveCount; | 2336 | + dstPosition += expertFinishNumTensor_(index); |
| 2283 | - if (expertLeftNumTensor_(index) == 0) { | 2337 | + CopyOut(xTmpTensor[tokenNum * packedTokenSize / sizeof(XOutType)], index, dstPosition, copyCnt, |
| 2284 | - finishNum++; | 2338 | + dataBufferId); |
| 2339 | + // 更新进度，全部输出下发后统一清理标志 | ||
| 2340 | + expertFinishNumTensor_(index) += copyCnt; | ||
| 2341 | + expertLeftNumTensor_(index) -= copyCnt; | ||
| 2342 | + if (expertLeftNumTensor_(index) == 0) { | ||
| 2343 | + finishNum++; | ||
| 2344 | + } | ||
| 2285 | } | 2345 | } |
| 2286 | - } else { | 2346 | + tokenNum += copyCnt; |
| 2287 | - index = (index + 1) % validNum; | ||
| 2288 | - } | ||
| 2289 | - if (validNum == finishNum) { | ||
| 2290 | - break; | ||
| 2291 | } | 2347 | } |
| 2348 | + AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(dataBufferId); | ||
| 2349 | + ++copyBatchId; | ||
| 2350 | + index = index == validNum ? 0U : index; | ||
| 2292 | } | 2351 | } |
| 2293 | 2352 | ||
| 2294 | uint32_t usedDataBufferNum = copyBatchId > LOCAL_COPY_BUFFER_NUM ? LOCAL_COPY_BUFFER_NUM : copyBatchId; | 2353 | uint32_t usedDataBufferNum = copyBatchId > LOCAL_COPY_BUFFER_NUM ? LOCAL_COPY_BUFFER_NUM : copyBatchId; |
| @@ -2452,11 +2511,10 @@ __aicore__ inline void MoeDistributeDispatchV2A5FullMesh<TemplateMC2A5FullMeshTy | |||
| 2452 | uint32_t expertNum = endLocalExpertId - startLocalExpertId; | 2511 | uint32_t expertNum = endLocalExpertId - startLocalExpertId; |
| 2453 | uint32_t statusNum = expertNum * epWorldSize_; | 2512 | uint32_t statusNum = expertNum * epWorldSize_; |
| 2454 | LocalTensor<float> countTensorFp32 = recvCountTensor_.ReinterpretCast<float>(); | 2513 | LocalTensor<float> countTensorFp32 = recvCountTensor_.ReinterpretCast<float>(); |
| 2455 | - DataCopyParams fullCountCopyParams{static_cast<uint16_t>(statusNum), 1, 0, 0}; | 2514 | + DataCopyParams fullCountCopyParams{1, static_cast<uint16_t>(statusNum), 0, 0}; |
| 2456 | DataCopy(countTensorFp32, | 2515 | DataCopy(countTensorFp32, |
| 2457 | windowInstatusFp32Tensor_[startLocalExpertId * epWorldSize_ * STATE_OFFSET / sizeof(float)], | 2516 | windowInstatusFp32Tensor_[startLocalExpertId * epWorldSize_ * STATE_OFFSET / sizeof(float)], |
| 2458 | fullCountCopyParams); | 2517 | fullCountCopyParams); |
| 2459 | - SyncFunc<AscendC::HardEvent::MTE2_V>(); | ||
| 2460 | SyncFunc<AscendC::HardEvent::MTE2_S>(); | 2518 | SyncFunc<AscendC::HardEvent::MTE2_S>(); |
| 2461 | LocalTensor<int32_t> statusBlocksTensor = countTensorFp32.ReinterpretCast<int32_t>(); | 2519 | LocalTensor<int32_t> statusBlocksTensor = countTensorFp32.ReinterpretCast<int32_t>(); |
| 2462 | LocalTensor<int64_t> expertTokenNumsLocalTensor = tokenNumBuf_.Get<int64_t>(); | 2520 | LocalTensor<int64_t> expertTokenNumsLocalTensor = tokenNumBuf_.Get<int64_t>(); |
【review】 问题分类：正确性 位置：mc2/moe_distribute_dispatch_v2/op_kernel/arch35/moe_distribute_dispatch_v2_a5_full_mesh.h:2140，代码：uint64_t flags = flagReadyTensor_.GetValue(beginFlag / flagsPerWord) >> (offset * BITS_PER_BYTE); 问题：CompareScalar 按每个浮点元素写入一个字节结果，而 beginFlag 表示 32B 标志块序号；当前索引没有乘以 UB_ALIGN_DATA_COUNT，导致第二个及后续标志块读取的是前一块结果中的填充字节，而不是对应 32B 块的首个比较结果。 可能导致的结果：当一个 token 或一批 token 包含多个 32B 标志块时，即使后续块已经就绪，检查也会读到填充的 0 值并判定未就绪，持续跳过已到达数据，造成输出长期等待或无法完成。 建议的修改方法：按 CompareScalar 结果的实际布局，将 beginFlag/offset 换算为每个 32B 块对应的 UB_ALIGN_DATA_COUNT 个字节（或在生成比较结果时直接压缩为每块一个字节），再按块序号构造 readyMask。