Star1115
2707
代码
代码
Issues205
Pull Requests371
Actions
项目讨论
Wiki
分析
Star1115
2707
已合并
修改关于nsaxx/SparseFlashMlaGrad/SparseLightningIndexerKLLossGrad算子的数值溢出问题 #6458
memorecool创建于 6月5日
修改关于nsaxx/SparseFlashMlaGrad/SparseLightningIndexerKLLossGrad算子的数值溢出问题 #6458
已合并
memorecool创建于 6月5日
从已删除 :fix合入到cann/ops-transformermaster
共 19 个文件变更+155-89
Mattention/nsa_compress/op_host/nsa_compress_tiling.h+3-3
@@ -24,9 +24,9 @@ const size_t MAX_CORE_NUM = 48; // 使用的CORE数
24 24 
25BEGIN_TILING_DATA_DEF(NsaCompressTilingData)25BEGIN_TILING_DATA_DEF(NsaCompressTilingData)
26TILING_DATA_FIELD_DEF(uint32_t, InputDtype); // 输入kv的总大小，初始化时指定kvGM的大小26TILING_DATA_FIELD_DEF(uint32_t, InputDtype); // 输入kv的总大小，初始化时指定kvGM的大小
27-TILING_DATA_FIELD_DEF(uint32_t, TotalKvSize); // 输入kv的总大小，初始化时指定kvGM的大小27+TILING_DATA_FIELD_DEF(uint64_t, TotalKvSize); // 输入kv的总大小，初始化时指定kvGM的大小
28-TILING_DATA_FIELD_DEF(uint32_t, TotalCompressSize); // 输出compress_kv的总大小，初始化时指定compressKvGM的大小28+TILING_DATA_FIELD_DEF(uint64_t, TotalCompressSize); // 输出compress_kv的总大小，初始化时指定compressKvGM的大小
29-TILING_DATA_FIELD_DEF(uint32_t, WeightSize); // 输入weight的总大小，初始化时指定weightGM的大小29+TILING_DATA_FIELD_DEF(uint64_t, WeightSize); // 输入weight的总大小，初始化时指定weightGM的大小
30TILING_DATA_FIELD_DEF(uint32_t, BatchSize); // 输入Kv的batch数，用于指定actseqlenGM的大小30TILING_DATA_FIELD_DEF(uint32_t, BatchSize); // 输入Kv的batch数，用于指定actseqlenGM的大小
31TILING_DATA_FIELD_DEF(uint32_t, CompressBlockSize); // 属性CompressBlockSize值31TILING_DATA_FIELD_DEF(uint32_t, CompressBlockSize); // 属性CompressBlockSize值
32TILING_DATA_FIELD_DEF(uint32_t, CompressStride); // 属性CompressStride值32TILING_DATA_FIELD_DEF(uint32_t, CompressStride); // 属性CompressStride值
Mattention/nsa_compress/op_host/nsa_compress_tiling_general.cpp+10-10
@@ -65,8 +65,8 @@ protected:
65 65 
66 headNum = inputShape->GetStorageShape().GetDim(One);66 headNum = inputShape->GetStorageShape().GetDim(One);
67 headDim = inputShape->GetStorageShape().GetDim(Two);67 headDim = inputShape->GetStorageShape().GetDim(Two);
68- headNumDim = headNum * headDim;68+ headNumDim = static_cast<uint64_t>(headNum) * headDim;
69- totalKvSize = inputShape->GetStorageShape().GetDim(0) * headNumDim;69+ totalKvSize = static_cast<uint64_t>(inputShape->GetStorageShape().GetDim(0)) * headNumDim;
70 return true;70 return true;
71 }71 }
72 72 
@@ -84,7 +84,7 @@ protected:
84 compressStride = static_cast<uint32_t>(*compressStridePtr);84 compressStride = static_cast<uint32_t>(*compressStridePtr);
85 actSeqLenType = *actSeqLenTypePtr;85 actSeqLenType = *actSeqLenTypePtr;
86 86 
87- weightSize = compressBlockSize * headNum;87+ weightSize = static_cast<uint64_t>(compressBlockSize) * headNum;
88 maxOverlap = 2u * (compressBlockSize + compressStride - 1u) / compressStride - 1u;88 maxOverlap = 2u * (compressBlockSize + compressStride - 1u) / compressStride - 1u;
89 return true;89 return true;
90 }90 }
@@ -148,7 +148,7 @@ protected:
148 actualOutputSeqLen.push_back(totalOutSeqLen);148 actualOutputSeqLen.push_back(totalOutSeqLen);
149 pre_seq_len += cur_seq_len;149 pre_seq_len += cur_seq_len;
150 }150 }
151- totalCompressSize = totalOutSeqLen * headNumDim;151+ totalCompressSize = static_cast<uint64_t>(totalOutSeqLen) * headNumDim;
152 return true;152 return true;
153 }153 }
154 154 
@@ -406,12 +406,12 @@ protected:
406 uint32_t compressStride;406 uint32_t compressStride;
407 uint32_t headNum;407 uint32_t headNum;
408 uint32_t headDim;408 uint32_t headDim;
409- uint32_t headNumDim;409+ uint64_t headNumDim;
410 410 
411- uint32_t totalKvSize; // 输入的总元素数量411+ uint64_t totalKvSize; // 输入的总元素数量
412- uint32_t weightSize; // weight的总元素数量412+ uint64_t weightSize; // weight的总元素数量
413- uint32_t totalOutSeqLen; // 输出的总seq_len数量413+ uint64_t totalOutSeqLen; // 输出的总seq_len数量
414- uint32_t totalCompressSize; // 输出的总元素数量414+ uint64_t totalCompressSize; // 输出的总元素数量
415 415 
416 uint32_t maxOverlap; // 当前compressBlockSize和compressStride下,中间结果overlap的数量416 uint32_t maxOverlap; // 当前compressBlockSize和compressStride下,中间结果overlap的数量
417 uint32_t maxCopyKVTokensNums; // 计算搬运至UB块的kv tokens数量417 uint32_t maxCopyKVTokensNums; // 计算搬运至UB块的kv tokens数量
@@ -435,7 +435,7 @@ protected:
435 uint32_t ubBlockFloatNum; // 当前计算的Dtype下(fp32),对齐32B所需的元素个数435 uint32_t ubBlockFloatNum; // 当前计算的Dtype下(fp32),对齐32B所需的元素个数
436 436 
437 std::vector<uint32_t> actualSeqLen; // 各batch的seq_len长度，前缀和方式437 std::vector<uint32_t> actualSeqLen; // 各batch的seq_len长度，前缀和方式
438- std::vector<uint32_t> actualOutputSeqLen; // 各batch可以输出的token数，前缀和方式438+ std::vector<uint64_t> actualOutputSeqLen; // 各batch可以输出的token数，前缀和方式
439 std::vector<uint32_t> divisorNums; // aivNum的因子439 std::vector<uint32_t> divisorNums; // aivNum的因子
440 uint32_t PerCoreHeadIdx[MAX_CORE_NUM] = {0}; // 记录每个核处理的head索引440 uint32_t PerCoreHeadIdx[MAX_CORE_NUM] = {0}; // 记录每个核处理的head索引
441 uint32_t PerCoreHeadNum[MAX_CORE_NUM] = {0}; // 记录每个核处理的head数量441 uint32_t PerCoreHeadNum[MAX_CORE_NUM] = {0}; // 记录每个核处理的head数量
Mattention/nsa_compress/op_kernel/nsa_compress.h+9-7
@@ -89,8 +89,8 @@ public:
89 subTiling.subSeqLen = tilingData.BlocksNums;89 subTiling.subSeqLen = tilingData.BlocksNums;
90 subTiling.subHeadDim = tiling.headDim;90 subTiling.subHeadDim = tiling.headDim;
91 subTiling.subHeadNum = coreInfo.coreHeadNums;91 subTiling.subHeadNum = coreInfo.coreHeadNums;
92- subTiling.subKvSize = subTiling.subSeqLen * coreInfo.coreHeadNums * tiling.headDim;92+ subTiling.subKvSize = static_cast<uint64_t>(subTiling.subSeqLen) * coreInfo.coreHeadNums * tiling.headDim;
93- subTiling.subCompressKvSize = coreInfo.coreHeadNums * tiling.headDim;93+ subTiling.subCompressKvSize = static_cast<uint64_t>(coreInfo.coreHeadNums) * tiling.headDim;
94 }94 }
95 }95 }
96 96 
@@ -347,11 +347,12 @@ __aicore__ inline void KernelNASCompress<T>::Init(__gm__ uint8_t *input, __gm__
347}347}
348 348 
349 349 
350-template <typename T> 350+template <typename T>
351__aicore__ inline void KernelNASCompress<T>::CopyIn()351__aicore__ inline void KernelNASCompress<T>::CopyIn()
352{352{
353- uint32_t block_start =353+ uint64_t block_start =
354- coreInfo.coreSeqIdx * tiling.headNum * tiling.headDim + coreInfo.coreHeadIdx * tiling.headDim;354+ static_cast<uint64_t>(coreInfo.coreSeqIdx) * tiling.headNum * tiling.headDim +
355+ static_cast<uint64_t>(coreInfo.coreHeadIdx) * tiling.headDim;
355 356 
356 AscendC::LocalTensor<T> kvCacheLocal = inQueueKvFp16.AllocTensor<T>();357 AscendC::LocalTensor<T> kvCacheLocal = inQueueKvFp16.AllocTensor<T>();
357 358 
@@ -499,8 +500,9 @@ __aicore__ inline void KernelNASCompress<T>::CopyOut()
499 // 成功压缩一个compress token. 数据UB->GM500 // 成功压缩一个compress token. 数据UB->GM
500 AscendC::LocalTensor<T> compressKvCacheLocal = outQueCompressKv.DeQue<T>();501 AscendC::LocalTensor<T> compressKvCacheLocal = outQueCompressKv.DeQue<T>();
501 if (coreInfo.coreCompressIdx < coreInfo.coreCompressNum) {502 if (coreInfo.coreCompressIdx < coreInfo.coreCompressNum) {
502- uint32_t compressOffset =503+ uint64_t compressOffset =
503- coreInfo.coreCompressOffset + coreInfo.coreCompressIdx * tiling.headNum * tiling.headDim;504+ coreInfo.coreCompressOffset +
505+ static_cast<uint64_t>(coreInfo.coreCompressIdx) * tiling.headNum * tiling.headDim;
504 AscendC::PipeBarrier<PIPE_ALL>();506 AscendC::PipeBarrier<PIPE_ALL>();
505 AscendC::DataCopy(compressKvGm[compressOffset], compressKvCacheLocal, coreInfo.coreCompressSize);507 AscendC::DataCopy(compressKvGm[compressOffset], compressKvCacheLocal, coreInfo.coreCompressSize);
506 AscendC::PipeBarrier<PIPE_ALL>();508 AscendC::PipeBarrier<PIPE_ALL>();
Mattention/nsa_compress/op_kernel/nsa_compress_seq_manager.h+3-3
@@ -193,14 +193,14 @@ struct SingleCompressionTokenMetadata_st {
193 return;193 return;
194 }194 }
195 195 
196- __aicore__ inline uint32_t _GetWeightOffset(WeightOffsetType type, SubTilingInfo *subtiling = nullptr)196+ __aicore__ inline uint64_t _GetWeightOffset(WeightOffsetType type, SubTilingInfo *subtiling = nullptr)
197 {197 {
198 if (type == WeightOffsetType::BUFFSET_OFFSET && subtiling == nullptr) {198 if (type == WeightOffsetType::BUFFSET_OFFSET && subtiling == nullptr) {
199 return WEIGHT_OFFSET_INVALID;199 return WEIGHT_OFFSET_INVALID;
200 }200 }
201 if (type == WeightOffsetType::BUFFSET_OFFSET && subtiling) {201 if (type == WeightOffsetType::BUFFSET_OFFSET && subtiling) {
202 uint32_t weight_dim = 8;202 uint32_t weight_dim = 8;
203- return weightOffset * subtiling->subHeadNum * weight_dim;203+ return static_cast<uint64_t>(weightOffset) * subtiling->subHeadNum * weight_dim;
204 }204 }
205 if (type == WeightOffsetType::TOKEN_OFFSET) {205 if (type == WeightOffsetType::TOKEN_OFFSET) {
206 return weightOffset;206 return weightOffset;
@@ -335,7 +335,7 @@ struct SequenceManager {
335 335 
336 /// @brief336 /// @brief
337 /// @return337 /// @return
338- __aicore__ inline uint32_t GetWeightOffset(WeightOffsetType offsetType, SubTilingInfo *subtiling = nullptr)338+ __aicore__ inline uint64_t GetWeightOffset(WeightOffsetType offsetType, SubTilingInfo *subtiling = nullptr)
339 {339 {
340 return compressMeta._GetWeightOffset(offsetType, subtiling);340 return compressMeta._GetWeightOffset(offsetType, subtiling);
341 }341 }
Mattention/nsa_compress_attention/op_host/nsa_compress_attention_tiling_general.cpp+13-4
@@ -950,10 +950,19 @@ protected:
950 {950 {
951 size_t *workspaces = context_->GetWorkspaceSizes(1);951 size_t *workspaces = context_->GetWorkspaceSizes(1);
952 int64_t mm1BaseCount = NSA_BASE_S1G_SIZE * alignedS2;952 int64_t mm1BaseCount = NSA_BASE_S1G_SIZE * alignedS2;
953- size_t pWorkspaceSize = static_cast<size_t>(mm1BaseCount * sizeof(float) * NSA_DOUBLE_BUFFER * aicNum); //alignedS2是S2最大值953+ 
954- size_t stage1ResWorkspaceSize = static_cast<size_t>(mm1BaseCount * sizeof(inputDtype) * aicNum);954+ uint64_t pWorkspaceSize64 =
955- size_t impScoreWorkspaceSize = static_cast<size_t>((NSA_BASE_S1G_SIZE / gSize) * CeilDiv(static_cast<uint64_t>(alignedS2), tilingData.importanceScoreParams.get_isM()) * sizeof(float) * NSA_DOUBLE_BUFFER * aicNum);955+ static_cast<uint64_t>(mm1BaseCount) * sizeof(float) * NSA_DOUBLE_BUFFER * static_cast<uint64_t>(aicNum);
956- workspaces[0] = pWorkspaceSize + impScoreWorkspaceSize + stage1ResWorkspaceSize + WORK_SPACE_RESERVE_SIZE;956+ uint64_t stage1ResWorkspaceSize64 =
957+ static_cast<uint64_t>(mm1BaseCount) * sizeof(inputDtype) * static_cast<uint64_t>(aicNum);
958+ uint64_t baseS1G = static_cast<uint64_t>(NSA_BASE_S1G_SIZE) / static_cast<uint64_t>(gSize);
959+ uint64_t alignedS2Div = CeilDiv(static_cast<uint64_t>(alignedS2), tilingData.importanceScoreParams.get_isM());
960+ uint64_t impScoreWorkspaceSize64 =
961+ baseS1G * alignedS2Div * sizeof(float) * NSA_DOUBLE_BUFFER * static_cast<uint64_t>(aicNum);
962+ 
963+ uint64_t totalWorkspace64 =
964+ pWorkspaceSize64 + impScoreWorkspaceSize64 + stage1ResWorkspaceSize64 + WORK_SPACE_RESERVE_SIZE;
965+ workspaces[0] = static_cast<size_t>(totalWorkspace64);
957 return ge::GRAPH_SUCCESS;966 return ge::GRAPH_SUCCESS;
958 }967 }
959 968 
Mattention/nsa_compress_attention/op_kernel/nsa_compress_attention_s1s2_bn2gs1_sab.h+50-19
@@ -281,14 +281,17 @@ NsaCompressAttentionS1s2Bn2gs1SameAB<layOutType, hasAtten, hasTopkMask, INPUT_T,
281 281 
282 // impScore282 // impScore
283 int64_t totalAicNum = AscendC::GetBlockNum();283 int64_t totalAicNum = AscendC::GetBlockNum();
284- int64_t impScoreOffset = totalOffset * totalAicNum;284+ uint64_t impScoreOffsetU64 = static_cast<uint64_t>(totalOffset) * static_cast<uint64_t>(totalAicNum);
285- int64_t perCoreOffset = (NSA_BASE_S1G_SIZE / this->tilingData->inputParams.gSize)285+ uint64_t baseS1G =
286- * (this->tilingData->inputParams.alignedS2 / this->tilingData->importanceScoreParams.isM)286+ static_cast<uint64_t>(NSA_BASE_S1G_SIZE) / static_cast<uint64_t>(this->tilingData->inputParams.gSize);
287- * sizeof(float) * GM_DOUBLE_BUFFER;287+ uint64_t alignedS2Div = static_cast<uint64_t>(this->tilingData->inputParams.alignedS2) /
288+ static_cast<uint64_t>(this->tilingData->importanceScoreParams.isM);
289+ uint64_t perCoreOffsetU64 = baseS1G * alignedS2Div * sizeof(float) * GM_DOUBLE_BUFFER;
288 this->impScoreRes[0].SetGlobalBuffer(290 this->impScoreRes[0].SetGlobalBuffer(
289- (__gm__ T *)(workspace + impScoreOffset + perCoreOffset * this->cubeBlockIdx));291+ (__gm__ T *)(workspace + impScoreOffsetU64 + perCoreOffsetU64 * static_cast<uint64_t>(this->cubeBlockIdx)));
290 this->impScoreRes[1].SetGlobalBuffer(292 this->impScoreRes[1].SetGlobalBuffer(
291- (__gm__ T *)(workspace + impScoreOffset + perCoreOffset * this->cubeBlockIdx + perCoreOffset / 2));293+ (__gm__ T *)(workspace + impScoreOffsetU64 + perCoreOffsetU64 *
294+ static_cast<uint64_t>(this->cubeBlockIdx) + perCoreOffsetU64 / 2));
292 295 
293 GetExtremeValue(this->negativeFloatScalar, this->positiveFloatScalar);296 GetExtremeValue(this->negativeFloatScalar, this->positiveFloatScalar);
294}297}
@@ -424,9 +427,16 @@ NsaCompressAttentionS1s2Bn2gs1SameAB<layOutType, hasAtten, hasTopkMask, INPUT_T,
424 int64_t n2Offset = 0;427 int64_t n2Offset = 0;
425 428 
426 if constexpr (layOutType == LayOutTypeEnum::LAYOUT_TND) {429 if constexpr (layOutType == LayOutTypeEnum::LAYOUT_TND) {
427- n2Offset = extraInfo.n2oIdx * this->s1TotalSize * this->gSize * this->dSize;430+ uint64_t n2OffsetU64 = static_cast<uint64_t>(extraInfo.n2oIdx) * static_cast<uint64_t>(this->s1TotalSize);
428- bOffset = extraInfo.s1SizeAcc * this->gSize * this->dSize;431+ n2OffsetU64 = n2OffsetU64 * static_cast<uint64_t>(this->gSize) * static_cast<uint64_t>(this->dSize);
429- s1gOffset = extraInfo.s1oIdx * this->tilingData->coreParams.s1BaseSize * this->gSize * this->dSize;432+ uint64_t bOffsetU64 = static_cast<uint64_t>(extraInfo.s1SizeAcc) *
433+ static_cast<uint64_t>(this->gSize) * static_cast<uint64_t>(this->dSize);
434+ uint64_t s1gOffsetU64 =
435+ static_cast<uint64_t>(extraInfo.s1oIdx) * static_cast<uint64_t>(this->tilingData->coreParams.s1BaseSize);
436+ s1gOffsetU64 = s1gOffsetU64 * static_cast<uint64_t>(this->gSize) * static_cast<uint64_t>(this->dSize);
437+ n2Offset = static_cast<int64_t>(n2OffsetU64);
438+ bOffset = static_cast<int64_t>(bOffsetU64);
439+ s1gOffset = static_cast<int64_t>(s1gOffsetU64);
430 }440 }
431 this->qCoreOffset = n2Offset + bOffset + s1gOffset;441 this->qCoreOffset = n2Offset + bOffset + s1gOffset;
432 extraInfo.qCoreOffset = this->qCoreOffset;442 extraInfo.qCoreOffset = this->qCoreOffset;
@@ -626,8 +636,11 @@ NsaCompressAttentionS1s2Bn2gs1SameAB<layOutType, hasAtten, hasTopkMask, INPUT_T,
626 AscendC::SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);636 AscendC::SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
627 if constexpr (hasTopkMask == true) {637 if constexpr (hasTopkMask == true) {
628 uint64_t maskOffset = isInfo.outerLoop * loopIdx;638 uint64_t maskOffset = isInfo.outerLoop * loopIdx;
629- uint64_t maskGmOffset = (extraInfo.s1oIdx * this->cubeS1BaseSize + extraInfo.vecCoreOffset / this->gSize)639+ uint64_t s1Offset = static_cast<uint64_t>(extraInfo.s1oIdx) * static_cast<uint64_t>(this->cubeS1BaseSize);
630- * CeilDiv(this->tilingData->inputParams.s2Size, isInfo.isM);640+ uint64_t vecOffset = static_cast<uint64_t>(extraInfo.vecCoreOffset) / static_cast<uint64_t>(this->gSize);
641+ uint64_t totalSeqOffset = s1Offset + vecOffset;
642+ uint64_t s2Div = CeilDiv(static_cast<uint64_t>(this->tilingData->inputParams.s2Size), isInfo.isM);
643+ uint64_t maskGmOffset = totalSeqOffset * s2Div;
631 uint64_t maskPad = maskLenAligned32B - scoreLoop;644 uint64_t maskPad = maskLenAligned32B - scoreLoop;
632 uint32_t maskSrcStride = static_cast<uint32_t>((this->s2Aligned64B - this->scoreLoop) / 32);645 uint32_t maskSrcStride = static_cast<uint32_t>((this->s2Aligned64B - this->scoreLoop) / 32);
633 if (loopIdx == this->s2Loop - 1) {646 if (loopIdx == this->s2Loop - 1) {
@@ -827,9 +840,17 @@ NsaCompressAttentionS1s2Bn2gs1SameAB<layOutType, hasAtten, hasTopkMask, INPUT_T,
827 int64_t s1gOffset = 0;840 int64_t s1gOffset = 0;
828 841 
829 if constexpr (layOutType == LayOutTypeEnum::LAYOUT_TND) {842 if constexpr (layOutType == LayOutTypeEnum::LAYOUT_TND) {
830- n2Offset = extraInfo.n2oIdx * this->s1TotalSize * k;843+ uint64_t n2OffsetU64 = static_cast<uint64_t>(extraInfo.n2oIdx) *
831- bOffset = extraInfo.s1SizeAcc * k;844+ static_cast<uint64_t>(this->s1TotalSize) * static_cast<uint64_t>(k);
832- s1gOffset = (extraInfo.s1oIdx * this->tilingData->coreParams.s1BaseSize + this->cubeSubIdx * ((extraInfo.cubeS1RealSize + 1) / 2)) * k;845+ uint64_t bOffsetU64 = static_cast<uint64_t>(extraInfo.s1SizeAcc) * static_cast<uint64_t>(k);
846+ uint64_t s1Base =
847+ static_cast<uint64_t>(extraInfo.s1oIdx) * static_cast<uint64_t>(this->tilingData->coreParams.s1BaseSize);
848+ uint64_t cubeOffset =
849+ static_cast<uint64_t>(this->cubeSubIdx) * ((static_cast<uint64_t>(extraInfo.cubeS1RealSize) + 1) / 2);
850+ uint64_t s1gOffsetU64 = (s1Base + cubeOffset) * static_cast<uint64_t>(k);
851+ n2Offset = static_cast<int64_t>(n2OffsetU64);
852+ bOffset = static_cast<int64_t>(bOffsetU64);
853+ s1gOffset = static_cast<int64_t>(s1gOffsetU64);
833 }854 }
834 int64_t TopkOutCoreOffset = n2Offset + bOffset + s1gOffset;855 int64_t TopkOutCoreOffset = n2Offset + bOffset + s1gOffset;
835 856 
@@ -1036,8 +1057,11 @@ __aicore__ inline void NsaCompressAttentionS1s2Bn2gs1SameAB<layOutType, hasAtten
1036 1057 
1037 if constexpr (layOutType == LayOutTypeEnum::LAYOUT_TND) {1058 if constexpr (layOutType == LayOutTypeEnum::LAYOUT_TND) {
1038 // b,s2,n2,d1059 // b,s2,n2,d
1039- bOffset = extraInfo.s2SizeAcc * this->tilingData->inputParams.n2Size * this->d2Size;1060+ uint64_t bOffsetU64 = static_cast<uint64_t>(extraInfo.s2SizeAcc) *
1040- n2Offset = extraInfo.n2oIdx * this->d2Size;1061+ static_cast<uint64_t>(this->tilingData->inputParams.n2Size) * static_cast<uint64_t>(this->d2Size);
1062+ uint64_t n2OffsetU64 = static_cast<uint64_t>(extraInfo.n2oIdx) * static_cast<uint64_t>(this->d2Size);
1063+ bOffset = static_cast<int64_t>(bOffsetU64);
1064+ n2Offset = static_cast<int64_t>(n2OffsetU64);
1041 }1065 }
1042 1066 
1043 int64_t vCoreOffset = n2Offset + bOffset + s2Offset;1067 int64_t vCoreOffset = n2Offset + bOffset + s2Offset;
@@ -1047,9 +1071,16 @@ __aicore__ inline void NsaCompressAttentionS1s2Bn2gs1SameAB<layOutType, hasAtten
1047 bmm2.SetTensorB(this->valueGm[vCoreOffset]);1071 bmm2.SetTensorB(this->valueGm[vCoreOffset]);
1048 1072 
1049 if constexpr (layOutType == LayOutTypeEnum::LAYOUT_TND) {1073 if constexpr (layOutType == LayOutTypeEnum::LAYOUT_TND) {
1050- n2Offset = extraInfo.n2oIdx * this->s1TotalSize * this->gSize * this->d2Size;1074+ uint64_t n2OffsetU64 = static_cast<uint64_t>(extraInfo.n2oIdx) * static_cast<uint64_t>(this->s1TotalSize);
1051- bOffset = extraInfo.s1SizeAcc * this->gSize * this->d2Size;1075+ n2OffsetU64 = n2OffsetU64 * static_cast<uint64_t>(this->gSize) * static_cast<uint64_t>(this->d2Size);
1052- s1gOffset = extraInfo.s1oIdx * this->tilingData->coreParams.s1BaseSize * this->gSize * this->d2Size;1076+ uint64_t bOffsetU64 = static_cast<uint64_t>(extraInfo.s1SizeAcc) *
1077+ static_cast<uint64_t>(this->gSize) * static_cast<uint64_t>(this->d2Size);
1078+ uint64_t s1gOffsetU64 =
1079+ static_cast<uint64_t>(extraInfo.s1oIdx) * static_cast<uint64_t>(this->tilingData->coreParams.s1BaseSize);
1080+ s1gOffsetU64 = s1gOffsetU64 * static_cast<uint64_t>(this->gSize) * static_cast<uint64_t>(this->d2Size);
1081+ n2Offset = static_cast<int64_t>(n2OffsetU64);
1082+ bOffset = static_cast<int64_t>(bOffsetU64);
1083+ s1gOffset = static_cast<int64_t>(s1gOffsetU64);
1053 }1084 }
1054 int64_t outCoreOffset = n2Offset + bOffset + s1gOffset;1085 int64_t outCoreOffset = n2Offset + bOffset + s1gOffset;
1055 1086 
Mattention/nsa_compress_grad/op_host/nsa_compress_grad_tiling.cpp+11-4
@@ -97,8 +97,8 @@ ge::graphStatus NsaCompressGradTiling::DoOpTiling()
97 97 
98 const auto& inputKvShape = GetShapeOfInput(INDEXONE);98 const auto& inputKvShape = GetShapeOfInput(INDEXONE);
99 tSeqLen_ = static_cast<uint32_t>(inputKvShape.GetDim(INDEXZERO));99 tSeqLen_ = static_cast<uint32_t>(inputKvShape.GetDim(INDEXZERO));
100- auto headToProcessPerCore = tSeqLen_ * nHeads / aivNum_;100+ auto headToProcessPerCore = static_cast<uint64_t>(tSeqLen_) * nHeads / aivNum_;
101- auto headRemainder = (tSeqLen_ * nHeads) % aivNum_;101+ auto headRemainder = (static_cast<uint64_t>(tSeqLen_) * nHeads) % aivNum_;
102 102 
103 auto batchSize = 1;103 auto batchSize = 1;
104 const auto& actSeqLenShape = context_->GetOptionalInputShape(INDEXTHREE);104 const auto& actSeqLenShape = context_->GetOptionalInputShape(INDEXTHREE);
@@ -143,8 +143,15 @@ ge::graphStatus NsaCompressGradTiling::GetWorkspaceSize() {
143 size_t *workspaces = context_->GetWorkspaceSizes(1);143 size_t *workspaces = context_->GetWorkspaceSizes(1);
144 size_t sysWorkspaceSize = WORKSIZE;144 size_t sysWorkspaceSize = WORKSIZE;
145 /* usrWorkspaceSize = workspace for other gm */145 /* usrWorkspaceSize = workspace for other gm */
146- size_t usrWorkspaceSize = tilingData_.get_blockSize() * tilingData_.get_numOfHead() * sizeof(uint32_t) * aivNum_ +146+ int64_t usrWorkspaceSizeTmp = static_cast<int64_t>(tilingData_.get_blockSize()) *
147- tSeqLen_ * tilingData_.get_numOfHead() * tilingData_.get_dimOfHead() * sizeof(uint32_t) + ONE_KILO * ONE_KILO;147+ static_cast<int64_t>(tilingData_.get_numOfHead()) *
148+ sizeof(uint32_t) * static_cast<int64_t>(aivNum_) +
149+ static_cast<int64_t>(tSeqLen_) *
150+ static_cast<int64_t>(tilingData_.get_numOfHead()) *
151+ static_cast<int64_t>(tilingData_.get_dimOfHead()) *
152+ sizeof(uint32_t) +
153+ static_cast<int64_t>(ONE_KILO) * ONE_KILO;
154+ size_t usrWorkspaceSize = static_cast<size_t>(usrWorkspaceSizeTmp);
148 workspaces[0] = sysWorkspaceSize + usrWorkspaceSize;155 workspaces[0] = sysWorkspaceSize + usrWorkspaceSize;
149 return ge::GRAPH_SUCCESS;156 return ge::GRAPH_SUCCESS;
150}157}
Mattention/nsa_compress_grad/op_kernel/nsa_compress_grad.h+17-16
@@ -109,7 +109,7 @@ __aicore__ inline void NsaCompressGradND<T>::InitWorkspace(GM_ADDR workspace) {
109 109 
110 if (isOverLap_) {110 if (isOverLap_) {
111 // overLap场景需要使用workspace累加111 // overLap场景需要使用workspace累加
112- uint32_t inputGradOffset = blockSize_ * headNum_ * coreNum_ * sizeof(float);112+ uint64_t inputGradOffset = static_cast<uint64_t>(blockSize_) * headNum_ * coreNum_ * sizeof(float);
113 inputGradAtomicWs_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(workspace + inputGradOffset));113 inputGradAtomicWs_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(workspace + inputGradOffset));
114 }114 }
115}115}
@@ -196,8 +196,8 @@ __aicore__ inline void NsaCompressGradND<T>::SetWsToZero() {
196 }196 }
197 197 
198 // 清零wtGradAtomicWs_198 // 清零wtGradAtomicWs_
199- uint32_t wtGradNum = blockSize_ * headNum_;199+ uint64_t wtGradNum = static_cast<uint64_t>(blockSize_) * headNum_;
200- Duplicate<float>(cmpGrad, ZERO_FLOAT, wtGradNum);200+ Duplicate<float>(cmpGrad, ZERO_FLOAT, static_cast<uint32_t>(wtGradNum));
201 201 
202 SetFlag<HardEvent::V_MTE3>(EVENT_ID3);202 SetFlag<HardEvent::V_MTE3>(EVENT_ID3);
203 WaitFlag<HardEvent::V_MTE3>(EVENT_ID3);203 WaitFlag<HardEvent::V_MTE3>(EVENT_ID3);
@@ -243,10 +243,10 @@ __aicore__ inline void NsaCompressGradND<T>::Process() {
243 PresetFlag();243 PresetFlag();
244 244 
245 uint32_t curBIdx = startBatchIdx_;245 uint32_t curBIdx = startBatchIdx_;
246- uint32_t offsetPerRow = headDim_ * headNum_;246+ uint64_t offsetPerRow = static_cast<uint64_t>(headDim_) * headNum_;
247- uint32_t offsetPerBlc = offsetPerRow * blockStride_;247+ uint64_t offsetPerBlc = offsetPerRow * blockStride_;
248 // inptKv当前处理的第一个batch的偏移248 // inptKv当前处理的第一个batch的偏移
249- uint32_t batchOffset = batchOffsetCurCore_ * offsetPerRow;249+ uint64_t batchOffset = static_cast<uint64_t>(batchOffsetCurCore_) * offsetPerRow;
250 250 
251 uint32_t blcStartIdxCurBch = blcIdxOfStartBatch_;251 uint32_t blcStartIdxCurBch = blcIdxOfStartBatch_;
252 uint32_t seqLenCurBch = GetSeqLenTarBch(curBIdx);252 uint32_t seqLenCurBch = GetSeqLenTarBch(curBIdx);
@@ -265,20 +265,21 @@ __aicore__ inline void NsaCompressGradND<T>::Process() {
265 nBlcCurBch = CompressBlkNum(seqLenCurBch, blockSize_, blockStride_);265 nBlcCurBch = CompressBlkNum(seqLenCurBch, blockSize_, blockStride_);
266 }266 }
267 // inputKV当前处理的block的偏移267 // inputKV当前处理的block的偏移
268- auto blcOffset = batchOffset + (blcId - blcStartIdxCurBch) * offsetPerBlc;268+ uint64_t blcOffset = batchOffset + static_cast<uint64_t>(blcId - blcStartIdxCurBch) * offsetPerBlc;
269 269 
270 // deal with current block270 // deal with current block
271 for (auto headId = 0; headId < headNum_; headId += nHeadOnce_) {271 for (auto headId = 0; headId < headNum_; headId += nHeadOnce_) {
272 // Block内按行处理272 // Block内按行处理
273 uint32_t headToProcess = (headId + nHeadOnce_ > headNum_) ? (headNum_ - headId) : nHeadOnce_;273 uint32_t headToProcess = (headId + nHeadOnce_ > headNum_) ? (headNum_ - headId) : nHeadOnce_;
274- uint32_t cmpOffset = blcId * offsetPerRow + headId * headDim_;274+ uint64_t cmpOffset = static_cast<uint64_t>(blcId) * offsetPerRow + static_cast<uint64_t>(headId) * headDim_;
275 275 
276 WaitFlag<HardEvent::V_MTE2>(EVENT_ID6);276 WaitFlag<HardEvent::V_MTE2>(EVENT_ID6);
277 CopyInCmpGrad(cmpOffset, headToProcess * headDim_);277 CopyInCmpGrad(cmpOffset, headToProcess * headDim_);
278 278
279 for (auto rowId = 0; rowId < blockSize_; rowId+=nRowsOnce_) {279 for (auto rowId = 0; rowId < blockSize_; rowId+=nRowsOnce_) {
280- uint32_t inputOffset = blcOffset + rowId * offsetPerRow + headId * headDim_;280+ uint64_t inputOffset =
281- uint32_t weightOffset = rowId * headNum_ + headId;281+ blcOffset + static_cast<uint64_t>(rowId) * offsetPerRow + static_cast<uint64_t>(headId) * headDim_;
282+ uint64_t weightOffset = static_cast<uint64_t>(rowId) * headNum_ + headId;
282 uint32_t rowsToProcess = (rowId + nRowsOnce_ > blockSize_) ? (blockSize_ - rowId) : nRowsOnce_;283 uint32_t rowsToProcess = (rowId + nRowsOnce_ > blockSize_) ? (blockSize_ - rowId) : nRowsOnce_;
283 284 
284 // 拷贝inputKV285 // 拷贝inputKV
@@ -442,10 +443,10 @@ __aicore__ inline void NsaCompressGradND<T>::DeterministicComputeSumWtGrad() {
442 return;443 return;
443 }444 }
444 445 
445- auto totalWeightLength = blockSize_ * headNum_;446+ uint64_t totalWeightLength = static_cast<uint64_t>(blockSize_) * headNum_;
446- auto minOffset = 16;447+ uint64_t minOffset = 16;
447- auto nLengthBase = (blockSize_ * headNum_ / minOffset) / coreNum_;448+ uint64_t nLengthBase = (static_cast<uint64_t>(blockSize_) * headNum_ / minOffset) / coreNum_;
448- auto nLengthReminder = (blockSize_ * headNum_ / minOffset) % coreNum_;449+ uint64_t nLengthReminder = (static_cast<uint64_t>(blockSize_) * headNum_ / minOffset) % coreNum_;
449 //当前core从第几个数据开始算450 //当前core从第几个数据开始算
450 auto startLengthIdx = (nLengthBase * coreId_ + (coreId_ < nLengthReminder ? coreId_ : nLengthReminder)) * minOffset;451 auto startLengthIdx = (nLengthBase * coreId_ + (coreId_ < nLengthReminder ? coreId_ : nLengthReminder)) * minOffset;
451 //当前core需要处理几个数据452 //当前core需要处理几个数据
@@ -501,7 +502,7 @@ __aicore__ inline void NsaCompressGradND<T>::DeterministicComputeSumWtGrad() {
501 502 
502template<typename T>503template<typename T>
503__aicore__ inline void NsaCompressGradND<T>::MoveWtGrad() {504__aicore__ inline void NsaCompressGradND<T>::MoveWtGrad() {
504- uint32_t wtNum = blockSize_ * headNum_;505+ uint64_t wtNum = static_cast<uint64_t>(blockSize_) * headNum_;
505 DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(wtNum * sizeof(float)), 0, 0, 0};506 DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(wtNum * sizeof(float)), 0, 0, 0};
506 DataCopyPadExtParams<float> padParams{false, 0, 0, 0};507 DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
507 DataCopyPad(transWtGrad_, wtGradAtomicWs_, dataCopyParams, padParams);508 DataCopyPad(transWtGrad_, wtGradAtomicWs_, dataCopyParams, padParams);
@@ -509,7 +510,7 @@ __aicore__ inline void NsaCompressGradND<T>::MoveWtGrad() {
509 SetFlag<HardEvent::MTE2_V>(EVENT_ID3);510 SetFlag<HardEvent::MTE2_V>(EVENT_ID3);
510 WaitFlag<HardEvent::MTE2_V>(EVENT_ID3);511 WaitFlag<HardEvent::MTE2_V>(EVENT_ID3);
511 512 
512- CastByCondition(halfTransWtGrad_, transWtGrad_, wtNum, true);513+ CastByCondition(halfTransWtGrad_, transWtGrad_, static_cast<uint32_t>(wtNum), true);
513 514 
514 SetFlag<HardEvent::V_MTE3>(EVENT_ID3);515 SetFlag<HardEvent::V_MTE3>(EVENT_ID3);
515 WaitFlag<HardEvent::V_MTE3>(EVENT_ID3);516 WaitFlag<HardEvent::V_MTE3>(EVENT_ID3);
Mattention/nsa_selected_attention/op_kernel/nsa_selected_attention.h+5-5
@@ -622,9 +622,9 @@ __aicore__ inline void NsaSelectedAttention<INPUT_T, ATTEN_ENABLE>::HandleMask(u
622{622{
623 event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE2_V>());623 event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE2_V>());
624 event_t eventIdMte2ToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));624 event_t eventIdMte2ToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
625- uint8_t dstStride = this->selectedLength / FP32_DATA_BLOCK;625+ uint32_t dstStride = this->selectedLength / FP32_DATA_BLOCK;
626 // when block size > 64, total bytes in one block will exceed 256626 // when block size > 64, total bytes in one block will exceed 256
627- // when selectedLength >= 2048, dstStride will be >= 256, can not be saved in uint8_t627+ // when selectedLength >= 2048, dstStride will be >= 256, useLoop branch avoids narrow stride param
628 bool useLoop = this->selectedBlockSize > UL_BIT_COUNT || (this->selectedLength >= FP32_DATA_BLOCK * NUM_256);628 bool useLoop = this->selectedBlockSize > UL_BIT_COUNT || (this->selectedLength >= FP32_DATA_BLOCK * NUM_256);
629 uint64_t diagBlockIdx = this->currentS1 / this->selectedBlockSize;629 uint64_t diagBlockIdx = this->currentS1 / this->selectedBlockSize;
630 int32_t currentBlockIdx = 0;630 int32_t currentBlockIdx = 0;
@@ -645,7 +645,7 @@ __aicore__ inline void NsaSelectedAttention<INPUT_T, ATTEN_ENABLE>::HandleMask(u
645 } else {645 } else {
646 GetShiftLeftMask(this->selectedBlockSize, saveNum, mask[0]);646 GetShiftLeftMask(this->selectedBlockSize, saveNum, mask[0]);
647 AscendC::Duplicate(this->bmm1ResUb[selectedBlockIdx * this->selectedBlockSize],647 AscendC::Duplicate(this->bmm1ResUb[selectedBlockIdx * this->selectedBlockSize],
648- this->negativeScalar, mask, currentProcessSize, 1, dstStride);648+ this->negativeScalar, mask, currentProcessSize, 1, static_cast<uint8_t>(dstStride));
649 }649 }
650 AscendC::PipeBarrier<PIPE_V>();650 AscendC::PipeBarrier<PIPE_V>();
651 } else if (currentBlockIdx > diagBlockIdx) {651 } else if (currentBlockIdx > diagBlockIdx) {
@@ -655,8 +655,8 @@ __aicore__ inline void NsaSelectedAttention<INPUT_T, ATTEN_ENABLE>::HandleMask(u
655 AscendC::Duplicate(this->bmm1ResUb[offset], this->negativeScalar, this->selectedBlockSize);655 AscendC::Duplicate(this->bmm1ResUb[offset], this->negativeScalar, this->selectedBlockSize);
656 }656 }
657 } else {657 } else {
658- AscendC::Duplicate(this->bmm1ResUb[selectedBlockIdx * this->selectedBlockSize],658+ AscendC::Duplicate(this->bmm1ResUb[selectedBlockIdx * this->selectedBlockSize], this->negativeScalar,
659- this->negativeScalar, this->selectedBlockSize, currentProcessSize, 1, dstStride);659+ this->selectedBlockSize, currentProcessSize, 1, static_cast<uint8_t>(dstStride));
660 }660 }
661 AscendC::PipeBarrier<PIPE_V>();661 AscendC::PipeBarrier<PIPE_V>();
662 }662 }
Mattention/nsa_selected_attention_grad/basic_modules/vec_op.h+1-1
@@ -452,7 +452,7 @@ __aicore__ inline void VecOp<NSAGT>::CalSoftmax(const int32_t loopIdx, const int
452 const int64_t mm345Addr, const uint64_t indicesGmOffset,452 const int64_t mm345Addr, const uint64_t indicesGmOffset,
453 const int64_t s1Index, const int32_t blkCntOffset)453 const int64_t s1Index, const int32_t blkCntOffset)
454{454{
455- int64_t dataSize = processM * params.sftBaseN;455+ int64_t dataSize = static_cast<int64_t>(processM) * static_cast<int64_t>(params.sftBaseN);
456 int64_t maxOffset = loopIdx * params.sftBaseM * BLOCK_FP32;456 int64_t maxOffset = loopIdx * params.sftBaseM * BLOCK_FP32;
457 auto tmpMaxTensor = maxTensor[maxOffset];457 auto tmpMaxTensor = maxTensor[maxOffset];
458 auto tmpSumTensor = sumTensor[maxOffset];458 auto tmpSumTensor = sumTensor[maxOffset];
Mattention/nsa_selected_attention_grad/op_kernel/nsa_selected_attention_grad_bs1.h+3-3
@@ -775,7 +775,7 @@ template <typename NSAGT>
775__aicore__ inline void SelectedAttentionGrad<NSAGT>::CalSoftmax(const int32_t loopIdx, const int32_t processM,775__aicore__ inline void SelectedAttentionGrad<NSAGT>::CalSoftmax(const int32_t loopIdx, const int32_t processM,
776 const int64_t mm12Addr, const int64_t mm345Addr)776 const int64_t mm12Addr, const int64_t mm345Addr)
777{777{
778- int64_t dataSize = processM * params.sftBaseN;778+ int64_t dataSize = static_cast<int64_t>(processM) * static_cast<int64_t>(params.sftBaseN);
779 int64_t maxOffset = loopIdx * params.sftBaseM * BLOCK_FP32;779 int64_t maxOffset = loopIdx * params.sftBaseM * BLOCK_FP32;
780 auto tmpMaxTensor = maxTensor[maxOffset];780 auto tmpMaxTensor = maxTensor[maxOffset];
781 auto tmpSumTensor = sumTensor[maxOffset];781 auto tmpSumTensor = sumTensor[maxOffset];
@@ -810,7 +810,7 @@ template <typename NSAGT>
810__aicore__ inline void SelectedAttentionGrad<NSAGT>::CalSoftmaxGrad(const int32_t loopIdx, const int32_t processM,810__aicore__ inline void SelectedAttentionGrad<NSAGT>::CalSoftmaxGrad(const int32_t loopIdx, const int32_t processM,
811 const int64_t mm12Addr, const int64_t mm345Addr)811 const int64_t mm12Addr, const int64_t mm345Addr)
812{812{
813- int64_t dataSize = processM * params.sftBaseN;813+ int64_t dataSize = static_cast<int64_t>(processM) * static_cast<int64_t>(params.sftBaseN);
814 int64_t rowsumOffset = loopIdx * params.sftBaseM * BLOCK_FP32;814 int64_t rowsumOffset = loopIdx * params.sftBaseM * BLOCK_FP32;
815 auto tmpRowSumOutTensor = rowSumOutTensor[rowsumOffset];815 auto tmpRowSumOutTensor = rowSumOutTensor[rowsumOffset];
816 816 
@@ -840,7 +840,7 @@ __aicore__ inline void SelectedAttentionGrad<NSAGT>::VectorProcess(int64_t mm12A
840 int32_t loop = (params.singleM + params.sftBaseM - 1) / params.sftBaseM;840 int32_t loop = (params.singleM + params.sftBaseM - 1) / params.sftBaseM;
841 int32_t processM = params.sftBaseM;841 int32_t processM = params.sftBaseM;
842 int32_t tailM = params.singleM % params.sftBaseM;842 int32_t tailM = params.singleM % params.sftBaseM;
843- int64_t dataSize = processM * params.sftBaseN;843+ int64_t dataSize = static_cast<int64_t>(processM) * static_cast<int64_t>(params.sftBaseN);
844 844 
845 for (int32_t i = 0; i < loop; i++) {845 for (int32_t i = 0; i < loop; i++) {
846 if (i == 0) {846 if (i == 0) {
Mattention/sparse_flash_mla_grad/op_host/sparse_flash_mla_grad_tiling_general.cpp+1-1
@@ -19,7 +19,7 @@
19 19 
20namespace optiling {20namespace optiling {
21namespace smlag {21namespace smlag {
22-constexpr uint32_t WORKSPACE_BASE_CAL = 32 * 1024 * 1024; // 100MB系统预留22+constexpr uint32_t WORKSPACE_BASE_CAL = 32 * 1024 * 1024; // 32MB系统预留
23constexpr uint32_t BLOCK = 32; // 32B23constexpr uint32_t BLOCK = 32; // 32B
24constexpr uint32_t B32 = 4; // 4B24constexpr uint32_t B32 = 4; // 4B
25constexpr uint32_t B16 = 2;25constexpr uint32_t B16 = 2;
Mattention/sparse_flash_mla_grad/op_kernel/arch22/basic_modules/cube_modules/cube3.h+2-1
@@ -48,7 +48,8 @@ CubeOp<T1>::cube3Process(const int64_t keyGmOffset,
48 dsL1CachedSize = 0;48 dsL1CachedSize = 0;
49 // selectedBlockSize较小时，nIdx需要+2或者更多49 // selectedBlockSize较小时，nIdx需要+2或者更多
50 for (int32_t nIdx = blkCntOffset; nIdx < blkCntOffset + selectedCntOffset; nIdx+=blockOffset) {50 for (int32_t nIdx = blkCntOffset; nIdx < blkCntOffset + selectedCntOffset; nIdx+=blockOffset) {
51- int32_t l1Offset = (nIdx - blkCntOffset) * selectedBlockSize * AlignTo<int64_t>(mmParam.singleM, SIZE_16);51+ int64_t l1Offset = static_cast<int64_t>(nIdx - blkCntOffset) *
52+ selectedBlockSize * AlignTo<int64_t>(mmParam.singleM, SIZE_16);
52 bool isFirstLoop = (nIdx == blkCntOffset);53 bool isFirstLoop = (nIdx == blkCntOffset);
53 bool isLastLoop = (nIdx + blockOffset >= blkCntOffset + selectedCntOffset);54 bool isLastLoop = (nIdx + blockOffset >= blkCntOffset + selectedCntOffset);
54 55 
Mattention/sparse_flash_mla_grad/op_kernel/arch22/basic_modules/cube_modules/cube4.h+2-1
@@ -50,7 +50,8 @@ CubeOp<T1>::cube4Process(const int64_t dsGmOffset, const int64_t queryGmOffset,
50 CopyGmToL1(l1_ds_tensor, dsWorkspaceGm[dsGmOffset], mmParam.singleK, selectedCntOffset * selectedBlockSize, singleN);50 CopyGmToL1(l1_ds_tensor, dsWorkspaceGm[dsGmOffset], mmParam.singleK, selectedCntOffset * selectedBlockSize, singleN);
51 }51 }
52 for (int32_t mIdx = blkCntOffset; mIdx < blkCntOffset + selectedCntOffset; mIdx+=blockOffset) {52 for (int32_t mIdx = blkCntOffset; mIdx < blkCntOffset + selectedCntOffset; mIdx+=blockOffset) {
53- int32_t l1Offset = (mIdx - blkCntOffset) * selectedBlockSize * AlignTo<int64_t>(mmParam.singleK, SIZE_16);53+ int64_t l1Offset =
54+ static_cast<int64_t>(mIdx - blkCntOffset) * selectedBlockSize * AlignTo<int64_t>(mmParam.singleK, SIZE_16);
54 55 
55 mmParam.singleN = perLoopDSize;56 mmParam.singleN = perLoopDSize;
56 mmParam.singleM = min(selectedBlockSize * blockOffset, selectedCntOffset * selectedBlockSize - (mIdx - blkCntOffset) * selectedBlockSize);57 mmParam.singleM = min(selectedBlockSize * blockOffset, selectedCntOffset * selectedBlockSize - (mIdx - blkCntOffset) * selectedBlockSize);
Mattention/sparse_flash_mla_grad/op_kernel/arch22/sfag_basic.h+2-2
@@ -216,7 +216,7 @@ __aicore__ inline void SelectedAttentionGradBasic<SMLAGT>::Process(
216 bool changeS1 = false;216 bool changeS1 = false;
217 for (int32_t i = 0; i < processBS1ByCore; i++) {217 for (int32_t i = 0; i < processBS1ByCore; i++) {
218 scatterTaskId = i % 2;218 scatterTaskId = i % 2;
219- int32_t t1Index = cubeBlockIdx + usedCoreNum * i;219+ int64_t t1Index = static_cast<int64_t>(cubeBlockIdx) + usedCoreNum * i;
220 GetTndSeqLen(cu_seqlens_q, cu_seqlens_ori_kv, cu_seqlens_cmp_kv, cmp_residual_kv, t1Index, bIndex);220 GetTndSeqLen(cu_seqlens_q, cu_seqlens_ori_kv, cu_seqlens_cmp_kv, cmp_residual_kv, t1Index, bIndex);
221 changePingpong = false;221 changePingpong = false;
222 for (n2Index = 0; n2Index < dimN2; n2Index++) {222 for (n2Index = 0; n2Index < dimN2; n2Index++) {
@@ -255,7 +255,7 @@ __aicore__ inline void SelectedAttentionGradBasic<SMLAGT>::Process(
255 int64_t task = 0;255 int64_t task = 0;
256 for (int32_t i = 0; i < processBS1ByCore; i++) {256 for (int32_t i = 0; i < processBS1ByCore; i++) {
257 scatterTaskId = i % 2;257 scatterTaskId = i % 2;
258- int32_t t1Index = cubeBlockIdx + usedCoreNum * i;258+ int64_t t1Index = static_cast<int64_t>(cubeBlockIdx) + usedCoreNum * i;
259 GetTndSeqLen(cu_seqlens_q, cu_seqlens_ori_kv, cu_seqlens_cmp_kv, cmp_residual_kv, t1Index, bIndex);259 GetTndSeqLen(cu_seqlens_q, cu_seqlens_ori_kv, cu_seqlens_cmp_kv, cmp_residual_kv, t1Index, bIndex);
260 changePingpong = false;260 changePingpong = false;
261 for (n2Index = 0; n2Index < dimN2; n2Index++) {261 for (n2Index = 0; n2Index < dimN2; n2Index++) {
Mattention/sparse_flash_mla_grad/op_kernel/arch22/smlag_basic.h+2-2
@@ -222,7 +222,7 @@ __aicore__ inline void SparseFlashMlaGrad<SMLAGT>::Process(
222 AllocEventID();222 AllocEventID();
223 int64_t task = 0;223 int64_t task = 0;
224 for (int32_t i = 0; i < processBS1ByCore; i++) {224 for (int32_t i = 0; i < processBS1ByCore; i++) {
225- int32_t t1Index = (cubeBlockIdx + usedCoreNum * i) * s1BasicSize;225+ int64_t t1Index = (static_cast<int64_t>(cubeBlockIdx) + usedCoreNum * i) * s1BasicSize;
226 GetTndSeqLen(t1Index, bIndex, s1Loop);226 GetTndSeqLen(t1Index, bIndex, s1Loop);
227 int32_t s1BasicAccum = 0;227 int32_t s1BasicAccum = 0;
228 for (int32_t j = 0; j < s1Loop; j++) {228 for (int32_t j = 0; j < s1Loop; j++) {
@@ -261,7 +261,7 @@ __aicore__ inline void SparseFlashMlaGrad<SMLAGT>::Process(
261 SyncAll();261 SyncAll();
262 int64_t task = 0;262 int64_t task = 0;
263 for (int32_t i = 0; i < processBS1ByCore; i++) {263 for (int32_t i = 0; i < processBS1ByCore; i++) {
264- int32_t t1Index = (cubeBlockIdx + usedCoreNum * i) * s1BasicSize;264+ int64_t t1Index = (static_cast<int64_t>(cubeBlockIdx) + usedCoreNum * i) * s1BasicSize;
265 GetTndSeqLen(t1Index, bIndex, s1Loop);265 GetTndSeqLen(t1Index, bIndex, s1Loop);
266 int32_t s1BasicAccum = 0;266 int32_t s1BasicAccum = 0;
267 for (int32_t j = 0; j < s1Loop; j++) {267 for (int32_t j = 0; j < s1Loop; j++) {
Mattention/sparse_lightning_indexer_kl_loss_grad/op_host/sparse_lightning_indexer_kl_loss_grad_tiling_general.cpp+17-3
@@ -328,7 +328,7 @@ void SparseLightningIndexerKLLossGradTilingBase::SetMultiCoreParamsRegbase(int64
328 }328 }
329}329}
330 330 
331-void SparseLightningIndexerKLLossGradTilingBase::InitOutputSplit()331+ge::graphStatus SparseLightningIndexerKLLossGradTilingBase::InitOutputSplit()
D
DDaiHuina16月8日

增加错误码返回后，在调用处增加校验，/ops-transformer/attention/sparse_lightning_indexer_grad_kl_loss/op_host/arch35/sparse_lightning_indexer_grad_kl_loss_tiling_general_regbase.cpp

likedislike
memorecool
memorecool
6月8日 评论：
332{332{
333 SLIKLLossGradInitOutputParams *initoutput = &tilingData->initOutputParams;333 SLIKLLossGradInitOutputParams *initoutput = &tilingData->initOutputParams;
334 auto &dKeyShape = context_->GetOutputShape(D_KEY_OUTPUT_INDEX)->GetStorageShape();334 auto &dKeyShape = context_->GetOutputShape(D_KEY_OUTPUT_INDEX)->GetStorageShape();
@@ -338,8 +338,19 @@ void SparseLightningIndexerKLLossGradTilingBase::InitOutputSplit()
338 } else if (tilingKeyLayout == LayoutType::LAYOUT_BSND) {338 } else if (tilingKeyLayout == LayoutType::LAYOUT_BSND) {
339 totalSize = static_cast<int64_t>(bSize) * dKeyShape.GetDim(1) * dKeyShape.GetDim(3);339 totalSize = static_cast<int64_t>(bSize) * dKeyShape.GetDim(1) * dKeyShape.GetDim(3);
340 }340 }
341- initoutput->set_singleCoreSize(static_cast<uint32_t>(CeilDivision(totalSize, static_cast<int64_t>(aivNum))));341+ int64_t singleCoreSize = CeilDivision(totalSize, static_cast<int64_t>(aivNum));
342+ if (singleCoreSize > UINT32_MAX) {
343+ OP_LOGE(context_, "singleCoreSize(%ld) exceeds UINT32_MAX limit.", singleCoreSize);
344+ return ge::GRAPH_FAILED;
345+ }
346+ initoutput->set_singleCoreSize(static_cast<uint32_t>(singleCoreSize));
347+ 
348+ if (totalSize > INT64_MAX / 2) {
349+ OP_LOGE(context_, "totalOutputSize(%ld) exceeds safe limit.", totalSize);
350+ return ge::GRAPH_FAILED;
351+ }
342 initoutput->set_totalOutputSize(totalSize);352 initoutput->set_totalOutputSize(totalSize);
353+ return ge::GRAPH_SUCCESS;
343}354}
344 355 
345ge::graphStatus SparseLightningIndexerKLLossGradTilingBase::DoOpTiling()356ge::graphStatus SparseLightningIndexerKLLossGradTilingBase::DoOpTiling()
@@ -357,7 +368,10 @@ ge::graphStatus SparseLightningIndexerKLLossGradTilingBase::DoOpTiling()
357 SoftMaxTilingFunc(srcShape, sizeof(float), softmaxTmpBufferSize, tilingData->vectorParams.softmaxYTilingData);368 SoftMaxTilingFunc(srcShape, sizeof(float), softmaxTmpBufferSize, tilingData->vectorParams.softmaxYTilingData);
358 SoftMaxTilingFunc(srcShape, sizeof(float), softmaxTmpBufferSize, tilingData->vectorParams.simpleSoftmaxPTilingData);369 SoftMaxTilingFunc(srcShape, sizeof(float), softmaxTmpBufferSize, tilingData->vectorParams.simpleSoftmaxPTilingData);
359 370 
360- InitOutputSplit();371+ auto ret = InitOutputSplit();
372+ if (ret != ge::GRAPH_SUCCESS) {
373+ return ret;
374+ }
361 OP_LOGD(context_, "ending template[%s]", templateName);375 OP_LOGD(context_, "ending template[%s]", templateName);
362 return ge::GRAPH_SUCCESS;376 return ge::GRAPH_SUCCESS;
363}377}
Mattention/sparse_lightning_indexer_kl_loss_grad/op_host/sparse_lightning_indexer_kl_loss_grad_tiling_general.h+1-1
@@ -287,7 +287,7 @@ protected:
287 bool AnalyzeLayout();287 bool AnalyzeLayout();
288 int64_t CalcTotalSize();288 int64_t CalcTotalSize();
289 void SetMultiCoreParamsRegbase(int64_t totalSize, int64_t coreNum);289 void SetMultiCoreParamsRegbase(int64_t totalSize, int64_t coreNum);
290- void InitOutputSplit();290+ ge::graphStatus InitOutputSplit();
291 291 
292 // 基础输入参数292 // 基础输入参数
293 int32_t bSize;293 int32_t bSize;
Mattention/sparse_lightning_indexer_kl_loss_grad/op_kernel/arch22/sparse_lightning_indexer_kl_loss_grad_vector.h+3-3
@@ -607,7 +607,7 @@ __aicore__ inline void SLIKLLossVectorService<SLIT>::MergeKv(const SLIKLLossGrad
607 if (s2GmLimit > gatherParams.s2ProcessSize) {607 if (s2GmLimit > gatherParams.s2ProcessSize) {
608 s2GmLimit = gatherParams.s2ProcessSize;608 s2GmLimit = gatherParams.s2ProcessSize;
609 }609 }
610- int64_t s2IdxOffset = runInfo.s2Idx * constInfo.s2BaseSize;610+ int64_t s2IdxOffset = static_cast<int64_t>(runInfo.s2Idx) * constInfo.s2BaseSize;
611 for (int64_t s2GmOffsetArray = s2GmStartOffset; s2GmOffsetArray < s2GmLimit; s2GmOffsetArray += 2 * constInfo.sparseBlockSize) {611 for (int64_t s2GmOffsetArray = s2GmStartOffset; s2GmOffsetArray < s2GmLimit; s2GmOffsetArray += 2 * constInfo.sparseBlockSize) {
612 if (needWaitMte3ToMte2) {612 if (needWaitMte3ToMte2) {
613 WaitFlag<AscendC::HardEvent::MTE3_MTE2>(mergeMte3Idx + IdStart);613 WaitFlag<AscendC::HardEvent::MTE3_MTE2>(mergeMte3Idx + IdStart);
@@ -875,7 +875,7 @@ __aicore__ inline void SLIKLLossVectorService<SLIT>::ProcessDeterVector2(SLIKLLo
875 int64_t realS2Idx1, realS2Idx2, s2GmOffset;875 int64_t realS2Idx1, realS2Idx2, s2GmOffset;
876 event_t eventIdArr[2] = {eventIdScatterAdd, eventIdScatterAddPong};876 event_t eventIdArr[2] = {eventIdScatterAdd, eventIdScatterAddPong};
877 runInfo.s2Idx = 0;877 runInfo.s2Idx = 0;
878- int64_t s2IdxOffset = runInfo.s2Idx * constInfo.s2BaseSize;878+ int64_t s2IdxOffset = static_cast<int64_t>(runInfo.s2Idx) * constInfo.s2BaseSize;
879 879 
880 CrossCoreWaitFlag<0, PIPE_MTE3>(SYNC_V2_TO_V2_DETER_SA_FLAG_MOD0);880 CrossCoreWaitFlag<0, PIPE_MTE3>(SYNC_V2_TO_V2_DETER_SA_FLAG_MOD0);
881 SetAtomicAdd<T>();881 SetAtomicAdd<T>();
@@ -979,7 +979,7 @@ __aicore__ inline void SLIKLLossVectorService<SLIT>::ProcessVector2(SLIKLLossGra
979 int64_t realS2Idx1, realS2Idx2, s2GmOffset;979 int64_t realS2Idx1, realS2Idx2, s2GmOffset;
980 event_t eventIdArr[2] = {eventIdScatterAdd, eventIdScatterAddPong};980 event_t eventIdArr[2] = {eventIdScatterAdd, eventIdScatterAddPong};
981 runInfo.s2Idx = 0;981 runInfo.s2Idx = 0;
982- int64_t s2IdxOffset = runInfo.s2Idx * constInfo.s2BaseSize;982+ int64_t s2IdxOffset = static_cast<int64_t>(runInfo.s2Idx) * constInfo.s2BaseSize;
983 983 
984 SetAtomicAdd<T>();984 SetAtomicAdd<T>();
985 for (int32_t kLoopIdx = 0; kLoopIdx < kLoopTimes; ++kLoopIdx) {985 for (int32_t kLoopIdx = 0; kLoopIdx < kLoopTimes; ++kLoopIdx) {