Star921
1841
代码
代码
Issues229
Pull Requests359
Actions
项目讨论
Wiki
分析
Star921
1841
已合并
修复swiglu_group_quant_grad bf16输入Cast hazard与group_index场景死锁 #8667
shilulu创建于 26 天前
修复swiglu_group_quant_grad bf16输入Cast hazard与group_index场景死锁 #8667
已合并
shilulu创建于 26 天前
从shilulu/ops-nn_3710swiglu_group_quant_grad合入到cann/ops-nnmaster
共 3 个文件变更+58-47
Mactivation/swiglu_group_quant_grad/op_host/swiglu_group_quant_grad_tiling_utils.h+11-1
@@ -41,10 +41,11 @@ constexpr uint32_t ZERO = 0;
41constexpr uint32_t ONE = 1;41constexpr uint32_t ONE = 1;
42constexpr uint32_t TWO = 2;42constexpr uint32_t TWO = 2;
43 43 
44-constexpr uint32_t UB_BASE_FACTOR = 12;44+constexpr uint32_t UB_BASE_FACTOR = 14;
45constexpr uint32_t UB_CLAMP_EXTRA_FACTOR = 8;45constexpr uint32_t UB_CLAMP_EXTRA_FACTOR = 8;
46constexpr uint32_t UB_WEIGHT_EXTRA_FACTOR = 4;46constexpr uint32_t UB_WEIGHT_EXTRA_FACTOR = 4;
47constexpr uint32_t UB_WEIGHT_EXTRA_TOKENS = 16;47constexpr uint32_t UB_WEIGHT_EXTRA_TOKENS = 16;
48+constexpr uint32_t UB_CAST_EXTRA_FACTOR = 2;
48 49 
49constexpr uint32_t MAX_H = 4096;50constexpr uint32_t MAX_H = 4096;
50constexpr uint32_t MIN_H = 512;51constexpr uint32_t MIN_H = 512;
@@ -403,12 +404,21 @@ inline void CalculateTilingParams(const gert::TilingContext* context, SwigluGrou
403 uint32_t hasClampLimit = compileInfo.hasClampLimit;404 uint32_t hasClampLimit = compileInfo.hasClampLimit;
404 uint32_t hasWeight = compileInfo.hasWeight;405 uint32_t hasWeight = compileInfo.hasWeight;
405 406 
407+ auto gradYDtype = context->GetInputDesc(INPUT_GRAD_Y_INDEX)->GetDataType();
408+ bool isCastInput = (gradYDtype != ge::DT_FLOAT);
409+ 
406 uint32_t ubFactor = UB_BASE_FACTOR;410 uint32_t ubFactor = UB_BASE_FACTOR;
411+ if (isCastInput) {
412+ ubFactor += UB_CAST_EXTRA_FACTOR;
413+ }
407 if (hasClampLimit) {414 if (hasClampLimit) {
408 ubFactor += UB_CLAMP_EXTRA_FACTOR;415 ubFactor += UB_CLAMP_EXTRA_FACTOR;
409 }416 }
410 if (hasWeight) {417 if (hasWeight) {
411 ubFactor += UB_WEIGHT_EXTRA_FACTOR;418 ubFactor += UB_WEIGHT_EXTRA_FACTOR;
419+ if (isCastInput) {
420+ ubFactor += UB_CAST_EXTRA_FACTOR;
421+ }
412 }422 }
413 423 
414 uint32_t ubAvailable = compileInfo.ubSize - BLOCK_SIZE - TMP_DATA_UB_SIZE;424 uint32_t ubAvailable = compileInfo.ubSize - BLOCK_SIZE - TMP_DATA_UB_SIZE;
Mactivation/swiglu_group_quant_grad/op_kernel/swiglu_group_quant_grad.h+42-42
@@ -100,7 +100,7 @@ __aicore__ inline void SwigluGroupQuantGrad<T>::Init(GM_ADDR gradY, GM_ADDR x, G
100 }100 }
101 ComputeTruncRelatedParams();101 ComputeTruncRelatedParams();
102 102 
103- InitBuffer();103+ InitBuffer(!std::is_same_v<T, float>);
104}104}
105 105 
106template <typename T>106template <typename T>
@@ -118,11 +118,12 @@ __aicore__ inline void SwigluGroupQuantGrad<T>::CopyInGradY(uint32_t tokenIdx, u
118 DataCopyPad(gradYTLocalTensor, gradYGm[gmOffset], gradYCopyParams, padParams);118 DataCopyPad(gradYTLocalTensor, gradYGm[gmOffset], gradYCopyParams, padParams);
119 gradYQueue.EnQue<float>(gradYTLocalTensor);119 gradYQueue.EnQue<float>(gradYTLocalTensor);
120 } else {120 } else {
121- DataCopyPad(gradYTLocalTensor, gradYGm[gmOffset], gradYCopyParams, padParams);121+ uint32_t castSrcBaseIdx = tileLength * sizeof(float) / sizeof(T);
122+ DataCopyPad(gradYTLocalTensor[castSrcBaseIdx], gradYGm[gmOffset], gradYCopyParams, padParams);
122 gradYQueue.EnQue<T>(gradYTLocalTensor);123 gradYQueue.EnQue<T>(gradYTLocalTensor);
123 gradYTLocalTensor = gradYQueue.DeQue<T>();124 gradYTLocalTensor = gradYQueue.DeQue<T>();
124 LocalTensor<float> gradYFloatLocalTensor = gradYTLocalTensor.template ReinterpretCast<float>();125 LocalTensor<float> gradYFloatLocalTensor = gradYTLocalTensor.template ReinterpretCast<float>();
125- Cast(gradYFloatLocalTensor, gradYTLocalTensor, RoundMode::CAST_NONE, computeSize);126+ Cast(gradYFloatLocalTensor, gradYTLocalTensor[castSrcBaseIdx], RoundMode::CAST_NONE, computeSize);
126 PipeBarrier<PIPE_V>();127 PipeBarrier<PIPE_V>();
127 gradYQueue.EnQue<float>(gradYFloatLocalTensor);128 gradYQueue.EnQue<float>(gradYFloatLocalTensor);
128 }129 }
@@ -147,16 +148,16 @@ __aicore__ inline void SwigluGroupQuantGrad<T>::CopyInX(uint32_t tokenIdx, uint3
147 DataCopyPad(xTLocalTensor[tileLength], xGm[x1GmOffset], copyParams, padParams);148 DataCopyPad(xTLocalTensor[tileLength], xGm[x1GmOffset], copyParams, padParams);
148 xQueue.EnQue<float>(xTLocalTensor);149 xQueue.EnQue<float>(xTLocalTensor);
149 } else {150 } else {
151+ uint32_t castSrcBaseIdx = TMP_BUFFER_INDEX * tileLength * sizeof(float) / sizeof(T);
150 uint32_t x0GmOffset = tokenIdx * dim2H + hTileIdx * tileH;152 uint32_t x0GmOffset = tokenIdx * dim2H + hTileIdx * tileH;
151- DataCopyPad(xTLocalTensor, xGm[x0GmOffset], copyParams, padParams);153+ DataCopyPad(xTLocalTensor[castSrcBaseIdx], xGm[x0GmOffset], copyParams, padParams);
152 uint32_t x1GmOffset = tokenIdx * dim2H + dimH + hTileIdx * tileH;154 uint32_t x1GmOffset = tokenIdx * dim2H + dimH + hTileIdx * tileH;
153- DataCopyPad(xTLocalTensor[tileLength * sizeof(float) / sizeof(T)], xGm[x1GmOffset], copyParams, padParams);155+ DataCopyPad(xTLocalTensor[castSrcBaseIdx + tileLength], xGm[x1GmOffset], copyParams, padParams);
154 xQueue.EnQue<T>(xTLocalTensor);156 xQueue.EnQue<T>(xTLocalTensor);
155 xTLocalTensor = xQueue.DeQue<T>();157 xTLocalTensor = xQueue.DeQue<T>();
156 LocalTensor<float> xFloatLocalTensor = xTLocalTensor.template ReinterpretCast<float>();158 LocalTensor<float> xFloatLocalTensor = xTLocalTensor.template ReinterpretCast<float>();
157- Cast(xFloatLocalTensor, xTLocalTensor, RoundMode::CAST_NONE, copySize);159+ Cast(xFloatLocalTensor, xTLocalTensor[castSrcBaseIdx], RoundMode::CAST_NONE, copySize);
158- Cast(xFloatLocalTensor[tileLength], xTLocalTensor[tileLength * sizeof(float) / sizeof(T)], RoundMode::CAST_NONE,160+ Cast(xFloatLocalTensor[tileLength], xTLocalTensor[castSrcBaseIdx + tileLength], RoundMode::CAST_NONE, copySize);
159- copySize);
160 PipeBarrier<PIPE_V>();161 PipeBarrier<PIPE_V>();
161 xQueue.EnQue<float>(xFloatLocalTensor);162 xQueue.EnQue<float>(xFloatLocalTensor);
162 }163 }
@@ -213,11 +214,12 @@ __aicore__ inline void SwigluGroupQuantGrad<T>::CopyInYOrigin(uint32_t tokenIdx,
213 DataCopyPad(yOriginTLocalTensor, yOriginGm[gmOffset], copyParams, padParams);214 DataCopyPad(yOriginTLocalTensor, yOriginGm[gmOffset], copyParams, padParams);
214 yOriginQueue.EnQue<float>(yOriginTLocalTensor);215 yOriginQueue.EnQue<float>(yOriginTLocalTensor);
215 } else {216 } else {
216- DataCopyPad(yOriginTLocalTensor, yOriginGm[gmOffset], copyParams, padParams);217+ uint32_t castSrcBaseIdx = tileLength * sizeof(float) / sizeof(T);
218+ DataCopyPad(yOriginTLocalTensor[castSrcBaseIdx], yOriginGm[gmOffset], copyParams, padParams);
217 yOriginQueue.EnQue<T>(yOriginTLocalTensor);219 yOriginQueue.EnQue<T>(yOriginTLocalTensor);
218 yOriginTLocalTensor = yOriginQueue.DeQue<T>();220 yOriginTLocalTensor = yOriginQueue.DeQue<T>();
219 LocalTensor<float> yOriginFloatLocalTensor = yOriginTLocalTensor.template ReinterpretCast<float>();221 LocalTensor<float> yOriginFloatLocalTensor = yOriginTLocalTensor.template ReinterpretCast<float>();
220- Cast(yOriginFloatLocalTensor, yOriginTLocalTensor, RoundMode::CAST_NONE, computeSize);222+ Cast(yOriginFloatLocalTensor, yOriginTLocalTensor[castSrcBaseIdx], RoundMode::CAST_NONE, computeSize);
221 PipeBarrier<PIPE_V>();223 PipeBarrier<PIPE_V>();
222 yOriginQueue.EnQue<float>(yOriginFloatLocalTensor);224 yOriginQueue.EnQue<float>(yOriginFloatLocalTensor);
223 }225 }
@@ -395,7 +397,7 @@ __aicore__ inline void SwigluGroupQuantGrad<T>::ZeroOutTrunc()
395 if (zeroTokenStart >= totalTokens) {397 if (zeroTokenStart >= totalTokens) {
396 return;398 return;
397 }399 }
398- uint32_t zeroTokenStep = usedCoreNum;400+ uint32_t zeroTokenStep = coreNumAll;
399 401 
400 LocalTensor<T> zeroXLocal = zeroQueue.AllocTensor<T>();402 LocalTensor<T> zeroXLocal = zeroQueue.AllocTensor<T>();
401 LocalTensor<float> zeroXFloatLocal;403 LocalTensor<float> zeroXFloatLocal;
@@ -513,40 +515,38 @@ __aicore__ inline void SwigluGroupQuantGrad<T>::ProcessTile(LocalTensor<float>&
513template <typename T>515template <typename T>
514__aicore__ inline void SwigluGroupQuantGrad<T>::Process()516__aicore__ inline void SwigluGroupQuantGrad<T>::Process()
515{517{
516- if (blockIdx >= usedCoreNum) {518+ if (blockIdx < usedCoreNum) {
517- return;519+ uint32_t tokenEnd = tokenStart + tokensPerCore;
518- }520+ if (tokenEnd > truncValue) {
519- 521+ tokenEnd = truncValue;
520- uint32_t tokenEnd = tokenStart + tokensPerCore;
521- if (tokenEnd > truncValue) {
522- tokenEnd = truncValue;
523- }
524- uint32_t tokenIdx = tokenStart;
525- while (tokenIdx < tokenEnd) {
526- uint32_t currentTileTokens = tileTokens;
527- if (tokenIdx + tileTokens > tokenEnd) {
528- currentTileTokens = tokenEnd - tokenIdx;
529 }522 }
530- 523+ uint32_t tokenIdx = tokenStart;
531- LocalTensor<float> weightLocalTensor;524+ while (tokenIdx < tokenEnd) {
532- if (hasWeight) {525+ uint32_t currentTileTokens = tileTokens;
533- CopyInTopkWeight(tokenIdx, currentTileTokens);526+ if (tokenIdx + tileTokens > tokenEnd) {
534- weightLocalTensor = weightQueue.DeQue<float>();527+ currentTileTokens = tokenEnd - tokenIdx;
535- }
536- 
537- for (uint32_t hTileIdx = 0; hTileIdx < numHTiles; hTileIdx++) {
538- uint32_t currentTileH = tileH;
539- if (hTileIdx == numHTiles - 1) {
540- currentTileH = dimH - hTileIdx * tileH;
541 }528 }
542- ProcessTile(weightLocalTensor, tokenIdx, hTileIdx, currentTileTokens, currentTileH);
543- }
544 529 
545- if (hasWeight) {530+ LocalTensor<float> weightLocalTensor;
546- CopyOutGradWeight(weightLocalTensor, tokenIdx, currentTileTokens);531+ if (hasWeight) {
547- weightQueue.FreeTensor<float>(weightLocalTensor);532+ CopyInTopkWeight(tokenIdx, currentTileTokens);
533+ weightLocalTensor = weightQueue.DeQue<float>();
534+ }
535+ 
536+ for (uint32_t hTileIdx = 0; hTileIdx < numHTiles; hTileIdx++) {
537+ uint32_t currentTileH = tileH;
538+ if (hTileIdx == numHTiles - 1) {
539+ currentTileH = dimH - hTileIdx * tileH;
540+ }
541+ ProcessTile(weightLocalTensor, tokenIdx, hTileIdx, currentTileTokens, currentTileH);
542+ }
543+ 
544+ if (hasWeight) {
545+ CopyOutGradWeight(weightLocalTensor, tokenIdx, currentTileTokens);
546+ weightQueue.FreeTensor<float>(weightLocalTensor);
547+ }
548+ tokenIdx += currentTileTokens;
548 }549 }
549- tokenIdx += currentTileTokens;
550 }550 }
551 SyncAll();551 SyncAll();
552 InitZeroOutTruncBuffer();552 InitZeroOutTruncBuffer();
@@ -555,4 +555,4 @@ __aicore__ inline void SwigluGroupQuantGrad<T>::Process()
555 555 
556} // namespace SwigluGroupQuantGradOp556} // namespace SwigluGroupQuantGradOp
557 557 
558-#endif // SWIGLU_GROUP_QUANT_GRAD_H558+#endif // SWIGLU_GROUP_QUANT_GRAD_H
Mactivation/swiglu_group_quant_grad/op_kernel/swiglu_group_quant_grad_base.h+5-4
@@ -66,9 +66,10 @@ public:
66 tokenStart = blockIdx * tokensPerCore;66 tokenStart = blockIdx * tokensPerCore;
67 }67 }
68 68 
69- __aicore__ inline void InitBuffer()69+ __aicore__ inline void InitBuffer(bool isCastInput)
70 {70 {
71- pipe.InitBuffer(gradYQueue, BUFFER_NUM, tileLength * sizeof(float));71+ uint32_t castBufferFactor = isCastInput ? 2 : 1;
72+ pipe.InitBuffer(gradYQueue, BUFFER_NUM, castBufferFactor * tileLength * sizeof(float));
72 if (hasClampLimit) {73 if (hasClampLimit) {
73 pipe.InitBuffer(xQueue, BUFFER_NUM, tileLength * HAS_CLAMP_SCENE_TILE_NUM * sizeof(float));74 pipe.InitBuffer(xQueue, BUFFER_NUM, tileLength * HAS_CLAMP_SCENE_TILE_NUM * sizeof(float));
74 } else {75 } else {
@@ -78,7 +79,7 @@ public:
78 if (hasWeight) {79 if (hasWeight) {
79 pipe.InitBuffer(weightQueue, BUFFER_NUM,80 pipe.InitBuffer(weightQueue, BUFFER_NUM,
80 AlignUp(tileTokens, FP32_32B_ALIGN_NUM) * sizeof(float) + tileLength * sizeof(float));81 AlignUp(tileTokens, FP32_32B_ALIGN_NUM) * sizeof(float) + tileLength * sizeof(float));
81- pipe.InitBuffer(yOriginQueue, BUFFER_NUM, tileLength * sizeof(float));82+ pipe.InitBuffer(yOriginQueue, BUFFER_NUM, castBufferFactor * tileLength * sizeof(float));
82 }83 }
83 }84 }
84 85 
@@ -143,4 +144,4 @@ protected:
143 144 
144} // namespace SwigluGroupQuantGradOp145} // namespace SwigluGroupQuantGradOp
145 146 
146-#endif // SWIGLU_GROUP_QUANT_GRAD_BASE_H147+#endif // SWIGLU_GROUP_QUANT_GRAD_BASE_H