mindspore 2.3.0rc1__cp39-none-any.whl → 2.3.0rc2__cp39-none-any.whl

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/apply_rotary_pos_emb/kernel/apply_rotary_pos_emb_base.h

@@ -0,0 +1,460 @@
+
+/**
+ * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef ROTARY_POS_EMB_BASE
+#define ROTARY_POS_EMB_BASE
+
+#include "apply_rotary_pos_emb_tiling.h"
+#include "apply_rotary_pos_emb_value.h"
+#include "kernel_operator.h"
+
+template <typename QkDtype, typename CosDtype, bool IF_COS_BROADCAST>
+class RopeBase {
+ public:
+  // QkDtype ：输入qk和输出qk的数据类型
+  // CosDtype ：输入cos/sin的数据类型
+  // IF_COS_BROADCAST ：cos sin是否已扩展
+  // 构造函数
+  __aicore__ inline RopeBase(RopeTilingData *tilingData) {
+    setCtrl_ = get_ctrl();
+#if __CCE_AICORE__ == 220
+    set_ctrl(sbitset0(get_ctrl(), REG_910B));
+#elif __CCE_AICORE__ == 200
+    set_ctrl(sbitset1(get_ctrl(), REG_310P));
+#endif
+    this->tilingData_ = tilingData;
+    batchSize_ = (tilingData_->cosFormat == 0)
+                   ? 0
+                   : ((tilingData_->batch + DEFAULT_REPEAT_STRIDE - 1) / DEFAULT_REPEAT_STRIDE) * DEFAULT_REPEAT_STRIDE;
+    hiddenSize_ =
+      tilingData_->hiddenSizeK > tilingData_->hiddenSizeQ ? tilingData_->hiddenSizeK : tilingData_->hiddenSizeQ;
+    nlCoreRun_ = (tilingData_->ntokens + tilingData_->realCore - 1) / tilingData_->realCore;
+    lCoreRun_ = tilingData_->ntokens - (tilingData_->realCore - 1) * nlCoreRun_;
+    headNum_ = tilingData_->headNumK > tilingData_->headNumQ ? tilingData_->headNumK : tilingData_->headNumQ;
+    rotateStride_ = tilingData_->headDim / tilingData_->rotaryCoeff;
+    dynamicRound_ = (block_idx == tilingData_->realCore - 1) ? lCoreRun_ : nlCoreRun_;
+    rotaryStrideOffset = (tilingData_->headDim == tilingData_->rotaryCoeff) ? 1 : rotateStride_;
+    alignRotary_ = rotateStride_ % ELE_NUM_FP16;
+    pipe_.InitBuffer(seqLenQueue_, 1, (batchSize_ * sizeof(int32_t)));
+  }
+
+  // 初始化Gm
+  __aicore__ inline void RopeInitGm(__gm__ uint8_t *q, __gm__ uint8_t *k, __gm__ uint8_t *cos, __gm__ uint8_t *sin,
+                                    __gm__ uint8_t *seqLen, __gm__ uint8_t *outQ, __gm__ uint8_t *outK) {
+    qGm_ = (__gm__ QkDtype *)q;
+    kGm_ = (__gm__ QkDtype *)k;
+    cosGm_ = (__gm__ CosDtype *)cos;
+    sinGm_ = (__gm__ CosDtype *)sin;
+    outQGm_ = (__gm__ QkDtype *)outQ;
+    outKGm_ = (__gm__ QkDtype *)outK;
+    seqLenGm_ = (__gm__ uint32_t *)seqLen;
+  }
+
+  template <typename T>
+  __aicore__ inline void Copy2Ub(__gm__ T *src, __ubuf__ T *dst, uint32_t copyLen) {
+#if defined(__CCE_KT_TEST__) || (__CCE_AICORE__ == 220)
+    if (g_coreType == AscendC::AIC) return;
+#endif
+    uint32_t blkSizeReal = BLK_SIZE / sizeof(T);
+    if (copyLen % blkSizeReal != 0) {
+      copy_gm_to_ubuf(dst, src, 0, 1, (copyLen + blkSizeReal - 1) / blkSizeReal, 0, 0);
+      pipe_barrier((PIPE_ALL));
+    } else {
+      copy_gm_to_ubuf(dst, src, 0, 1, copyLen / blkSizeReal, 0, 0);
+      pipe_barrier((PIPE_ALL));
+    }
+  }
+
+  template <typename T>
+  __aicore__ inline void Copy2Gm(__ubuf__ T *src, __gm__ T *dst, uint32_t hiddenSizeLen) {
+#if defined(__CCE_KT_TEST__) || (__CCE_AICORE__ == 220)
+    if (g_coreType == AscendC::AIC) return;
+#endif
+    uint32_t blkSizeReal = BLK_SIZE / sizeof(T);
+    if (hiddenSizeLen % blkSizeReal != 0) {
+      copy_ubuf_to_gm(dst, src, 0, 1, (hiddenSizeLen + blkSizeReal - 1) / blkSizeReal, 0, 0);
+    } else {
+      copy_ubuf_to_gm(dst, src, 0, 1, hiddenSizeLen / blkSizeReal, 0, 0);
+    }
+  }
+
+  // 此函数用来复用unpad情況下的cos和sin
+  // 例：cos[0~7] cos[0~3]用于第一个batch， cos[0~4]用于第二个batch
+  __aicore__ inline void ExpandCosSin(__ubuf__ CosDtype *tempBuf, __gm__ CosDtype *src, __gm__ CosDtype *extraGm) {
+#if defined(__CCE_KT_TEST__) || (__CCE_AICORE__ == 220)
+    if (g_coreType == AscendC::AIC) return;
+#endif
+    // cos or sin,[maxseqlen,headsize]-->[sumseqlen,hiddensize]
+    AscendC::LocalTensor<int32_t> seqLenLocal = seqLenQueue_.AllocTensor<int32_t>();
+    // copy_gm_to_ubuf((__ubuf__ int32_t *)seqLenLocal.GetPhyAddr(), seqLenGm_, 0, 1,
+    //     batchSize_ * sizeof(int32_t) / 32, 0, 0);
+    int32_t seqLenTmp = this->tilingData_->ntokens / this->tilingData_->batch;
+    for (uint32_t i = 0; i < this->tilingData_->batch; i++) {
+      seqLenLocal.SetValue(i, seqLenTmp);
+    }
+    pipe_barrier((PIPE_ALL));
+    int32_t rowsPerLoop = (maxProcessNum_ - batchSize_ * NUM_TWO) / tilingData_->headDim;
+    int32_t cosoffset = 0;
+    for (uint32_t perBatch = 0; perBatch < tilingData_->batch; perBatch++) {
+      int32_t rowsRepeat = seqLenLocal.GetValue(perBatch) / rowsPerLoop;
+      int32_t rowsRemain = seqLenLocal.GetValue(perBatch) % rowsPerLoop;
+      for (int32_t j = 0; j < rowsRepeat; j++) {
+        Copy2Ub(src + (j * rowsPerLoop) * tilingData_->headDim, tempBuf, rowsPerLoop * tilingData_->headDim);
+        Copy2Gm(tempBuf, (extraGm + (cosoffset + j * rowsPerLoop) * tilingData_->headDim),
+                rowsPerLoop * tilingData_->headDim);
+        pipe_barrier((PIPE_ALL));
+      }
+      if (rowsRemain > 0) {
+        Copy2Ub(src + (rowsRepeat * rowsPerLoop) * tilingData_->headDim, tempBuf, rowsRemain * tilingData_->headDim);
+        Copy2Gm(tempBuf, (extraGm + (cosoffset + rowsRepeat * rowsPerLoop) * tilingData_->headDim),
+                rowsRemain * tilingData_->headDim);
+        pipe_barrier((PIPE_ALL));
+      }
+      cosoffset += seqLenLocal.GetValue(perBatch);
+    }
+    seqLenQueue_.FreeTensor(seqLenLocal);
+    pipe_barrier((PIPE_ALL));
+  }
+
+  // 构建tensor -1 -1 -1 0 0 0
+  // 构建tensor 0 0 0 1 1 1
+  template <typename BUF_TYPE>
+  __aicore__ inline void ExpandNeg(__ubuf__ BUF_TYPE *tempBuf, uint32_t bufPos, uint32_t headNumTemp,
+                                   uint32_t repeatTimeTemp) {
+    if (tilingData_->headDim != tilingData_->rotaryCoeff) {
+      if (alignRotary_ == 0) {  // 对齐直接 -1 1
+        for (uint32_t i = 0; i < rotateStride_; ++i) {
+          *(tempBuf + negOne_ + i) = (BUF_TYPE)-1;
+          *(tempBuf + negOne_ + i + rotateStride_) = (BUF_TYPE)1;
+        }
+        set_flag(PIPE_S, PIPE_V, EVENT_ID1);
+        wait_flag(PIPE_S, PIPE_V, EVENT_ID1);
+        for (uint32_t i = 1; i < headNumTemp * tilingData_->rotaryCoeff / NUM_TWO; ++i) {
+          // halfHeadDim = rotateStride_ * 2
+          copy_ubuf_to_ubuf(tempBuf + negOne_ + rotateStride_ * NUM_TWO * i, tempBuf + negOne_, 0, 1,
+                            rotateStride_ * sizeof(BUF_TYPE) / ELE_NUM_FP16, 0, 0);
+        }
+      } else {
+        for (uint32_t i = 0; i < rotateStride_; ++i) {  // 非对齐 -1 0
+          *(tempBuf + negOne_ + i) = (BUF_TYPE)-1;
+          *(tempBuf + negOne_ + i + rotateStride_) = (BUF_TYPE)0;
+        }
+        set_flag(PIPE_S, PIPE_V, EVENT_ID1);
+        wait_flag(PIPE_S, PIPE_V, EVENT_ID1);
+        for (uint32_t i = 0; i < headNumTemp * tilingData_->rotaryCoeff / NUM_TWO; ++i) {
+          if ((rotateStride_ * NUM_TWO) * sizeof(BUF_TYPE) % BLK_SIZE == 0) {
+            copy_ubuf_to_ubuf(tempBuf + negOne_ + rotateStride_ * NUM_TWO * i, tempBuf + negOne_, 0, 1,
+                              rotateStride_ * NUM_TWO * sizeof(BUF_TYPE) / ELE_NUM_FP16, 0, 0);
+          } else {
+            for (uint32_t j = 0; j < rotateStride_ * NUM_TWO; j++) {
+              *(tempBuf + negOne_ + rotateStride_ * NUM_TWO * i + j) = *(tempBuf + negOne_ + j);
+            }
+          }
+        }
+        set_flag(PIPE_S, PIPE_V, EVENT_ID1);
+        wait_flag(PIPE_S, PIPE_V, EVENT_ID1);
+        pipe_barrier(PIPE_V);
+        vadds(tempBuf + bufPos, tempBuf + negOne_, (BUF_TYPE)1, repeatTimeTemp, 1, 1, DEFAULT_REPEAT_STRIDE,
+              DEFAULT_REPEAT_STRIDE);
+      }
+    } else {
+      set_vector_mask((uint64_t)-1, (uint64_t)-1);
+      vector_dup(tempBuf + negOne_, (BUF_TYPE)-1.0, repeatTimeTemp, 1, 1, (uint16_t)DEFAULT_REPEAT_STRIDE,
+                 (uint16_t)DEFAULT_REPEAT_STRIDE);
+      set_vector_mask(0xaaaaaaaaaaaaaaaa, 0xaaaaaaaaaaaaaaaa);
+      vector_dup(tempBuf + negOne_, (BUF_TYPE)0.0, repeatTimeTemp, 1, 1, (uint16_t)DEFAULT_REPEAT_STRIDE,
+                 (uint16_t)DEFAULT_REPEAT_STRIDE);
+      set_vector_mask((uint64_t)-1, (uint64_t)-1);
+      pipe_barrier((PIPE_V));
+      vadds(tempBuf + bufPos, tempBuf + negOne_, (BUF_TYPE)1, repeatTimeTemp, 1, 1, DEFAULT_REPEAT_STRIDE,
+            DEFAULT_REPEAT_STRIDE);
+    }
+  }
+
+  // 从(tilingData_->headDim)->(heads*tilingData_->headDim)
+  __aicore__ inline void CosSinCommonBroardcast(__gm__ uint8_t *extraGm, uint32_t z, __ubuf__ CosDtype *tempBuf,
+                                                uint32_t calcLen) {
+    // 永远的先拷一次
+    uint32_t cosOffset = block_idx * nlCoreRun_ * tilingData_->headDim + z * tilingData_->headDim;
+    uint32_t sinOffset = block_idx * nlCoreRun_ * tilingData_->headDim + z * tilingData_->headDim;
+    set_flag(PIPE_S, PIPE_MTE2, EVENT_ID1);
+    wait_flag(PIPE_S, PIPE_MTE2, EVENT_ID1);
+    copy_gm_to_ubuf(tempBuf + cosPad_, cosGm_ + cosOffset, 0, 1,
+                    (tilingData_->headDim * sizeof(CosDtype) + BLK_SIZE - 1) / BLK_SIZE, 0, 0);
+    copy_gm_to_ubuf(tempBuf + sinPad_, sinGm_ + sinOffset, 0, 1,
+                    (tilingData_->headDim * sizeof(CosDtype) + BLK_SIZE - 1) / BLK_SIZE, 0, 0);
+    if (tilingData_->cosFormat == 1) {
+      pipe_barrier(PIPE_ALL);
+    }
+    set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID3);
+    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3);
+    if ((tilingData_->headDim * sizeof(CosDtype)) % BLK_SIZE != 0) {
+      wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID3);
+      // 补齐cos，从(tilingData_->headDim)->(heads*tilingData_->headDim)
+      // headnum
+      for (uint32_t i = 0; i < calcLen / tilingData_->headDim; ++i) {
+        copy_ubuf_to_gm((__gm__ CosDtype *)extraGm + offsetExtraGm_ + tilingData_->headDim * i, tempBuf + cosPad_, 0, 1,
+                        (tilingData_->headDim * sizeof(CosDtype) + BLK_SIZE - 1) / BLK_SIZE, 0, 0);
+        pipe_barrier((PIPE_ALL));
+      }
+      Copy2Ub<CosDtype>((__gm__ CosDtype *)extraGm + offsetExtraGm_, tempBuf + cosPad_, calcLen);
+      // 补齐sin，从(tilingData_->headDim)->(heads*tilingData_->headDim)
+      for (uint32_t i = 0; i < calcLen / tilingData_->headDim; ++i) {
+        copy_ubuf_to_gm((__gm__ CosDtype *)extraGm + offsetExtraGm_ + tilingData_->headDim * i, tempBuf + sinPad_, 0, 1,
+                        (tilingData_->headDim * sizeof(CosDtype) + BLK_SIZE - 1) / BLK_SIZE, 0, 0);
+        pipe_barrier((PIPE_ALL));
+      }
+      Copy2Ub<CosDtype>((__gm__ CosDtype *)extraGm + offsetExtraGm_, tempBuf + sinPad_, calcLen);
+      wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3);
+    } else {
+      wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3);
+      for (uint32_t i = 1; i < calcLen / tilingData_->headDim; ++i) {
+        copy_ubuf_to_ubuf(tempBuf + cosPad_ + tilingData_->headDim * i, tempBuf + cosPad_, 0, 1,
+                          tilingData_->headDim * sizeof(CosDtype) / BLK_SIZE, 0, 0);
+        copy_ubuf_to_ubuf(tempBuf + sinPad_ + tilingData_->headDim * i, tempBuf + sinPad_, 0, 1,
+                          tilingData_->headDim * sizeof(CosDtype) / BLK_SIZE, 0, 0);
+      }
+      wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID3);
+    }
+  }
+
+  // 满足 cos sin 多头输入
+  template <typename BUF_TYPE>
+  __aicore__ inline void CosSinBroadcast(__gm__ uint8_t *extraGm, uint32_t z, __ubuf__ BUF_TYPE *tempBuf,
+                                         uint32_t Calclen) {
+    if constexpr (IF_COS_BROADCAST) {
+      copy_gm_to_ubuf(tempBuf + cosPad_,
+                      cosGm_ + block_idx * nlCoreRun_ * tilingData_->hiddenSizeQ + z * tilingData_->hiddenSizeQ, 0, 1,
+                      Calclen * sizeof(BUF_TYPE) / BLK_SIZE, 0, 0);
+      copy_gm_to_ubuf(tempBuf + sinPad_,
+                      sinGm_ + block_idx * nlCoreRun_ * tilingData_->hiddenSizeQ + z * tilingData_->hiddenSizeQ, 0, 1,
+                      Calclen * sizeof(BUF_TYPE) / BLK_SIZE, 0, 0);
+    } else {
+      CosSinCommonBroardcast(extraGm, z, tempBuf, Calclen);
+    }
+  }
+
+  // qk 公用函数
+  template <typename BUF_TYPE>
+  __aicore__ inline void QkComm(__gm__ BUF_TYPE *src, __gm__ uint8_t *extraGm1, uint32_t hiddenSizeTmp,
+                                __ubuf__ BUF_TYPE *tempBuf, uint32_t headNumTemp) {
+    uint32_t hiddenSizeBlk = hiddenSizeTmp / ELE_NUM_FP16;
+    set_flag(PIPE_S, PIPE_MTE2, EVENT_ID1);
+    wait_flag(PIPE_S, PIPE_MTE2, EVENT_ID1);
+    copy_gm_to_ubuf(tempBuf + oriPos_,  // gm -> ub
+                    src, 0, 1, hiddenSizeBlk, 0, 0);
+    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+    set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID2);
+    if (alignRotary_ == 0) {
+      wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+      wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID2);
+      copy_ubuf_to_ubuf(tempBuf + removeBefore_ + rotaryStrideOffset, tempBuf + oriPos_, 0,
+                        headNumTemp * tilingData_->rotaryCoeff / 2, rotaryStrideOffset / ELE_NUM_FP16,
+                        rotaryStrideOffset / ELE_NUM_FP16, rotaryStrideOffset / ELE_NUM_FP16);
+
+      copy_ubuf_to_ubuf(tempBuf + removeBefore_, tempBuf + oriPos_ + rotaryStrideOffset, 0,
+                        headNumTemp * tilingData_->rotaryCoeff / 2, rotaryStrideOffset / ELE_NUM_FP16,
+                        rotaryStrideOffset / ELE_NUM_FP16, rotaryStrideOffset / ELE_NUM_FP16);
+    } else {
+      wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+      wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID2);
+      // ub -> workspace[0~hiddensize]
+      copy_ubuf_to_gm((__gm__ BUF_TYPE *)extraGm1 + offsetExtraGm_, tempBuf + oriPos_, 0, 1, hiddenSizeBlk, 0, 0);
+      // ub -> workspace[hiddensize ~ 2 * hiddensize]
+      copy_ubuf_to_gm((__gm__ BUF_TYPE *)extraGm1 + offsetExtraGm_ + hiddenSizeTmp, tempBuf + oriPos_, 0, 1,
+                      hiddenSizeBlk, 0, 0);
+      // workspace[rotary ~ hiddensize + rotary] -> ub[hiddensize ~ 2 * hiddensize]
+      pipe_barrier((PIPE_ALL));
+      copy_gm_to_ubuf(tempBuf + removeBefore_, (__gm__ BUF_TYPE *)extraGm1 + offsetExtraGm_ + rotateStride_, 0, 1,
+                      hiddenSizeBlk, 0, 0);
+      // gm[hiddensize - rotary ~ 2 * hiddensize - rotary] -> ub[2 *hiddensize ~ 3 * hiddensize]
+      copy_gm_to_ubuf(tempBuf + padBefore_,
+                      (__gm__ BUF_TYPE *)extraGm1 + offsetExtraGm_ + hiddenSizeTmp - rotateStride_, 0, 1, hiddenSizeBlk,
+                      0, 0);
+    }
+  }
+
+  // 主体计算逻辑
+  template <typename BUF_TYPE>
+  __aicore__ inline void CalcRope(__ubuf__ BUF_TYPE *tempBuf, uint32_t repeatTimes1, uint32_t oriPosTemp,
+                                  uint32_t removeTemp, uint32_t padTemp, uint32_t posTemp, uint32_t res) {
+    set_vector_mask((uint64_t)-1, (uint64_t)-1);
+#if defined(__CCE_KT_TEST__) || (__CCE_AICORE__ == 220)
+    if (g_coreType == AscendC::AIC) return;
+#endif
+
+    vmul(tempBuf + oriPosTemp, tempBuf + cosPad_, tempBuf + oriPosTemp,
+         repeatTimes1,           // repeat times
+         1,                      // dstBlockStride
+         1,                      // src0BlockStride
+         1,                      // src1BlockStride
+         DEFAULT_REPEAT_STRIDE,  // dstRepeatStride
+         DEFAULT_REPEAT_STRIDE,  // src0RepeatStride
+         DEFAULT_REPEAT_STRIDE   // src1RepeatStride
+    );
+    vmul(tempBuf + padTemp, tempBuf + posTemp, tempBuf + padTemp,
+         repeatTimes1,           // repeat times
+         1,                      // dstBlockStride
+         1,                      // src0BlockStride
+         1,                      // src1BlockStride
+         DEFAULT_REPEAT_STRIDE,  // dstRepeatStride
+         DEFAULT_REPEAT_STRIDE,  // src0RepeatStride
+         DEFAULT_REPEAT_STRIDE   // src1RepeatStride
+    );
+    pipe_barrier((PIPE_V));
+    vmul(tempBuf + removeTemp, tempBuf + sinPad_, tempBuf + removeTemp,
+         repeatTimes1,           // repeat times
+         1,                      // dstBlockStride
+         1,                      // src0BlockStride
+         1,                      // src1BlockStride
+         DEFAULT_REPEAT_STRIDE,  // dstRepeatStride
+         DEFAULT_REPEAT_STRIDE,  // src0RepeatStride
+         DEFAULT_REPEAT_STRIDE   // src1RepeatStride
+    );
+    vmul(tempBuf + padTemp, tempBuf + sinPad_, tempBuf + padTemp,
+         repeatTimes1,           // repeat times
+         1,                      // dstBlockStride
+         1,                      // src0BlockStride
+         1,                      // src1BlockStride
+         DEFAULT_REPEAT_STRIDE,  // dstRepeatStride
+         DEFAULT_REPEAT_STRIDE,  // src0RepeatStride
+         DEFAULT_REPEAT_STRIDE   // src1RepeatStride
+    );
+    pipe_barrier((PIPE_V));
+
+    vmul(tempBuf + removeTemp, tempBuf + negOne_, tempBuf + removeTemp,
+         repeatTimes1,           // repeat times
+         1,                      // dstBlockStride
+         1,                      // src0BlockStride
+         1,                      // src1BlockStride
+         DEFAULT_REPEAT_STRIDE,  // dstRepeatStride
+         DEFAULT_REPEAT_STRIDE,  // src0RepeatStride
+         DEFAULT_REPEAT_STRIDE   // src1RepeatStride
+    );
+    vadd(tempBuf + padTemp, tempBuf + oriPosTemp, tempBuf + padTemp,
+         repeatTimes1,           // repeat times
+         1,                      // dstBlockStride
+         1,                      // src0BlockStride
+         1,                      // src1BlockStride
+         DEFAULT_REPEAT_STRIDE,  // dstRepeatStride
+         DEFAULT_REPEAT_STRIDE,  // src0RepeatStride
+         DEFAULT_REPEAT_STRIDE   // src1RepeatStride
+    );
+    pipe_barrier((PIPE_V));
+
+    vadd(tempBuf + res, tempBuf + removeTemp, tempBuf + padTemp,
+         repeatTimes1,           // repeat times
+         1,                      // dstBlockStride
+         1,                      // src0BlockStride
+         1,                      // src1BlockStride
+         DEFAULT_REPEAT_STRIDE,  // dstRepeatStride
+         DEFAULT_REPEAT_STRIDE,  // src0RepeatStride
+         DEFAULT_REPEAT_STRIDE   // src1RepeatStride
+    );
+    pipe_barrier((PIPE_V));
+    set_ctrl(setCtrl_);
+  }
+
+  // 主体计算逻辑
+  template <typename BUF_TYPE>
+  __aicore__ inline void CalcRopeAlign(__ubuf__ BUF_TYPE *tempBuf, uint32_t repeatTimes1, uint32_t oriPosTemp,
+                                       uint32_t removeTemp, uint32_t padTemp) {
+    set_vector_mask((uint64_t)-1, (uint64_t)-1);
+#if defined(__CCE_KT_TEST__) || (__CCE_AICORE__ == 220)
+    if (g_coreType == AscendC::AIC) return;
+#endif
+    vmul(tempBuf + oriPosTemp, tempBuf + cosPad_, tempBuf + oriPosTemp,
+         repeatTimes1,           // repeat times
+         1,                      // dstBlockStride
+         1,                      // src0BlockStride
+         1,                      // src1BlockStride
+         DEFAULT_REPEAT_STRIDE,  // dstRepeatStride
+         DEFAULT_REPEAT_STRIDE,  // src0RepeatStride
+         DEFAULT_REPEAT_STRIDE   // src1RepeatStride
+    );
+    vmul(tempBuf + removeTemp, tempBuf + negOne_, tempBuf + removeTemp,
+         repeatTimes1,           // repeat times
+         1,                      // dstBlockStride
+         1,                      // src0BlockStride
+         1,                      // src1BlockStride
+         DEFAULT_REPEAT_STRIDE,  // dstRepeatStride
+         DEFAULT_REPEAT_STRIDE,  // src0RepeatStride
+         DEFAULT_REPEAT_STRIDE   // src1RepeatStride
+    );
+    pipe_barrier((PIPE_V));
+    vmul(tempBuf + removeTemp, tempBuf + sinPad_, tempBuf + removeTemp,
+         repeatTimes1,           // repeat times
+         1,                      // dstBlockStride
+         1,                      // src0BlockStride
+         1,                      // src1BlockStride
+         DEFAULT_REPEAT_STRIDE,  // dstRepeatStride
+         DEFAULT_REPEAT_STRIDE,  // src0RepeatStride
+         DEFAULT_REPEAT_STRIDE   // src1RepeatStride
+    );
+    pipe_barrier((PIPE_V));
+    vadd(tempBuf + padTemp, tempBuf + removeTemp, tempBuf + oriPosTemp,
+         repeatTimes1,           // repeat times
+         1,                      // dstBlockStride
+         1,                      // src0BlockStride
+         1,                      // src1BlockStride
+         DEFAULT_REPEAT_STRIDE,  // dstRepeatStride
+         DEFAULT_REPEAT_STRIDE,  // src0RepeatStride
+         DEFAULT_REPEAT_STRIDE   // src1RepeatStride
+    );
+    pipe_barrier((PIPE_V));
+    set_ctrl(setCtrl_);
+  }
+
+ public:
+  RopeTilingData *tilingData_ = nullptr;
+  __gm__ QkDtype *qGm_{nullptr};
+  __gm__ QkDtype *kGm_{nullptr};
+  __gm__ CosDtype *cosGm_{nullptr};
+  __gm__ CosDtype *sinGm_{nullptr};
+  __gm__ uint32_t *seqLenGm_{nullptr};
+  __gm__ QkDtype *outQGm_{nullptr};
+  __gm__ QkDtype *outKGm_{nullptr};
+  AscendC::TPipe pipe_;
+  AscendC::TQue<AscendC::QuePosition::VECIN, 1> seqLenQueue_;
+
+  uint32_t cosPad_{0};             // broadcast的cos在uB中的位置
+  uint32_t sinPad_{0};             // broadcast的sin在uB中的位置
+  uint32_t negOne_{0};             // -1 -1 -1 0 0 0在uB中的位置
+  uint32_t oriPos_{0};             // q,k在uB中的位置
+  uint32_t padBefore_{0};          // 保存qk[-x : hiddensize - x]
+  uint32_t removeBefore_{0};       // 保存qk[x : hiddensize + x]
+  uint32_t repeatSize_{0};         // 一拍做几个元素
+  uint32_t maxProcessNum_{0};      // 最大处理元素个数
+  uint32_t repeatTimesQ_{0};       // q重复次数
+  uint32_t repeatTimesK_{0};       // k重复次数
+  uint32_t hiddenSizeAlign_{0};    // 对齐后的hiddensize
+  uint32_t repeatTimes_{0};        // 对齐后重复次数
+  uint32_t headNum_{0};            // 几个头
+  uint32_t hiddenSize_{0};         // hiddensizeQ,K的最大值
+  uint32_t nlCoreRun_{0};          // 非最后一个核需要跑几次
+  uint32_t lCoreRun_{0};           // 最后一个核需要跑几次
+  uint32_t batchSize_{0};          // batch向上取整
+  uint32_t rotateStride_{0};       // headdim / 旋转系数
+  uint32_t offsetExtraGm_{0};      // 使用workspace需要的offset
+  uint32_t dynamicRound_{0};       // 每个核做几轮
+  uint32_t setCtrl_;               // 复位寄存器
+  uint32_t alignHalfHeadDim_{0};   // headDim / 旋转系数 * 2 是否对齐
+  uint32_t rotaryStrideOffset{0};  // 每次旋转长度
+  uint32_t alignRotary_;           // 旋转距离是否对齐
+  uint32_t syncOffset_;            // 每个核使用workspace的offset
+};
+
+#endif

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/apply_rotary_pos_emb/kernel/apply_rotary_pos_emb_bf16.h

@@ -0,0 +1,217 @@
+/**
+ * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef ROTARY_POS_EMB_BF16
+#define ROTARY_POS_EMB_BF16
+#include "apply_rotary_pos_emb_base.h"
+template <typename QK_DTYPE, typename COS_DTYPE, bool IF_COS_BROADCAST>
+class RopeBf16 : public RopeBase<QK_DTYPE, COS_DTYPE, IF_COS_BROADCAST> {
+ public:
+  __aicore__ inline RopeBf16(RopeTilingData *tilingData) : RopeBase<QK_DTYPE, COS_DTYPE, IF_COS_BROADCAST>(tilingData) {
+    this->repeatSize_ = 64;                   // 64 = 256B / sizeof(float)
+    this->maxProcessNum_ = 3 * MAX_LEN_FP16;  // 3 is fp16 space needed
+    this->repeatTimesQ_ = (this->tilingData_->hiddenSizeQ + this->repeatSize_ - 1) / this->repeatSize_;
+    this->repeatTimesK_ = (this->tilingData_->hiddenSizeK + this->repeatSize_ - 1) / this->repeatSize_;
+    headDimAlign_ = ((this->tilingData_->headDim + ELE_NUM_FP16 - 1) / ELE_NUM_FP16) * ELE_NUM_FP16;
+    this->alignHalfHeadDim_ = (this->rotateStride_ * NUM_TWO) % ELE_NUM_FP32;
+    this->hiddenSizeAlign_ = ((this->hiddenSize_ + this->repeatSize_ - 1) / this->repeatSize_) * this->repeatSize_;
+    this->syncOffset_ =
+      (this->tilingData_->headDim % ELE_NUM_FP16 == 0) ? this->hiddenSizeAlign_ : this->headNum_ * headDimAlign_;
+    this->offsetExtraGm_ = NUM_TWO * block_idx * this->syncOffset_;
+
+    sliceSizeTmp_ = (SLICE_SIZE / this->tilingData_->headDim) * this->tilingData_->headDim;  // 向下取整
+
+    // fp16
+    this->oriPos_ = 0;
+    this->removeBefore_ = this->oriPos_ + sliceSizeTmp_;
+    this->padBefore_ = this->removeBefore_ + sliceSizeTmp_;
+
+    // fp32
+    this->cosPad_ = 0;
+    this->sinPad_ = this->cosPad_ + sliceSizeTmp_;
+    this->negOne_ = this->sinPad_ + sliceSizeTmp_;
+    oriPosF32_ = this->negOne_ + sliceSizeTmp_;
+    PadBeforeF32_ = oriPosF32_ + sliceSizeTmp_;
+    removeBeforeF32_ = PadBeforeF32_ + sliceSizeTmp_;
+    posOneF32_ = removeBeforeF32_ + sliceSizeTmp_;
+
+    this->pipe_.InitBuffer(
+      qkfp32QueueCO2_, 1,
+      (this->tilingData_->maxUbSize - (this->batchSize_ + this->maxProcessNum_) * sizeof(half)));  // 留給fp32的
+    AscendC::LocalTensor<float> qkfp32_perloop_ub = qkfp32QueueCO2_.AllocTensor<float>();
+    qkfp32Ubuf_ = (__ubuf__ float *)qkfp32_perloop_ub.GetPhyAddr();
+    this->pipe_.InitBuffer(outQueueCO2_, 1, ((this->maxProcessNum_) * sizeof(half)));
+    AscendC::LocalTensor<QK_DTYPE> cache_perloop_ub2 = outQueueCO2_.AllocTensor<QK_DTYPE>();
+    commonUbuf_ = (__ubuf__ QK_DTYPE *)cache_perloop_ub2.GetPhyAddr();
+
+    // 判断是否需要切块计算
+    if (this->tilingData_->hiddenSizeQ > sliceSizeTmp_) {
+      sliceTimeQ_ = (this->tilingData_->hiddenSizeQ + sliceSizeTmp_ - 1) / sliceSizeTmp_;    // 向上取整
+      lastSliceSizeQ_ = this->tilingData_->hiddenSizeQ - (sliceTimeQ_ - 1) * sliceSizeTmp_;  // 1024
+    } else {
+      sliceTimeQ_ = 1;
+      lastSliceSizeQ_ = this->tilingData_->hiddenSizeQ;
+    }
+
+    if (this->tilingData_->hiddenSizeK > sliceSizeTmp_) {
+      sliceTimeK_ = (this->tilingData_->hiddenSizeK + sliceSizeTmp_ - 1) / sliceSizeTmp_;  // 向上取整
+      lastSliceSizeK_ = this->tilingData_->hiddenSizeK - (sliceTimeK_ - 1) * sliceSizeTmp_;
+    } else {
+      sliceTimeK_ = 1;
+      lastSliceSizeK_ = this->tilingData_->hiddenSizeK;
+    }
+  }
+
+  __aicore__ inline void ConvertCos(uint32_t repeatTimes) {
+    vconv_bf162f32(qkfp32Ubuf_ + this->cosPad_, commonUbuf_ + this->cosPad_, repeatTimes, 1, 1, DEFAULT_REPEAT_STRIDE,
+                   DEFAULT_REPEAT_STRIDE / NUM_TWO);
+    vconv_bf162f32(qkfp32Ubuf_ + this->sinPad_, commonUbuf_ + this->sinPad_, repeatTimes, 1, 1, DEFAULT_REPEAT_STRIDE,
+                   DEFAULT_REPEAT_STRIDE / NUM_TWO);
+  }
+
+  __aicore__ inline void CastB162F32(uint32_t repeatTimes1) {
+    vconv_bf162f32(qkfp32Ubuf_ + oriPosF32_, commonUbuf_ + this->oriPos_, repeatTimes1, 1, 1, DEFAULT_REPEAT_STRIDE,
+                   DEFAULT_REPEAT_STRIDE / NUM_TWO);
+    vconv_bf162f32(qkfp32Ubuf_ + removeBeforeF32_, commonUbuf_ + this->removeBefore_, repeatTimes1, 1, 1,
+                   DEFAULT_REPEAT_STRIDE, DEFAULT_REPEAT_STRIDE / NUM_TWO);
+    vconv_bf162f32(qkfp32Ubuf_ + PadBeforeF32_, commonUbuf_ + this->padBefore_, repeatTimes1, 1, 1,
+                   DEFAULT_REPEAT_STRIDE, DEFAULT_REPEAT_STRIDE / NUM_TWO);
+  }
+
+  __aicore__ inline void CastF322B16(__gm__ QK_DTYPE *dst, __ubuf__ QK_DTYPE *src1, __ubuf__ float *src,
+                                     uint32_t repeatTimes1, uint32_t hiddenSize1) {
+    vconv_f322bf16r(src1, src, repeatTimes1, 1, 1, DEFAULT_REPEAT_STRIDE / NUM_TWO, DEFAULT_REPEAT_STRIDE);
+    set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
+    wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
+    copy_ubuf_to_gm(dst, src1, 0, 1, hiddenSize1 / ELE_NUM_FP16, 0, 0);
+  }
+
+  __aicore__ inline void Process(__gm__ uint8_t *extraGm) {
+    if (this->tilingData_->cosFormat == 1) {
+      pipe_barrier((PIPE_ALL));
+      this->ExpandCosSin(commonUbuf_, this->cosGm_, (__gm__ COS_DTYPE *)extraGm);
+      this->cosGm_ = (__gm__ COS_DTYPE *)extraGm;
+      pipe_barrier((PIPE_ALL));
+      this->ExpandCosSin(commonUbuf_, this->sinGm_,
+                         (__gm__ COS_DTYPE *)extraGm + this->tilingData_->ntokens * this->tilingData_->headDim);
+      this->sinGm_ = (__gm__ COS_DTYPE *)extraGm + this->tilingData_->ntokens * this->tilingData_->headDim;
+      extraGm =
+        extraGm + this->tilingData_->ntokens * this->tilingData_->headDim * 4;  // sizeof(uint8_t) * 2 = sizeof(half)
+      pipe_barrier((PIPE_ALL));
+    }
+
+    uint32_t dynamicSliceQ =
+      this->tilingData_->hiddenSizeQ > sliceSizeTmp_ ? sliceSizeTmp_ : this->tilingData_->hiddenSizeQ;
+    uint32_t headNumTempQ = dynamicSliceQ / this->tilingData_->headDim;
+
+    uint32_t dynamicSliceK =
+      this->tilingData_->hiddenSizeK > sliceSizeTmp_ ? sliceSizeTmp_ : this->tilingData_->hiddenSizeK;
+    uint32_t headNumTempK = dynamicSliceK / this->tilingData_->headDim;
+    uint32_t repeatTemp = (dynamicSliceQ + this->repeatSize_ - 1) / this->repeatSize_;
+    this->ExpandNeg(qkfp32Ubuf_, posOneF32_, headNumTempQ, repeatTemp);
+    for (uint32_t zz = 0; zz < this->dynamicRound_; ++zz) {
+      this->CosSinBroadcast(extraGm, zz, commonUbuf_, dynamicSliceQ);
+      if (this->tilingData_->headDim % ELE_NUM_FP16 == 0) {
+        pipe_barrier(PIPE_V);
+        ConvertCos(repeatTemp);
+      } else {
+        set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+        wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+        ConvertCos(repeatTemp);
+      }
+      for (uint32_t perSlice = 0; perSlice < sliceTimeQ_; ++perSlice) {  // 核内每块
+        uint32_t dynamicSliceQTemp = (perSlice == sliceTimeQ_ - 1) ? lastSliceSizeQ_ : sliceSizeTmp_;
+        headNumTempQ = dynamicSliceQTemp / this->tilingData_->headDim;
+        uint32_t repeatTimeOnce = (dynamicSliceQTemp + this->repeatSize_ - 1) / this->repeatSize_;
+        pipe_barrier(PIPE_MTE2);
+        set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
+        wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
+        this->QkComm(this->qGm_ + block_idx * this->nlCoreRun_ * this->tilingData_->hiddenSizeQ +
+                       zz * this->tilingData_->hiddenSizeQ + perSlice * sliceSizeTmp_,
+                     extraGm, dynamicSliceQTemp, commonUbuf_, headNumTempQ);
+
+        if (this->alignRotary_ == 0) {
+          pipe_barrier((PIPE_V));
+          CastB162F32(repeatTimeOnce);
+
+          pipe_barrier((PIPE_V));
+          this->CalcRopeAlign(qkfp32Ubuf_, repeatTimeOnce, oriPosF32_, removeBeforeF32_, PadBeforeF32_);
+        } else {
+          set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+          wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+
+          CastB162F32(repeatTimeOnce);
+          pipe_barrier((PIPE_V));
+          this->CalcRope(qkfp32Ubuf_, repeatTimeOnce, oriPosF32_, removeBeforeF32_, PadBeforeF32_, posOneF32_,
+                         PadBeforeF32_);
+        }
+
+        CastF322B16(this->outQGm_ + block_idx * this->nlCoreRun_ * this->tilingData_->hiddenSizeQ +
+                      zz * this->tilingData_->hiddenSizeQ + perSlice * sliceSizeTmp_,
+                    commonUbuf_ + this->padBefore_, qkfp32Ubuf_ + PadBeforeF32_, repeatTimeOnce, dynamicSliceQTemp);
+        pipe_barrier((PIPE_ALL));
+      }
+
+      for (uint32_t perSlice = 0; perSlice < sliceTimeK_; ++perSlice) {  // 核内每块
+        uint32_t dynamicSliceKTemp = (perSlice == sliceTimeK_ - 1) ? lastSliceSizeK_ : sliceSizeTmp_;
+        headNumTempK = dynamicSliceKTemp / this->tilingData_->headDim;
+        uint32_t repeatTimeOnce = (dynamicSliceKTemp + this->repeatSize_ - 1) / this->repeatSize_;
+        pipe_barrier(PIPE_MTE2);
+        this->QkComm(this->kGm_ + block_idx * this->nlCoreRun_ * this->tilingData_->hiddenSizeK +
+                       zz * this->tilingData_->hiddenSizeK + perSlice * sliceSizeTmp_,
+                     extraGm, dynamicSliceKTemp, commonUbuf_, headNumTempK);
+
+        if (this->alignRotary_ == 0) {
+          pipe_barrier((PIPE_V));
+          CastB162F32(repeatTimeOnce);
+
+          pipe_barrier((PIPE_V));
+          this->CalcRopeAlign(qkfp32Ubuf_, repeatTimeOnce, oriPosF32_, removeBeforeF32_, PadBeforeF32_);
+        } else {
+          set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+          wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+          CastB162F32(repeatTimeOnce);
+
+          pipe_barrier((PIPE_V));
+          this->CalcRope(qkfp32Ubuf_, repeatTimeOnce, oriPosF32_, removeBeforeF32_, PadBeforeF32_, posOneF32_,
+                         PadBeforeF32_);
+        }
+
+        CastF322B16(this->outKGm_ + block_idx * this->nlCoreRun_ * this->tilingData_->hiddenSizeK +
+                      zz * this->tilingData_->hiddenSizeK + perSlice * sliceSizeTmp_,
+                    commonUbuf_ + this->padBefore_, qkfp32Ubuf_ + PadBeforeF32_, repeatTimeOnce, dynamicSliceKTemp);
+        pipe_barrier((PIPE_ALL));
+      }
+    }
+  }
+
+ private:
+  AscendC::TQue<AscendC::QuePosition::VECIN, 1> outQueueCO2_;
+  AscendC::TQue<AscendC::QuePosition::VECIN, 1> qkfp32QueueCO2_;
+  __ubuf__ QK_DTYPE *commonUbuf_{nullptr};
+  __ubuf__ float *qkfp32Ubuf_{nullptr};
+  uint32_t oriPosF32_{0};        // fp32的buf中qk的位置
+  uint32_t PadBeforeF32_{0};     // fp32的buf中保存qk[-x : hiddensize - x]
+  uint32_t removeBeforeF32_{0};  // fp32的buf中保存qk[x : hiddensize + x]
+  uint32_t posOneF32_{0};        // fp32的buf中0 0 0 1 1 1的位置
+  uint32_t headDimAlign_;        // 对齐的headDim
+  uint32_t sliceTimeQ_;          // 切分块的次数
+  uint32_t lastSliceSizeQ_;      // 最后一块的大小
+  uint32_t sliceTimeK_;
+  uint32_t lastSliceSizeK_;
+  uint32_t sliceSizeTmp_;
+};
+
+#endif