EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1009
csrc/attention/inplace_partial_rotary_mul/op_kernel/apply_rotary_pos_emb_common.h Normal file
View File

File diff suppressed because it is too large Load Diff

322
csrc/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul.cpp Normal file
View File

@@ -0,0 +1,322 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotary_position_embedding.cpp
* \brief
*/
#if defined(__DAV_C310__)
#include "kernel_operator.h"
#include "rotary_position_embedding_reg_bab.h"
#include "rotary_position_embedding_reg_ab.h"
#include "rotary_position_embedding_reg_aba_and_ba.h"
#include "rotary_position_embedding_reg_a_and_b.h"
#include "rotary_position_embedding_reg_bab_mixed.h"
#include "rotary_position_embedding_reg_aba_and_ba_mixed.h"
#include "rotary_position_embedding_reg_a_and_b_mixed.h"
#include "rotary_position_embedding_reg_ab_mixed.h"
#else
#include "kernel_operator.h"
#include "kernel_tiling/kernel_tiling.h"
#include "inplace_partial_rotary_mul.h"
#include "rotate_interleaved_split_s.h"
#include "rotate_interleaved_split_bs.h"
#include "rotate_interleaved_split_bsn.h"
#include "rotate_interleaved_split_s_pad.h"
#include "rotate_interleaved_split_bs_pad.h"
#include "rotate_interleaved_split_bsn_pad.h"
using namespace AscendC;
using namespace RotateInterleavedN;
#endif
#define TILING_KEY_ABA 20010
#define TILING_KEY_BA 20011
#define TILING_KEY_BAB 20020
#define TILING_KEY_AB 20030
#define TILING_KEY_A 20040
#define TILING_KEY_B 20041
#define TILING_KEY_ABA_BF16_FP32_MIXED 20110
#define TILING_KEY_ABA_FP16_FP32_MIXED 20210
#define TILING_KEY_BA_BF16_FP32_MIXED 20111
#define TILING_KEY_BA_FP16_FP32_MIXED 20211
#define TILING_KEY_BAB_BF16_FP32_MIXED 20120
#define TILING_KEY_BAB_FP16_FP32_MIXED 20220
#define TILING_KEY_AB_BF16_FP32_MIXED 20130
#define TILING_KEY_AB_FP16_FP32_MIXED 20230
#define TILING_KEY_A_BF16_FP32_MIXED 20140
#define TILING_KEY_A_FP16_FP32_MIXED 20240
#define TILING_KEY_B_BF16_FP32_MIXED 20141
#define TILING_KEY_B_FP16_FP32_MIXED 20241
#define TILING_KEY1 1
#define TILING_KEY2 2
#define TILING_KEY1_FP32_ROPE 11
#define TILING_KEY2_FP32_ROPE 12
using namespace AscendC;
using namespace InplacePartialRotaryMul;
extern "C" __global__ __aicore__ void inplace_partial_rotary_mul(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
GM_ADDR workspace, GM_ADDR tiling)
{
AscendC::TPipe pipe;
#if defined(__DAV_C310__)
if (TILING_KEY_IS(TILING_KEY_ABA))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBA<DTYPE_X, false> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_BA))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBA<DTYPE_X, true> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_BAB))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingBAB<DTYPE_X> op(&pipe, tilingData);
op.Init(x, cos, sin, y);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_AB))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingAB<DTYPE_X> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_A))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingAAndB<DTYPE_X, false> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_B))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingAAndB<DTYPE_X, true> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
// Mixed precision: x is half/bfloat16, cos/sin are float32
else if (TILING_KEY_IS(TILING_KEY_BAB_FP16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingBABMixed<half> op(&pipe, tilingData);
op.Init(x, cos, sin, y);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_BAB_BF16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingBABMixed<bfloat16_t> op(&pipe, tilingData);
op.Init(x, cos, sin, y);
op.Process();
}
// Mixed precision ABA/BA kernels
else if (TILING_KEY_IS(TILING_KEY_ABA_FP16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBAMixed<half, false> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_ABA_BF16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBAMixed<bfloat16_t, false> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_BA_FP16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBAMixed<half, true> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_BA_BF16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBAMixed<bfloat16_t, true> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
// Mixed precision AAndB kernels
else if (TILING_KEY_IS(TILING_KEY_A_FP16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingAAndBMixed<half, false> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_A_BF16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingAAndBMixed<bfloat16_t, false> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_B_FP16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingAAndBMixed<half, true> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_B_BF16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingAAndBMixed<bfloat16_t, true> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
// Mixed precision AB kernels
else if (TILING_KEY_IS(TILING_KEY_AB_FP16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingABMixed<half> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
else if (TILING_KEY_IS(TILING_KEY_AB_BF16_FP32_MIXED))
{
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling);
const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in;
InplacePartialRotaryMul::RotaryPositionEmbeddingABMixed<bfloat16_t> op;
op.Init(x, cos, sin, y, workspace, tilingData, &pipe);
op.Process();
}
#else
GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tilingData, tiling);
const RopeRegbaseTilingData* __restrict__ tilingData1 = &tilingData;
if (TILING_KEY_IS(TILING_KEY1)) {
InplacePartialRotaryMul::InplacePartialRotaryMulABA<DTYPE_X, true> op;
op.Init(x, cos, sin, y, workspace, tilingData1, &pipe);
op.Process();
return;
}
if (TILING_KEY_IS(TILING_KEY2)) {
InplacePartialRotaryMul::InplacePartialRotaryMulABA<DTYPE_X, false> op;
op.Init(x, cos, sin, y, workspace, tilingData1, &pipe);
op.Process();
return;
}
if (TILING_KEY_IS(TILING_KEY1_FP32_ROPE)) {
InplacePartialRotaryMul::InplacePartialRotaryMulABA<DTYPE_X, true, float> op;
op.Init(x, cos, sin, y, workspace, tilingData1, &pipe);
op.Process();
return;
}
if (TILING_KEY_IS(TILING_KEY2_FP32_ROPE)) {
InplacePartialRotaryMul::InplacePartialRotaryMulABA<DTYPE_X, false, float> op;
op.Init(x, cos, sin, y, workspace, tilingData1, &pipe);
op.Process();
return;
}
// mode: rotate_interleaved
if (TILING_KEY_IS(2000)) {
InterleavedSplitS<half> interleavedSplitS;
interleavedSplitS.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitS.Process();
} else if (TILING_KEY_IS(2010)) {
InterleavedSplitS<bfloat16_t> interleavedSplitS;
interleavedSplitS.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitS.Process();
} else if (TILING_KEY_IS(2020)) {
InterleavedSplitS<float> interleavedSplitS;
interleavedSplitS.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitS.Process();
} else if (TILING_KEY_IS(2100)) {
InterleavedSplitBS<half> interleavedSplitBS;
interleavedSplitBS.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBS.Process();
} else if (TILING_KEY_IS(2110)) {
InterleavedSplitBS<bfloat16_t> interleavedSplitBS;
interleavedSplitBS.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBS.Process();
} else if (TILING_KEY_IS(2120)) {
InterleavedSplitBS<float> interleavedSplitBS;
interleavedSplitBS.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBS.Process();
} else if (TILING_KEY_IS(2200)) {
InterleavedSplitBSN<half> interleavedSplitBSN;
interleavedSplitBSN.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBSN.Process();
} else if (TILING_KEY_IS(2210)) {
InterleavedSplitBSN<bfloat16_t> interleavedSplitBSN;
interleavedSplitBSN.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBSN.Process();
} else if (TILING_KEY_IS(2220)) {
InterleavedSplitBSN<float> interleavedSplitBSN;
interleavedSplitBSN.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBSN.Process();
} else if (TILING_KEY_IS(2001)) {
InterleavedSplitSPad<half> interleavedSplitSPad;
interleavedSplitSPad.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitSPad.Process();
} else if (TILING_KEY_IS(2011)) {
InterleavedSplitSPad<bfloat16_t> interleavedSplitSPad;
interleavedSplitSPad.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitSPad.Process();
} else if (TILING_KEY_IS(2021)) {
InterleavedSplitSPad<float> interleavedSplitSPad;
interleavedSplitSPad.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitSPad.Process();
} else if (TILING_KEY_IS(2101)) {
InterleavedSplitBSPad<half> interleavedSplitBSPad;
interleavedSplitBSPad.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBSPad.Process();
} else if (TILING_KEY_IS(2111)) {
InterleavedSplitBSPad<bfloat16_t> interleavedSplitBSPad;
interleavedSplitBSPad.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBSPad.Process();
} else if (TILING_KEY_IS(2121)) {
InterleavedSplitBSPad<float> interleavedSplitBSPad;
interleavedSplitBSPad.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBSPad.Process();
} else if (TILING_KEY_IS(2201)) {
InterleavedSplitBSNPad<half> interleavedSplitBSNPad;
interleavedSplitBSNPad.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBSNPad.Process();
} else if (TILING_KEY_IS(2211)) {
InterleavedSplitBSNPad<bfloat16_t> interleavedSplitBSNPad;
interleavedSplitBSNPad.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBSNPad.Process();
} else if (TILING_KEY_IS(2221)) {
InterleavedSplitBSNPad<float> interleavedSplitBSNPad;
interleavedSplitBSNPad.Init(x, cos, sin, y, tilingData1, &pipe);
interleavedSplitBSNPad.Process();
}
#endif
}

302
csrc/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul.h Normal file
View File

@@ -0,0 +1,302 @@
/**
* Copyright (c) Huawei Technologies Co., Ltd. 2025. 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.
*/
/*!
* \file inplace_partial_rotary_mul.h
* \brief
*/
#ifndef INPLACE_PARITAL_ROTARY_MUL_H
#define INPLACE_PARITAL_ROTARY_MUL_H
#include "kernel_operator.h"
namespace InplacePartialRotaryMul {
using namespace AscendC;
template <typename T, bool isBrc, typename R = T>
class InplacePartialRotaryMulABA {
public:
__aicore__ inline InplacePartialRotaryMulABA() {};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR r1, GM_ADDR r2, GM_ADDR y, GM_ADDR workspace,
const RopeRegbaseTilingData *tilingData, TPipe *pipe);
__aicore__ inline void Process();
__aicore__ inline void CopyInData(LocalTensor<T> &xUb, GlobalTensor<T> &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset);
__aicore__ inline void CopyInDataR(LocalTensor<R> &xUb, GlobalTensor<R> &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset);
__aicore__ inline void SetGatherSrcOffset(LocalTensor<int32_t> &idsUb, int64_t count);
__aicore__ inline void ComputeMul(LocalTensor<float> &dtsUb, LocalTensor<float> & src0Ub, LocalTensor<float> &src1Ub, int64_t onceA, int64_t numHead,int64_t headDim);
__aicore__ inline void InterleavedInversion(int64_t count,LocalTensor<float> &ub);
__aicore__ inline void DataCopyOut(LocalTensor<T> &yUb, GlobalTensor<T> &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset);
private:
TPipe* pipe_;
const RopeRegbaseTilingData* tiling_;
int32_t blockIdx_ = 0;
//需要的tilingdata
int64_t halfNumx_ = 0;
int64_t ropeUbOffset_ = 0;
int64_t xNum_= 0;
int64_t r1Num_= 0;
int64_t count_ = 0;
static constexpr int32_t ONE_BLOCK_SIZE = 32;
int32_t perBlock32 = ONE_BLOCK_SIZE / sizeof(float);
GlobalTensor<T> xGm_;
GlobalTensor<R> r1Gm_;
GlobalTensor<R> r2Gm_;
GlobalTensor<T> yGm_;
TQue<QuePosition::VECIN, 1> xQue_;
TQue<QuePosition::VECIN, 1> r1Que_;
TQue<QuePosition::VECIN, 1> r2Que_;
TQue<QuePosition::VECOUT, 1> yQue_;
TBuf<QuePosition::VECCALC> idsBuf_;
BinaryRepeatParams repeatParams_{1, 1, 1, 0, 0, 0};
DataCopyPadExtParams<T> dataCopyPadParams_{false, 0, 0, 0};
DataCopyPadExtParams<R> dataCopyRPadParams_{false, 0, 0, 0};
int64_t usedCoreNum_=0;
int64_t numHead_ =0;
int64_t headDim_ =0;
int64_t allHeadDim_ = 0;
int64_t coreTUbLoopTime_ =0;
int64_t coreBUbLoopTime_ = 0; //b分核，每个核处理多少个b，b就是shape0
int64_t coreTUbLoopTail_ =0;
int64_t coreBUbLoopTail_ = 0;
int64_t ubFactor_ = 0;
int64_t start_ = 0;
int64_t blockFactor_=0;
};
template <typename T, bool isBrc, typename R>
__aicore__ inline void InplacePartialRotaryMulABA<T, isBrc, R>::Init(GM_ADDR x, GM_ADDR r1, GM_ADDR r2, GM_ADDR y, GM_ADDR workspace,
const RopeRegbaseTilingData *tilingData, TPipe *pipe)
{
blockIdx_ = GetBlockIdx();
tiling_ = tilingData;
pipe_ = pipe;
usedCoreNum_ =tiling_->usedCoreNum;
numHead_ = tiling_->numHead;
headDim_ = tiling_->headDim;
allHeadDim_ = tiling_->allHeadDim;
coreTUbLoopTime_ = tiling_->coreTUbLoopTime;
coreBUbLoopTime_ = tiling_->coreBUbLoopTime; //b分核，每个核处理多少个b，b就是shape0
coreTUbLoopTail_ = tiling_->coreTUbLoopTail;
coreBUbLoopTail_ = tiling_->coreBUbLoopTail;
ubFactor_ = tiling_->ubFactor;
blockFactor_ = tiling_->blockFactor;
start_ = tiling_->start;
xGm_.SetGlobalBuffer((__gm__ T *)x);
r1Gm_.SetGlobalBuffer((__gm__ R *)r1);
r2Gm_.SetGlobalBuffer((__gm__ R *)r2);
yGm_.SetGlobalBuffer((__gm__ T *)y);
count_ = numHead_ * headDim_;
xNum_ = ubFactor_ * count_;
r1Num_ = ubFactor_ * headDim_;
pipe_->InitBuffer(xQue_, 2, xNum_ * sizeof(float));
pipe_->InitBuffer(r1Que_, 2, r1Num_ * sizeof(float));
pipe_->InitBuffer(r2Que_, 2, r1Num_ * sizeof(float));
pipe_->InitBuffer(yQue_, 2, xNum_ * sizeof(float));
pipe_->InitBuffer(idsBuf_, count_ *sizeof(uint32_t));
if constexpr(sizeof(T) != sizeof(float)) {
halfNumx_ = xNum_;
}
if constexpr(sizeof(R) != sizeof(float)) {
ropeUbOffset_ = r1Num_;
}
}
template <typename T, bool isBrc, typename R>
__aicore__ inline void InplacePartialRotaryMulABA<T, isBrc, R>::CopyInData(LocalTensor<T> &xUb, GlobalTensor<T> &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset)
{
DataCopyExtParams copyParams;
copyParams.blockCount = blockCout;
copyParams.blockLen = blockLen * sizeof(T);
copyParams.srcStride = (allHeadDim_ - headDim_)*sizeof(T); //整个输入的大小
copyParams.dstStride = 0;
DataCopyPad(xUb[ubOffset], xGm[gmOffset], copyParams, dataCopyPadParams_);
}
template <typename T, bool isBrc, typename R>
__aicore__ inline void InplacePartialRotaryMulABA<T, isBrc, R>::CopyInDataR(LocalTensor<R> &xUb, GlobalTensor<R> &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset)
{
DataCopyExtParams copyParams;
copyParams.blockCount = blockCout;
copyParams.blockLen = blockLen * sizeof(R);
copyParams.srcStride = 0; //整个输入的大小
copyParams.dstStride = 0;
DataCopyPad(xUb[ubOffset], xGm[gmOffset], copyParams, dataCopyRPadParams_);
}
template <typename T, bool isBrc, typename R>
__aicore__ inline void InplacePartialRotaryMulABA<T, isBrc, R>::SetGatherSrcOffset(LocalTensor<int32_t> &idsUb, int64_t count)
{
for (int32_t i = 0; i < 8; ++i) {
idsUb.SetValue(i, i ^ 1); // XOR with 1 to swap even and odd indices
}
event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
SetFlag<HardEvent::S_V>(eventIdSToV);
WaitFlag<HardEvent::S_V>(eventIdSToV);
int32_t scalarValue = 8;
int32_t onceNum = 8;
while (scalarValue < count) {
int32_t nextValue = scalarValue * 2;
if (nextValue < count) {
Adds(idsUb[scalarValue], idsUb, scalarValue, scalarValue);
} else {
Adds(idsUb[scalarValue], idsUb, scalarValue, count - scalarValue);
break;
}
scalarValue = nextValue;
}
Muls(idsUb, idsUb, 4, count);
}
template <typename T, bool isBrc, typename R>
__aicore__ inline void InplacePartialRotaryMulABA<T, isBrc, R>::ComputeMul(LocalTensor<float> &dtsUb, LocalTensor<float> & src0Ub, LocalTensor<float> &src1Ub, int64_t onceA, int64_t numHead,int64_t headDim)
{
int64_t count = numHead *headDim;
if constexpr(!isBrc) {
int64_t xtotalNum = onceA * count;
Mul(dtsUb, src0Ub, src1Ub, xtotalNum); //x*cos 非brc elewise乘
} else {
if (headDim <= 64) {
int32_t mask = headDim;
repeatParams_.dstBlkStride = 1;
repeatParams_.src0BlkStride = 1;
repeatParams_.src1BlkStride = 1;
repeatParams_.dstRepStride = 8;
repeatParams_.src0RepStride = 8;
repeatParams_.src1RepStride = 0;
for (int64_t j =0; j<onceA; j++) {
Mul(dtsUb[j*count], src0Ub[j*count], src1Ub[j*headDim], mask, numHead, repeatParams_); //x*cos 非brc elewise乘
}
} else {
for (int64_t j =0; j < onceA; j++) {
for (int64_t jj =0; jj < numHead; jj++) {
Mul(dtsUb[j*count + jj * headDim], src0Ub[j*count + jj * headDim], src1Ub[j*headDim], headDim); //x*cos 非brc elewise乘
}
}
}
}
PipeBarrier<PIPE_V>();
}
template <typename T, bool isBrc, typename R>
__aicore__ inline void InplacePartialRotaryMulABA<T, isBrc, R>::InterleavedInversion(int64_t count,LocalTensor<float> &ub)
{
// 做奇数位的*-1
SetMaskNorm();
int64_t fp32Mask = 64;
int64_t repeatTimes = count / 64;
int64_t remain = count % 64;
uint64_t fullMask = 0x5555555555555555; //0101010101010101
uint64_t tailMask = 0x55;
SetVectorMask<float, MaskMode::NORMAL>(0, fullMask);
Muls<float, false>(ub, ub, -1.0f, MASK_PLACEHOLDER, repeatTimes, {1,1,8,8});
if (remain != 0) {
int32_t tailTimes = count % 64 / 8;
SetVectorMask<float, MaskMode::NORMAL>(0, tailMask);
Muls<float, false>(ub[repeatTimes * 64], ub[repeatTimes * 64], -1.0f, MASK_PLACEHOLDER, tailTimes, {1,1,1,1});
}
ResetMask();
}
template <typename T, bool isBrc, typename R>
__aicore__ inline void InplacePartialRotaryMulABA<T, isBrc, R>::DataCopyOut(LocalTensor<T> &yUb, GlobalTensor<T> &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset)
{
DataCopyExtParams copyParams;
copyParams.blockCount = blockCout;
copyParams.blockLen = blockLen * sizeof(T);
copyParams.srcStride = 0;
copyParams.dstStride = (allHeadDim_ - headDim_)*sizeof(T);
DataCopyPad(yGm_[gmOffset], yUb[ubOffset], copyParams);
}
template <typename T, bool isBrc, typename R>
__aicore__ inline void InplacePartialRotaryMulABA<T, isBrc, R>::Process()
{
if (blockIdx_ >= usedCoreNum_) {
return;
}
LocalTensor<int32_t> idsUb = idsBuf_.Get<int32_t>();
int32_t count = numHead_ * headDim_;
SetGatherSrcOffset(idsUb, count);
LocalTensor<uint32_t> idsUbUint32 = idsBuf_.Get<uint32_t>();
// A分核分ub, 一次搬入x个 cout, x就是ubFactor_
int64_t ubLoopTimes = blockIdx_ == tiling_->usedCoreNum - 1 ? coreTUbLoopTime_ :coreBUbLoopTime_; //b分核，每个核处理多少个b，b就是shape0
int64_t ubLoopTailNum = blockIdx_ == tiling_->usedCoreNum - 1 ? coreTUbLoopTail_ :coreBUbLoopTail_;
int64_t ysCount = numHead_ *allHeadDim_;
for(int64_t i=0; i< ubLoopTimes; i++) {
int64_t ubSize = i == ubLoopTimes -1 ? ubLoopTailNum : ubFactor_;
int64_t xtotalNum = ubSize * count;
int64_t r1totalNum = ubSize * headDim_;
LocalTensor<T> xUb = xQue_.AllocTensor<T>();
int64_t gmOffset = blockIdx_ * blockFactor_ * ysCount + i * ubFactor_ * ysCount + start_;
int64_t blockCout = ubSize * numHead_;
CopyInData(xUb, xGm_, blockCout , headDim_, gmOffset, halfNumx_);
xQue_.EnQue<T>(xUb);
LocalTensor<R> r1Ub = r1Que_.AllocTensor<R>();
int64_t r1Offset = blockIdx_ * blockFactor_ * headDim_ + i * ubFactor_ * headDim_;
CopyInDataR(r1Ub, r1Gm_, ubSize , headDim_, r1Offset, ropeUbOffset_);
r1Que_.EnQue<R>(r1Ub);
LocalTensor<R> r2Ub = r2Que_.AllocTensor<R>();
CopyInDataR(r2Ub, r2Gm_, ubSize , headDim_, r1Offset, ropeUbOffset_);
r2Que_.EnQue<R>(r2Ub);
xUb = xQue_.DeQue<T>();
r1Ub = r1Que_.DeQue<R>();
r2Ub = r2Que_.DeQue<R>();
LocalTensor<float> xUbFp32 = xUb.template ReinterpretCast<float>();
LocalTensor<float> r1UbFp32 = r1Ub.template ReinterpretCast<float>();
LocalTensor<float> r2UbFp32 = r2Ub.template ReinterpretCast<float>();
if constexpr(sizeof(T) != sizeof(float)) {
// 非fp32时需要做cast
Cast(xUbFp32, xUb[halfNumx_], RoundMode::CAST_NONE, xtotalNum);
}
if constexpr(sizeof(R) != sizeof(float)) {
Cast(r1UbFp32, r1Ub[ropeUbOffset_], RoundMode::CAST_NONE, r1totalNum);
Cast(r2UbFp32, r2Ub[ropeUbOffset_], RoundMode::CAST_NONE, r1totalNum);
}
PipeBarrier<PIPE_V>();
LocalTensor<T> yUb = yQue_.AllocTensor<T>();
LocalTensor<float> yUbFp32 = yUb.template ReinterpretCast<float>();
ComputeMul(yUbFp32 , xUbFp32, r1UbFp32, ubSize, numHead_,headDim_);
// 开始做选择的取数，基偶取数
for (int64_t k =0; k < ubSize; k++) {
Gather(xUbFp32[k*count], xUbFp32[k*count], idsUbUint32, uint32_t(0), uint32_t(count)); //奇偶交换完成
}
PipeBarrier<PIPE_V>();
ComputeMul(xUbFp32 , xUbFp32, r2UbFp32, ubSize, numHead_,headDim_);
r1Que_.FreeTensor(r1Ub);
r2Que_.FreeTensor(r2Ub);
InterleavedInversion(xtotalNum, xUbFp32);
// 做最后的Add
Add(yUbFp32, yUbFp32, xUbFp32, xtotalNum);
xQue_.FreeTensor(xUb);
PipeBarrier<PIPE_V>();
if constexpr(sizeof(T) != sizeof(float)) {
Cast(yUb, yUbFp32, RoundMode::CAST_RINT, xtotalNum);
}
yQue_.EnQue<T>(yUb);
yUb = yQue_.DeQue<T>();
DataCopyOut(yUb, yGm_, blockCout, headDim_, gmOffset, 0);
yQue_.FreeTensor(yUb);
}
}
}
#endif

214
csrc/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul_common.h Normal file
View File

@@ -0,0 +1,214 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file common.h
* \brief
*/
#ifndef MOE_GATING_TOP_K_COMMON_H_V35
#define MOE_GATING_TOP_K_COMMON_H_V35
#include "kernel_operator.h"
namespace ops {
using namespace AscendC;
constexpr int32_t MIN_FP32 = 0xFF800000;
constexpr int64_t ONE_REPEAT_SORT_NUM = 32;
constexpr int64_t BLOCK_BYTES = 32;
constexpr int64_t MERGE_LIST_TWO = 2;
constexpr int64_t MERGE_LIST_THREE = 3;
constexpr int64_t MERGE_LIST_FOUR = 4;
constexpr int64_t MERGE_LIST_IDX_TWO = 2;
constexpr int64_t MERGE_LIST_IDX_THREE = 3;
__aicore__ inline int64_t Ceil(int64_t a, int64_t b)
{
if (b == 0) {
return 0;
}
return (a + b - 1) / b;
}
__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes)
{
if (bytes == 0) {
return 0;
}
return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes;
}
__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes)
{
return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES;
}
template <typename T>
__aicore__ inline T Min(T a, T b)
{
return a > b ? b : a;
}
template <typename T>
__aicore__ inline T Max(T a, T b)
{
return a < b ? b : a;
}
template <typename T1, typename T2>
__aicore__ inline T1 CeilAlign(T1 a, T2 b) {
if (b == 0) {
return 0;
}
return (a + b - 1) / b * b;
}
template <typename T1, typename T2>
__aicore__ inline T1 CeilDiv(T1 x, T2 y)
{
if (y != 0 && x != 0) {
const T1 quotient = x / y;
return (x % y != 0 && ((x ^ y) >= 0)) ? (quotient + 1) : quotient;
}
return x;
}
template <HardEvent event>
__aicore__ inline void SetWaitFlag(HardEvent evt)
{
event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(evt));
SetFlag<event>(eventId);
WaitFlag<event>(eventId);
}
constexpr AscendC::MicroAPI::CastTrait castTraitB162B32 = {
AscendC::MicroAPI::RegLayout::ZERO,
AscendC::MicroAPI::SatMode::UNKNOWN,
AscendC::MicroAPI::MaskMergeMode::ZEROING,
AscendC::RoundMode::UNKNOWN,
};
constexpr AscendC::MicroAPI::CastTrait castTraitB322B16 = {
AscendC::MicroAPI::RegLayout::ZERO,
AscendC::MicroAPI::SatMode::NO_SAT,
AscendC::MicroAPI::MaskMergeMode::ZEROING,
AscendC::RoundMode::CAST_RINT,
};
constexpr AscendC::MicroAPI::CastTrait castTraitB322Int32 = {
AscendC::MicroAPI::RegLayout::UNKNOWN,
AscendC::MicroAPI::SatMode::NO_SAT,
AscendC::MicroAPI::MaskMergeMode::ZEROING,
AscendC::RoundMode::CAST_TRUNC,
};
constexpr AscendC::MicroAPI::CastTrait castTraitB322Int16 = {
AscendC::MicroAPI::RegLayout::ZERO,
AscendC::MicroAPI::SatMode::NO_SAT,
AscendC::MicroAPI::MaskMergeMode::ZEROING,
AscendC::RoundMode::CAST_TRUNC,
};
constexpr AscendC::MicroAPI::CastTrait castTraitB162Int8 = {
AscendC::MicroAPI::RegLayout::ZERO,
AscendC::MicroAPI::SatMode::NO_SAT,
AscendC::MicroAPI::MaskMergeMode::ZEROING,
AscendC::RoundMode::CAST_TRUNC,
};
// load 对齐的 bfloat16,float16,bfloat32类型的 input(ub中)数据到 float32类型的dst(寄存器)中
template <typename T>
__aicore__ inline void LoadOneTensorForDtypeT(__local_mem__ T *input, MicroAPI::RegTensor<float> &dst,
MicroAPI::MaskReg &preg, uint32_t offset)
{
if constexpr (IsSameType<T, half>::value) {
MicroAPI::RegTensor<half> xFp16;
DataCopy<half, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ half *)(input) + (offset)));
Cast<float, half, castTraitB162B32>(dst, xFp16, preg);
} else if constexpr (IsSameType<T, bfloat16_t>::value) {
MicroAPI::RegTensor<bfloat16_t> xBf16;
DataCopy<bfloat16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>(xBf16,
((__local_mem__ bfloat16_t *)(input) + (offset)));
Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg);
} else {
DataCopy(dst, ((__local_mem__ float *)(input) + (offset)));
}
}
// load 2个对齐的Tensor 到寄存器中
template <typename T>
__aicore__ inline void LoadTwoTensorForDtypeT(__local_mem__ T *src1, __local_mem__ T *src2,
MicroAPI::RegTensor<float> &dst1, MicroAPI::RegTensor<float> &dst2,
MicroAPI::MaskReg &dst1Preg, MicroAPI::MaskReg &dst2Preg,
uint32_t src1Offset, uint32_t src2Offset)
{
if constexpr (IsSameType<T, half>::value) {
MicroAPI::RegTensor<half> xFp16Q;
MicroAPI::RegTensor<half> xFp16R;
DataCopy<half, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__local_mem__ half *)(src1) + (src1Offset)));
DataCopy<half, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16R, ((__local_mem__ half *)(src2) + (src2Offset)));
Cast<float, half, castTraitB162B32>(dst1, xFp16Q, dst1Preg);
Cast<float, half, castTraitB162B32>(dst2, xFp16R, dst2Preg);
} else if constexpr (IsSameType<T, bfloat16_t>::value) {
MicroAPI::RegTensor<bfloat16_t> xFp16Q;
MicroAPI::RegTensor<bfloat16_t> xFp16R;
DataCopy<bfloat16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__local_mem__ bfloat16_t *)(src1) + (src1Offset)));
DataCopy<bfloat16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16R, ((__local_mem__ bfloat16_t *)(src2) + (src2Offset)));
Cast<float, bfloat16_t, castTraitB162B32>(dst1, xFp16Q, dst1Preg);
Cast<float, bfloat16_t, castTraitB162B32>(dst2, xFp16R, dst2Preg);
} else {
DataCopy(dst1, ((__local_mem__ float *)(src1) + (src1Offset)));
DataCopy(dst2, ((__local_mem__ float *)(src2) + (src2Offset)));
}
}
// store 对齐的float32类型的src(寄存器)数据到output(ub)中，output数据类型支持bfloat16,float16,bfloat32,int32_t,int16_t,int8_t,uint8_t
template <typename T>
__aicore__ inline void StoreOneTensorForDtypeT(__local_mem__ T *output, MicroAPI::RegTensor<float> &src,
MicroAPI::MaskReg &preg, uint32_t offset)
{
if constexpr (IsSameType<T, half>::value) {
MicroAPI::RegTensor<half> yFp16;
Cast<half, float, castTraitB322B16>(yFp16, src, preg);
DataCopy<half, MicroAPI::StoreDist::DIST_PACK_B32>(((__local_mem__ half *)output + offset), yFp16, preg);
} else if constexpr (IsSameType<T, bfloat16_t>::value) {
MicroAPI::RegTensor<bfloat16_t> xBf16;
Cast<bfloat16_t, float, castTraitB322B16>(xBf16, src, preg);
DataCopy<bfloat16_t, MicroAPI::StoreDist::DIST_PACK_B32>(((__local_mem__ bfloat16_t *)output + offset),
xBf16, preg);
} else if constexpr (IsSameType<T, int32_t>::value) {
MicroAPI::RegTensor<int32_t> zInt32;
MicroAPI::Cast<int32_t, float, castTraitB322Int32>(zInt32, src, preg);
DataCopy<int32_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(((__local_mem__ int32_t *)output + offset), zInt32, preg);
} else if constexpr (IsSameType<T, int16_t>::value) {
MicroAPI::RegTensor<int16_t> zInt16;
MicroAPI::Cast<int16_t, float, castTraitB322Int16>(zInt16, src, preg);
DataCopy<int16_t, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(((__local_mem__ int16_t *)output + offset), zInt16, preg);
} else if constexpr (IsSameType<T, int8_t>::value) {
MicroAPI::RegTensor<half> yFp16;
MicroAPI::RegTensor<int8_t> zInt8;
Cast<half, float, castTraitB322Int16>(yFp16, src, preg);
Cast<int8_t, half, castTraitB162Int8>(zInt8, yFp16, preg);
DataCopy<int8_t, MicroAPI::StoreDist::DIST_PACK4_B32>(((__local_mem__ int8_t *)output + offset), zInt8, preg);
} else if constexpr (IsSameType<T, uint8_t>::value) {
MicroAPI::RegTensor<half> yFp16;
MicroAPI::RegTensor<uint8_t> zUint8;
Cast<half, float, castTraitB322Int16>(yFp16, src, preg);
Cast<uint8_t, half, castTraitB162Int8>(zUint8, yFp16, preg);
DataCopy<uint8_t, MicroAPI::StoreDist::DIST_PACK4_B32>(((__local_mem__ uint8_t *)output + offset), zUint8, preg);
} else {
DataCopy(((__local_mem__ float *)output + offset), src, preg);
}
}
} // namespace ops
#endif // MOE_GATING_TOP_K_COMMON_H_V35

294
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_a_and_b.h Normal file
View File

@@ -0,0 +1,294 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotary_position_embedding_reg_a_and_b.h
* \brief
*/
#ifndef ROTARY_POSITION_EMBEDDING_REG_A_AND_B_H
#define ROTARY_POSITION_EMBEDDING_REG_A_AND_B_H
// #include "op_kernel/math_util.h"
#include "apply_rotary_pos_emb_common.h"
namespace InplacePartialRotaryMul {
using namespace AscendC;
template <typename T, bool IsBoardCast>
class RotaryPositionEmbeddingAAndB
{
public:
__aicore__ inline RotaryPositionEmbeddingAAndB(){};
__aicore__ inline ~RotaryPositionEmbeddingAAndB(){};
__aicore__ inline void Init(
GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData,
TPipe* pipe);
__aicore__ inline void Process();
private:
// Init过程中使用的内部函数
__aicore__ inline void InitAllGlobalBuffer(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut);
__aicore__ inline void InitAllBuffer();
__aicore__ inline void InitLoopParams();
// 各个层级的Process函数
__aicore__ inline void ProcessInLoop(LocalTensor<T>& cos, LocalTensor<T>& sin, int64_t bStart, int64_t bLength);
// 拷入拷出函数
__aicore__ inline void CopyInCosAndSin(int64_t bStart, int64_t bLength);
__aicore__ inline void CopyInQ(GlobalTensor<T>& source, int64_t bStart, int64_t bLength);
__aicore__ inline void CopyOutQ(GlobalTensor<T>& target, int64_t bStart, int64_t bLength);
// 计算函数
__aicore__ inline void Compute(LocalTensor<T>& cos, LocalTensor<T>& sin, int64_t bLength);
private:
constexpr static uint32_t COS_DB_BUFFER = IsBoardCast ? 1 : DOUBLE_BUFFER;
TPipe* pipe_;
// GlobalMemory
GlobalTensor<T> qGm_;
GlobalTensor<T> cosGm_;
GlobalTensor<T> sinGm_;
GlobalTensor<T> qOutGm_;
// UB
TQue<QuePosition::VECIN, DOUBLE_BUFFER> qInQueue_;
TQue<QuePosition::VECIN, COS_DB_BUFFER> cosInQueue_;
TQue<QuePosition::VECIN, COS_DB_BUFFER> sinInQueue_;
TQue<QuePosition::VECOUT, DOUBLE_BUFFER> qOutQueue_;
// Split core info
int64_t blockIdx_ = 0;
int64_t bBlockStart_ = 0;
int64_t bBlockLength_ = 0;
// TilingData
const RopeRegbaseTilingData* tilingData_;
int64_t ubFactorB_ = 0;
int64_t D_ = 0;
int64_t dAlign_ = 0;
// 拷贝参数
uint8_t dSplitCoef_ = 1;
uint8_t copyInQSplitCoef_ = 1; // 拷贝q时使用的splitCoef
uint64_t ubCopyInStride = 0; // 输入在ub中的stride，deepseek_interleave中不为0
};
template <typename T, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndB<T, IsBoardCast>::Init(
GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData,
TPipe* pipe)
{
this->tilingData_ = tilingData;
this->pipe_ = pipe;
this->blockIdx_ = GetBlockIdx();
this->InitAllGlobalBuffer(q, cos, sin, qOut);
this->InitAllBuffer();
this->InitLoopParams();
}
template <typename T, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndB<T, IsBoardCast>::InitAllGlobalBuffer(
GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut)
{
this->qGm_.SetGlobalBuffer((__gm__ T*)q);
this->cosGm_.SetGlobalBuffer((__gm__ T*)cos);
this->sinGm_.SetGlobalBuffer((__gm__ T*)sin);
this->qOutGm_.SetGlobalBuffer((__gm__ T*)qOut);
}
template <typename T, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndB<T, IsBoardCast>::InitAllBuffer()
{
this->ubFactorB_ = this->tilingData_->ubFactorB;
this->D_ = this->tilingData_->D;
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF) ||
tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
this->dSplitCoef_ = HALF_INTERLEAVE_COEF;
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
this->dSplitCoef_ = QUARTER_MODE_COEF;
}
this->copyInQSplitCoef_ = dSplitCoef_;
this->dAlign_ = ops::CeilAlign<int64_t>(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(T)) * dSplitCoef_;
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
this->copyInQSplitCoef_ = 1;
// 非boardcast时，使用批量计算API，需要拷贝时添加stride
if constexpr (!IsBoardCast) {
this->ubCopyInStride =
(this->dAlign_ * sizeof(T) - ops::CeilAlign<int64_t>(tilingData_->sliceLength * sizeof(T), BLOCK_TYPE_SIZE)) /
BLOCK_TYPE_SIZE;
}
}
this->pipe_->InitBuffer(this->qInQueue_, DOUBLE_BUFFER, ubFactorB_ * dAlign_ * sizeof(T));
this->pipe_->InitBuffer(this->qOutQueue_, DOUBLE_BUFFER, ubFactorB_ * dAlign_ * sizeof(T));
if constexpr (IsBoardCast) {
this->pipe_->InitBuffer(this->cosInQueue_, COS_DB_BUFFER, dAlign_ * sizeof(T));
this->pipe_->InitBuffer(this->sinInQueue_, COS_DB_BUFFER, dAlign_ * sizeof(T));
} else {
this->pipe_->InitBuffer(this->cosInQueue_, COS_DB_BUFFER, ubFactorB_ * dAlign_ * sizeof(T));
this->pipe_->InitBuffer(this->sinInQueue_, COS_DB_BUFFER, ubFactorB_ * dAlign_ * sizeof(T));
}
}
template <typename T, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndB<T, IsBoardCast>::InitLoopParams()
{
this->bBlockLength_ = tilingData_->blockFactorB;
if (blockIdx_ == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) {
this->bBlockLength_ = tilingData_->B % tilingData_->blockFactorB;
}
this->bBlockStart_ = blockIdx_ * tilingData_->blockFactorB;
}
template <typename T, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndB<T, IsBoardCast>::Process()
{
// 在B轴进行循环
int64_t ubLoopCount = ops::CeilDiv(bBlockLength_, ubFactorB_);
if constexpr (IsBoardCast) {
this->CopyInCosAndSin(0, 1);
LocalTensor<T> cosUb = this->cosInQueue_.template DeQue<T>();
LocalTensor<T> sinUb = this->sinInQueue_.template DeQue<T>();
for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) {
this->ProcessInLoop(
cosUb, sinUb, bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_);
}
this->cosInQueue_.FreeTensor(cosUb);
this->sinInQueue_.FreeTensor(sinUb);
} else {
for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) {
this->CopyInCosAndSin(
bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_);
LocalTensor<T> cosUb = this->cosInQueue_.template DeQue<T>();
LocalTensor<T> sinUb = this->sinInQueue_.template DeQue<T>();
this->ProcessInLoop(
cosUb, sinUb, bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_);
this->cosInQueue_.FreeTensor(cosUb);
this->sinInQueue_.FreeTensor(sinUb);
}
}
}
template <typename T, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndB<T, IsBoardCast>::ProcessInLoop(
LocalTensor<T>& cos, LocalTensor<T>& sin, int64_t bUbStart, int64_t bUbLength)
{
CopyInQ(qGm_, bUbStart, bUbLength);
Compute(cos, sin, bUbLength);
CopyOutQ(qOutGm_, bUbStart, bUbLength);
}
template <typename T, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndB<T, IsBoardCast>::CopyInCosAndSin(int64_t bStart, int64_t bLength)
{
LocalTensor<T> cosUb = this->cosInQueue_.template AllocTensor<T>();
LocalTensor<T> sinUb = this->sinInQueue_.template AllocTensor<T>();
DataCopyPadExtParams<T> copyPadExtparams;
copyPadExtparams.isPad = false;
copyPadExtparams.leftPadding = 0;
copyPadExtparams.rightPadding = 0;
copyPadExtparams.paddingValue = 0;
DataCopyExtParams copyExtParams;
copyExtParams.blockCount = bLength * dSplitCoef_;
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / dSplitCoef_;
copyExtParams.srcStride = 0;
copyExtParams.dstStride = 0;
DataCopyPad(cosUb, this->cosGm_[bStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams);
DataCopyPad(sinUb, this->sinGm_[bStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams);
this->cosInQueue_.template EnQue(cosUb);
this->sinInQueue_.template EnQue(sinUb);
}
template <typename T, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndB<T, IsBoardCast>::CopyInQ(
GlobalTensor<T>& source, int64_t bStart, int64_t bLength)
{
LocalTensor<T> target = this->qInQueue_.template AllocTensor<T>();
DataCopyExtParams copyExtParams;
copyExtParams.blockCount = bLength * copyInQSplitCoef_;
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / copyInQSplitCoef_;
copyExtParams.srcStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(T);
copyExtParams.dstStride = ubCopyInStride;
DataCopyPadExtParams<T> copyPadExtparams;
copyPadExtparams.isPad = false;
copyPadExtparams.leftPadding = 0;
copyPadExtparams.rightPadding = 0;
copyPadExtparams.paddingValue = 0;
DataCopyPad(target, source[bStart * D_ + tilingData_->sliceStart], copyExtParams, copyPadExtparams);
this->qInQueue_.template EnQue(target);
}
template <typename T, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndB<T, IsBoardCast>::CopyOutQ(
GlobalTensor<T>& target, int64_t bStart, int64_t bLength)
{
LocalTensor<T> source = this->qOutQueue_.template DeQue<T>();
DataCopyExtParams copyExtParams;
copyExtParams.blockCount = bLength * dSplitCoef_;
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / dSplitCoef_;
copyExtParams.srcStride = 0;
copyExtParams.dstStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(T);
DataCopyPad(target[bStart * D_+ tilingData_->sliceStart], source, copyExtParams);
this->qOutQueue_.FreeTensor(source);
}
template <typename T, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndB<T, IsBoardCast>::Compute(
LocalTensor<T>& cos, LocalTensor<T>& sin, int64_t bLength)
{
LocalTensor<T> inUb = this->qInQueue_.template DeQue<T>();
LocalTensor<T> outUb = this->qOutQueue_.template AllocTensor<T>();
if constexpr (IsBoardCast) {
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF)) {
HalfAlignVF<T>(sin, cos, inUb, outUb, tilingData_->sliceLength, dAlign_, 1, bLength);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::INTERLEAVE)) {
InterleaveModeVF<T>(sin, cos, inUb, outUb, tilingData_->sliceLength, 1, bLength);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
QuarterAlignVF<T>(sin, cos, inUb, outUb, tilingData_->sliceLength, dAlign_, 1, bLength);
} else {
DeepSeekInterleaveModeVF<T>(sin, cos, inUb, outUb, tilingData_->sliceLength, 1, bLength);
}
} else {
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF)) {
BatchHalfAlignVF<T, IsBoardCast>(
(__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(),
(__local_mem__ T*)sin.GetPhyAddr(), (__local_mem__ T*)outUb.GetPhyAddr(), bLength, 1, 1, tilingData_->sliceLength, dAlign_,
ubFactorB_, 1);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::INTERLEAVE)) {
BatchInterleaveModeVF<T, IsBoardCast>(
(__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(),
(__local_mem__ T*)sin.GetPhyAddr(), (__local_mem__ T*)outUb.GetPhyAddr(), bLength, 1, 1, tilingData_->sliceLength, dAlign_,
ubFactorB_, 1);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
BatchQuarterAlignVF<T, IsBoardCast>(
(__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(),
(__local_mem__ T*)sin.GetPhyAddr(), (__local_mem__ T*)outUb.GetPhyAddr(), bLength, 1, 1, tilingData_->sliceLength, dAlign_,
ubFactorB_, 1);
} else {
BatchDeepSeekInterleaveModeVF<T, IsBoardCast>(
(__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(),
(__local_mem__ T*)sin.GetPhyAddr(), (__local_mem__ T*)outUb.GetPhyAddr(), bLength, 1, 1, tilingData_->sliceLength, dAlign_,
ubFactorB_, 1);
}
}
this->qInQueue_.FreeTensor(inUb);
this->qOutQueue_.template EnQue(outUb);
}
} // namespace InplacePartialRotaryMul
#endif

273
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_a_and_b_mixed.h Normal file
View File

@@ -0,0 +1,273 @@
/**
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
* CANN Open Software License Agreement Version 2.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
*/
/*!
* \file rotary_position_embedding_reg_a_and_b_mixed.h
* \brief Mixed precision kernel for AAndB layout: x is half/bfloat16, cos/sin are float
*/
#ifndef ROTARY_POSITION_EMBEDDING_REG_A_AND_B_MIXED_H
#define ROTARY_POSITION_EMBEDDING_REG_A_AND_B_MIXED_H
#include "apply_rotary_pos_emb_common.h"
namespace InplacePartialRotaryMul {
using namespace AscendC;
template <typename TX, bool IsBoardCast>
class RotaryPositionEmbeddingAAndBMixed {
public:
__aicore__ inline RotaryPositionEmbeddingAAndBMixed(){};
__aicore__ inline ~RotaryPositionEmbeddingAAndBMixed(){};
__aicore__ inline void Init(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace,
const RopeRegbaseTilingData *tilingData, TPipe *pipe);
__aicore__ inline void Process();
private:
__aicore__ inline void InitAllGlobalBuffer(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut);
__aicore__ inline void InitAllBuffer();
__aicore__ inline void InitLoopParams();
__aicore__ inline void ProcessInLoop(
LocalTensor<float> &cos, LocalTensor<float> &sin, int64_t bStart, int64_t bLength);
__aicore__ inline void CopyInCosAndSin(int64_t bStart, int64_t bLength);
__aicore__ inline void CopyInQ(GlobalTensor<TX> &source, int64_t bStart, int64_t bLength);
__aicore__ inline void CopyOutQ(GlobalTensor<TX> &target, int64_t bStart, int64_t bLength);
__aicore__ inline void Compute(LocalTensor<float> &cos, LocalTensor<float> &sin, int64_t bLength);
private:
static constexpr uint32_t COS_DB_BUFFER = IsBoardCast ? 1 : DOUBLE_BUFFER;
TPipe *pipe_;
GlobalTensor<TX> qGm_;
GlobalTensor<float> cosGm_;
GlobalTensor<float> sinGm_;
GlobalTensor<TX> qOutGm_;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> qInQueue_;
TQue<QuePosition::VECIN, COS_DB_BUFFER> cosInQueue_;
TQue<QuePosition::VECIN, COS_DB_BUFFER> sinInQueue_;
TQue<QuePosition::VECOUT, DOUBLE_BUFFER> qOutQueue_;
int64_t blockIdx_ = 0;
int64_t bBlockStart_ = 0;
int64_t bBlockLength_ = 0;
const RopeRegbaseTilingData *tilingData_;
int64_t ubFactorB_ = 0;
int64_t D_ = 0;
int64_t dAlign_ = 0;
int64_t dAlignFloat_ = 0;
uint8_t dSplitCoef_ = 1;
uint8_t copyInQSplitCoef_ = 1;
uint64_t ubCopyInStride = 0;
};
template <typename TX, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndBMixed<TX, IsBoardCast>::Init(GM_ADDR q, GM_ADDR cos, GM_ADDR sin,
GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData *tilingData, TPipe *pipe)
{
this->tilingData_ = tilingData;
this->pipe_ = pipe;
this->blockIdx_ = GetBlockIdx();
this->InitAllGlobalBuffer(q, cos, sin, qOut);
this->InitAllBuffer();
this->InitLoopParams();
}
template <typename TX, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndBMixed<TX, IsBoardCast>::InitAllGlobalBuffer(
GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut)
{
this->qGm_.SetGlobalBuffer((__gm__ TX *)q);
this->cosGm_.SetGlobalBuffer((__gm__ float *)cos);
this->sinGm_.SetGlobalBuffer((__gm__ float *)sin);
this->qOutGm_.SetGlobalBuffer((__gm__ TX *)qOut);
}
template <typename TX, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndBMixed<TX, IsBoardCast>::InitAllBuffer()
{
this->ubFactorB_ = this->tilingData_->ubFactorB;
this->D_ = this->tilingData_->D;
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF) ||
tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
this->dSplitCoef_ = HALF_INTERLEAVE_COEF;
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
this->dSplitCoef_ = QUARTER_MODE_COEF;
}
this->copyInQSplitCoef_ = dSplitCoef_;
this->dAlign_ =
ops::CeilAlign<int64_t>(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(TX)) * dSplitCoef_;
this->dAlignFloat_ =
ops::CeilAlign<int64_t>(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(float)) * dSplitCoef_;
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
this->copyInQSplitCoef_ = 1;
if constexpr (!IsBoardCast) {
this->ubCopyInStride =
(this->dAlign_ * sizeof(TX) -
ops::CeilAlign<int64_t>(tilingData_->sliceLength * sizeof(TX), BLOCK_TYPE_SIZE)) /
BLOCK_TYPE_SIZE;
}
}
this->pipe_->InitBuffer(this->qInQueue_, DOUBLE_BUFFER, ubFactorB_ * dAlign_ * sizeof(TX));
this->pipe_->InitBuffer(this->qOutQueue_, DOUBLE_BUFFER, ubFactorB_ * dAlign_ * sizeof(TX));
if constexpr (IsBoardCast) {
this->pipe_->InitBuffer(this->cosInQueue_, COS_DB_BUFFER, dAlignFloat_ * sizeof(float));
this->pipe_->InitBuffer(this->sinInQueue_, COS_DB_BUFFER, dAlignFloat_ * sizeof(float));
} else {
this->pipe_->InitBuffer(this->cosInQueue_, COS_DB_BUFFER, ubFactorB_ * dAlignFloat_ * sizeof(float));
this->pipe_->InitBuffer(this->sinInQueue_, COS_DB_BUFFER, ubFactorB_ * dAlignFloat_ * sizeof(float));
}
}
template <typename TX, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndBMixed<TX, IsBoardCast>::InitLoopParams()
{
this->bBlockLength_ = tilingData_->blockFactorB;
if (blockIdx_ == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) {
this->bBlockLength_ = tilingData_->B % tilingData_->blockFactorB;
}
this->bBlockStart_ = blockIdx_ * tilingData_->blockFactorB;
}
template <typename TX, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndBMixed<TX, IsBoardCast>::Process()
{
int64_t ubLoopCount = ops::CeilDiv(bBlockLength_, ubFactorB_);
if constexpr (IsBoardCast) {
this->CopyInCosAndSin(0, 1);
LocalTensor<float> cosUb = this->cosInQueue_.template DeQue<float>();
LocalTensor<float> sinUb = this->sinInQueue_.template DeQue<float>();
for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) {
this->ProcessInLoop(cosUb,
sinUb,
bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_);
}
this->cosInQueue_.FreeTensor(cosUb);
this->sinInQueue_.FreeTensor(sinUb);
} else {
for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) {
this->CopyInCosAndSin(bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_);
LocalTensor<float> cosUb = this->cosInQueue_.template DeQue<float>();
LocalTensor<float> sinUb = this->sinInQueue_.template DeQue<float>();
this->ProcessInLoop(cosUb,
sinUb,
bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_);
this->cosInQueue_.FreeTensor(cosUb);
this->sinInQueue_.FreeTensor(sinUb);
}
}
}
template <typename TX, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndBMixed<TX, IsBoardCast>::ProcessInLoop(
LocalTensor<float> &cos, LocalTensor<float> &sin, int64_t bUbStart, int64_t bUbLength)
{
CopyInQ(qGm_, bUbStart, bUbLength);
Compute(cos, sin, bUbLength);
CopyOutQ(qOutGm_, bUbStart, bUbLength);
}
template <typename TX, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndBMixed<TX, IsBoardCast>::CopyInCosAndSin(
int64_t bStart, int64_t bLength)
{
LocalTensor<float> cosUb = this->cosInQueue_.template AllocTensor<float>();
LocalTensor<float> sinUb = this->sinInQueue_.template AllocTensor<float>();
DataCopyPadExtParams<float> copyPadExtparams;
copyPadExtparams.isPad = false;
copyPadExtparams.leftPadding = 0;
copyPadExtparams.rightPadding = 0;
copyPadExtparams.paddingValue = 0;
DataCopyExtParams copyExtParams;
if constexpr (IsBoardCast) {
copyExtParams.blockCount = 1 * dSplitCoef_;
} else {
copyExtParams.blockCount = bLength * dSplitCoef_;
}
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(float) / dSplitCoef_;
copyExtParams.srcStride = 0;
copyExtParams.dstStride = 0;
DataCopyPad(cosUb, this->cosGm_[bStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams);
DataCopyPad(sinUb, this->sinGm_[bStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams);
this->cosInQueue_.template EnQue(cosUb);
this->sinInQueue_.template EnQue(sinUb);
}
template <typename TX, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndBMixed<TX, IsBoardCast>::CopyInQ(
GlobalTensor<TX> &source, int64_t bStart, int64_t bLength)
{
LocalTensor<TX> target = this->qInQueue_.template AllocTensor<TX>();
DataCopyExtParams copyExtParams;
copyExtParams.blockCount = bLength * copyInQSplitCoef_;
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(TX) / copyInQSplitCoef_;
copyExtParams.srcStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(TX);
copyExtParams.dstStride = ubCopyInStride;
DataCopyPadExtParams<TX> copyPadExtparams;
copyPadExtparams.isPad = false;
copyPadExtparams.leftPadding = 0;
copyPadExtparams.rightPadding = 0;
copyPadExtparams.paddingValue = 0;
DataCopyPad(target, source[bStart * D_ + tilingData_->sliceStart], copyExtParams, copyPadExtparams);
this->qInQueue_.template EnQue(target);
}
template <typename TX, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndBMixed<TX, IsBoardCast>::CopyOutQ(
GlobalTensor<TX> &target, int64_t bStart, int64_t bLength)
{
LocalTensor<TX> source = this->qOutQueue_.template DeQue<TX>();
DataCopyExtParams copyExtParams;
copyExtParams.blockCount = bLength * dSplitCoef_;
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(TX) / dSplitCoef_;
copyExtParams.srcStride = 0;
copyExtParams.dstStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(TX);
DataCopyPad(target[bStart * D_ + tilingData_->sliceStart], source, copyExtParams);
this->qOutQueue_.FreeTensor(source);
}
template <typename TX, bool IsBoardCast>
__aicore__ inline void RotaryPositionEmbeddingAAndBMixed<TX, IsBoardCast>::Compute(
LocalTensor<float> &cos, LocalTensor<float> &sin, int64_t bLength)
{
LocalTensor<TX> inUb = this->qInQueue_.template DeQue<TX>();
LocalTensor<TX> outUb = this->qOutQueue_.template AllocTensor<TX>();
if constexpr (IsBoardCast) {
InterleaveModeVFMixed<TX>(inUb, cos, sin, outUb, tilingData_->sliceLength, 1, bLength);
} else {
BatchInterleaveModeVFMixed<TX, IsBoardCast>((__local_mem__ TX *)inUb.GetPhyAddr(),
(__local_mem__ float *)cos.GetPhyAddr(),
(__local_mem__ float *)sin.GetPhyAddr(),
(__local_mem__ TX *)outUb.GetPhyAddr(),
bLength,
1,
1,
tilingData_->sliceLength,
dAlign_,
dAlignFloat_,
ubFactorB_,
1);
}
this->qInQueue_.FreeTensor(inUb);
this->qOutQueue_.template EnQue(outUb);
}
} // namespace InplacePartialRotaryMul
#endif

164
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_ab.h Normal file
View File

@@ -0,0 +1,164 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotary_position_embedding_reg_ab.h
* \brief
*/
#ifndef ROTARY_POSITION_EMBEDDING_REG_AB_H
#define ROTARY_POSITION_EMBEDDING_REG_AB_H
// #include "op_kernel/math_util.h"
// #include "op_kernel/platform_util.h"
// #include "op_kernel/load_store_utils.h"
#include "apply_rotary_pos_emb_common.h"
namespace InplacePartialRotaryMul {
using namespace AscendC;
template <typename T>
class RotaryPositionEmbeddingAB
{
public:
__aicore__ inline RotaryPositionEmbeddingAB(){};
__aicore__ inline void Init(
GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData,
TPipe* pipe);
__aicore__ inline void Process();
private:
__aicore__ inline void ProcessLoop(
int64_t xGmOffset, LocalTensor<T> cosBuffer, LocalTensor<T> sinBuffer, int64_t ubIdx, int64_t bsCount,
int64_t nCount);
private:
TPipe* pipe_;
TQue<QuePosition::VECIN, 1> xInQueue_;
TQue<QuePosition::VECIN, 1> cosInQueue_;
TQue<QuePosition::VECIN, 1> sinInQueue_;
TQue<QuePosition::VECOUT, 1> yOutQueue_;
GlobalTensor<T> xGm_;
GlobalTensor<T> cosGm_;
GlobalTensor<T> sinGm_;
GlobalTensor<T> yGm_;
const RopeRegbaseTilingData* tilingData_;
DataCopyPadExtParams<T> padParams_ = {false, 0, 0, static_cast<T>(0)};
uint8_t DB_FLAG = 2;
uint32_t dSplitSize_ = 0;
int64_t bsBlockCount_ = 0;
int64_t nBlockCount_ = 0;
int64_t sliceAlign_ = 0;
};
template <typename T>
__aicore__ inline void RotaryPositionEmbeddingAB<T>::Init(
GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData,
TPipe* pipe)
{
pipe_ = pipe;
tilingData_ = tilingData;
dSplitSize_ = tilingData_->sliceLength / tilingData_->dSplitCoef * sizeof(T);
int64_t blockDimBS = GetBlockIdx() / tilingData_->blockNumN;
int64_t blockDimN = GetBlockIdx() % tilingData_->blockNumN;
bsBlockCount_ = (blockDimBS == tilingData_->blockNumBS - 1) ? tilingData_->blockTailBS : tilingData_->blockFactorBS;
nBlockCount_ = (blockDimN == tilingData_->blockNumN - 1) ? tilingData_->blockTailN : tilingData_->blockFactorN;
int64_t cosOffset = blockDimBS * tilingData_->blockFactorBS * tilingData_->sliceLength;
int64_t offset = blockDimBS * tilingData_->blockFactorBS * tilingData_->D;
int64_t xOffset = offset * tilingData_->N + blockDimN * tilingData_->blockFactorN * tilingData_->D + tilingData_->sliceStart;
this->cosGm_.SetGlobalBuffer((__gm__ T*)cos + cosOffset);
this->sinGm_.SetGlobalBuffer((__gm__ T*)sin + cosOffset);
this->xGm_.SetGlobalBuffer((__gm__ T*)x + xOffset);
this->yGm_.SetGlobalBuffer((__gm__ T*)y + xOffset);
sliceAlign_ = ops::CeilDiv(tilingData_->sliceLength * sizeof(T), GetUbBlockSize()) * GetUbBlockSize() / sizeof(T);
int64_t bufferSize = sliceAlign_ * sizeof(T) * tilingData_->ubFactorBS;
pipe_->InitBuffer(xInQueue_, DB_FLAG, bufferSize * tilingData_->ubFactorN);
pipe_->InitBuffer(cosInQueue_, DB_FLAG, bufferSize);
pipe_->InitBuffer(sinInQueue_, DB_FLAG, bufferSize);
pipe_->InitBuffer(yOutQueue_, DB_FLAG, bufferSize * tilingData_->ubFactorN);
}
template <typename T>
__aicore__ inline void RotaryPositionEmbeddingAB<T>::Process()
{
uint32_t bsLoopCnt = ops::CeilDiv(bsBlockCount_, tilingData_->ubFactorBS);
uint32_t nLoopCnt = ops::CeilDiv(nBlockCount_, tilingData_->ubFactorN);
for (uint32_t bsLoopIdx = 0; bsLoopIdx < bsLoopCnt; bsLoopIdx++) {
int64_t xGmOffset = bsLoopIdx * tilingData_->ubFactorBS * tilingData_->N * tilingData_->D;
uint32_t currBSNum = (bsLoopIdx != bsLoopCnt - 1) ? tilingData_->ubFactorBS :
bsBlockCount_ - (bsLoopIdx * tilingData_->ubFactorBS);
DataCopyExtParams cosParams = {
static_cast<uint16_t>(currBSNum * tilingData_->dSplitCoef), dSplitSize_, 0, 0, 0};
LocalTensor<T> cosBuffer = cosInQueue_.AllocTensor<T>();
LocalTensor<T> sinBuffer = sinInQueue_.AllocTensor<T>();
DataCopyPad(cosBuffer, cosGm_[bsLoopIdx * tilingData_->ubFactorBS * tilingData_->sliceLength], cosParams, padParams_);
cosInQueue_.EnQue(cosBuffer);
cosBuffer = cosInQueue_.DeQue<T>();
DataCopyPad(sinBuffer, sinGm_[bsLoopIdx * tilingData_->ubFactorBS * tilingData_->sliceLength], cosParams, padParams_);
sinInQueue_.EnQue(sinBuffer);
sinBuffer = sinInQueue_.DeQue<T>();
for (int64_t nLoopIdx = 0; nLoopIdx < nLoopCnt; nLoopIdx++) {
int64_t currNNum = (nLoopIdx != nLoopCnt - 1) ? tilingData_->ubFactorN :
nBlockCount_ - (nLoopIdx * tilingData_->ubFactorN);
ProcessLoop(xGmOffset, cosBuffer, sinBuffer, nLoopIdx, currBSNum, currNNum);
}
cosInQueue_.FreeTensor(cosBuffer);
sinInQueue_.FreeTensor(sinBuffer);
}
}
template <typename T>
__aicore__ inline void RotaryPositionEmbeddingAB<T>::ProcessLoop(
int64_t xGmOffset, LocalTensor<T> cosBuffer, LocalTensor<T> sinBuffer, int64_t ubIdx, int64_t bsCount,
int64_t nCount)
{
int64_t totalCount = bsCount * nCount;
DataCopyExtParams inParams = {static_cast<uint16_t>(totalCount * tilingData_->dSplitCoef), dSplitSize_, static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0, 0};
DataCopyExtParams outParams = {static_cast<uint16_t>(totalCount * tilingData_->dSplitCoef), dSplitSize_, 0, static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0};
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
inParams = {static_cast<uint16_t>(totalCount), tilingData_->D * sizeof(T), static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0, 0};
}
LocalTensor<T> inBuffer = xInQueue_.AllocTensor<T>();
LocalTensor<T> outBuffer = yOutQueue_.AllocTensor<T>();
DataCopyPad(inBuffer, xGm_[xGmOffset + ubIdx * tilingData_->ubFactorN * tilingData_->D], inParams, padParams_);
xInQueue_.EnQue(inBuffer);
inBuffer = xInQueue_.DeQue<T>();
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF)) {
HalfAlignVF(sinBuffer, cosBuffer, inBuffer, outBuffer, tilingData_->sliceLength, sliceAlign_, bsCount, nCount);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::INTERLEAVE)) {
InterleaveModeVF(sinBuffer, cosBuffer, inBuffer, outBuffer, tilingData_->sliceLength, bsCount, nCount);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
QuarterAlignVF(sinBuffer, cosBuffer, inBuffer, outBuffer, tilingData_->sliceLength, sliceAlign_, bsCount, nCount);
} else {
DeepSeekInterleaveModeVF<T>(sinBuffer, cosBuffer, inBuffer, outBuffer, tilingData_->sliceLength, bsCount, nCount);
}
yOutQueue_.EnQue(outBuffer);
outBuffer = yOutQueue_.DeQue<T>();
xInQueue_.FreeTensor(inBuffer);
DataCopyPad(yGm_[xGmOffset + ubIdx * tilingData_->ubFactorN * tilingData_->D], outBuffer, outParams);
yOutQueue_.FreeTensor(outBuffer);
}
} // namespace InplacePartialRotaryMul
#endif // ROTARY_POSITION_EMBEDDING_REG_AB_H

175
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_ab_mixed.h Normal file
View File

@@ -0,0 +1,175 @@
/**
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
* CANN Open Software License Agreement Version 2.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
*/
/*!
* \file rotary_position_embedding_reg_ab_mixed.h
* \brief Mixed precision kernel for AB layout: x is half/bfloat16, cos/sin are float
*/
#ifndef ROTARY_POSITION_EMBEDDING_REG_AB_MIXED_H
#define ROTARY_POSITION_EMBEDDING_REG_AB_MIXED_H
#include "apply_rotary_pos_emb_common.h"
namespace InplacePartialRotaryMul {
using namespace AscendC;
template <typename TX>
class RotaryPositionEmbeddingABMixed {
public:
__aicore__ inline RotaryPositionEmbeddingABMixed(){};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, GM_ADDR workspace,
const RopeRegbaseTilingData *tilingData, TPipe *pipe);
__aicore__ inline void Process();
private:
__aicore__ inline void ProcessLoop(int64_t xGmOffset, LocalTensor<float> cosBuffer, LocalTensor<float> sinBuffer,
int64_t ubIdx, int64_t bsCount, int64_t nCount);
private:
TPipe *pipe_;
TQue<QuePosition::VECIN, 1> xInQueue_;
TQue<QuePosition::VECIN, 1> cosInQueue_;
TQue<QuePosition::VECIN, 1> sinInQueue_;
TQue<QuePosition::VECOUT, 1> yOutQueue_;
GlobalTensor<TX> xGm_;
GlobalTensor<float> cosGm_;
GlobalTensor<float> sinGm_;
GlobalTensor<TX> yGm_;
const RopeRegbaseTilingData *tilingData_;
DataCopyPadExtParams<TX> padParams_ = {false, 0, 0, static_cast<TX>(0)};
DataCopyPadExtParams<float> padParamsFloat_ = {false, 0, 0, 0};
uint8_t DB_FLAG = 2;
uint32_t dSplitSizeTX_ = 0;
uint32_t dSplitSizeFloat_ = 0;
int64_t bsBlockCount_ = 0;
int64_t nBlockCount_ = 0;
int64_t sliceAlignTX_ = 0;
int64_t sliceAlignFloat_ = 0;
};
template <typename TX>
__aicore__ inline void RotaryPositionEmbeddingABMixed<TX>::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
GM_ADDR workspace, const RopeRegbaseTilingData *tilingData, TPipe *pipe)
{
pipe_ = pipe;
tilingData_ = tilingData;
dSplitSizeTX_ = tilingData_->sliceLength / tilingData_->dSplitCoef * sizeof(TX);
dSplitSizeFloat_ = tilingData_->sliceLength / tilingData_->dSplitCoef * sizeof(float);
int64_t blockDimBS = GetBlockIdx() / tilingData_->blockNumN;
int64_t blockDimN = GetBlockIdx() % tilingData_->blockNumN;
bsBlockCount_ = (blockDimBS == tilingData_->blockNumBS - 1) ? tilingData_->blockTailBS : tilingData_->blockFactorBS;
nBlockCount_ = (blockDimN == tilingData_->blockNumN - 1) ? tilingData_->blockTailN : tilingData_->blockFactorN;
int64_t cosOffset = blockDimBS * tilingData_->blockFactorBS * tilingData_->sliceLength;
int64_t offset = blockDimBS * tilingData_->blockFactorBS * tilingData_->D;
int64_t xOffset =
offset * tilingData_->N + blockDimN * tilingData_->blockFactorN * tilingData_->D + tilingData_->sliceStart;
this->cosGm_.SetGlobalBuffer((__gm__ float *)cos + cosOffset);
this->sinGm_.SetGlobalBuffer((__gm__ float *)sin + cosOffset);
this->xGm_.SetGlobalBuffer((__gm__ TX *)x + xOffset);
this->yGm_.SetGlobalBuffer((__gm__ TX *)y + xOffset);
sliceAlignTX_ =
ops::CeilDiv(tilingData_->sliceLength * sizeof(TX), GetUbBlockSize()) * GetUbBlockSize() / sizeof(TX);
sliceAlignFloat_ =
ops::CeilDiv(tilingData_->sliceLength * sizeof(float), GetUbBlockSize()) * GetUbBlockSize() / sizeof(float);
int64_t bufferSizeTX = sliceAlignTX_ * sizeof(TX) * tilingData_->ubFactorBS;
int64_t bufferSizeFloat = sliceAlignFloat_ * sizeof(float) * tilingData_->ubFactorBS;
pipe_->InitBuffer(xInQueue_, DB_FLAG, bufferSizeTX * tilingData_->ubFactorN);
pipe_->InitBuffer(yOutQueue_, DB_FLAG, bufferSizeTX * tilingData_->ubFactorN);
pipe_->InitBuffer(cosInQueue_, DB_FLAG, bufferSizeFloat);
pipe_->InitBuffer(sinInQueue_, DB_FLAG, bufferSizeFloat);
}
template <typename TX>
__aicore__ inline void RotaryPositionEmbeddingABMixed<TX>::Process()
{
uint32_t bsLoopCnt = ops::CeilDiv(bsBlockCount_, tilingData_->ubFactorBS);
uint32_t nLoopCnt = ops::CeilDiv(nBlockCount_, tilingData_->ubFactorN);
for (uint32_t bsLoopIdx = 0; bsLoopIdx < bsLoopCnt; bsLoopIdx++) {
int64_t xGmOffset = bsLoopIdx * tilingData_->ubFactorBS * tilingData_->N * tilingData_->D;
uint32_t currBSNum = (bsLoopIdx != bsLoopCnt - 1) ? tilingData_->ubFactorBS
: bsBlockCount_ - (bsLoopIdx * tilingData_->ubFactorBS);
DataCopyExtParams cosParams = {
static_cast<uint16_t>(currBSNum * tilingData_->dSplitCoef), dSplitSizeFloat_, 0, 0, 0};
LocalTensor<float> cosBuffer = cosInQueue_.AllocTensor<float>();
LocalTensor<float> sinBuffer = sinInQueue_.AllocTensor<float>();
DataCopyPad(cosBuffer,
cosGm_[bsLoopIdx * tilingData_->ubFactorBS * tilingData_->sliceLength],
cosParams,
padParamsFloat_);
cosInQueue_.EnQue(cosBuffer);
cosBuffer = cosInQueue_.DeQue<float>();
DataCopyPad(sinBuffer,
sinGm_[bsLoopIdx * tilingData_->ubFactorBS * tilingData_->sliceLength],
cosParams,
padParamsFloat_);
sinInQueue_.EnQue(sinBuffer);
sinBuffer = sinInQueue_.DeQue<float>();
for (int64_t nLoopIdx = 0; nLoopIdx < nLoopCnt; nLoopIdx++) {
int64_t currNNum = (nLoopIdx != nLoopCnt - 1) ? tilingData_->ubFactorN
: nBlockCount_ - (nLoopIdx * tilingData_->ubFactorN);
ProcessLoop(xGmOffset, cosBuffer, sinBuffer, nLoopIdx, currBSNum, currNNum);
}
cosInQueue_.FreeTensor(cosBuffer);
sinInQueue_.FreeTensor(sinBuffer);
}
}
template <typename TX>
__aicore__ inline void RotaryPositionEmbeddingABMixed<TX>::ProcessLoop(int64_t xGmOffset, LocalTensor<float> cosBuffer,
LocalTensor<float> sinBuffer, int64_t ubIdx, int64_t bsCount, int64_t nCount)
{
int64_t totalCount = bsCount * nCount;
DataCopyExtParams inParams = {static_cast<uint16_t>(totalCount * tilingData_->dSplitCoef),
dSplitSizeTX_,
static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)),
0,
0};
DataCopyExtParams outParams = {static_cast<uint16_t>(totalCount * tilingData_->dSplitCoef),
dSplitSizeTX_,
0,
static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)),
0};
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
inParams = {static_cast<uint16_t>(totalCount),
tilingData_->D * sizeof(TX),
static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)),
0,
0};
}
LocalTensor<TX> inBuffer = xInQueue_.AllocTensor<TX>();
LocalTensor<TX> outBuffer = yOutQueue_.AllocTensor<TX>();
DataCopyPad(inBuffer, xGm_[xGmOffset + ubIdx * tilingData_->ubFactorN * tilingData_->D], inParams, padParams_);
xInQueue_.EnQue(inBuffer);
inBuffer = xInQueue_.DeQue<TX>();
InterleaveModeVFMixed<TX>(inBuffer, cosBuffer, sinBuffer, outBuffer, tilingData_->sliceLength, bsCount, nCount);
yOutQueue_.EnQue(outBuffer);
outBuffer = yOutQueue_.DeQue<TX>();
xInQueue_.FreeTensor(inBuffer);
DataCopyPad(yGm_[xGmOffset + ubIdx * tilingData_->ubFactorN * tilingData_->D], outBuffer, outParams);
yOutQueue_.FreeTensor(outBuffer);
}
} // namespace InplacePartialRotaryMul
#endif // ROTARY_POSITION_EMBEDDING_REG_AB_MIXED_H

364
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_aba_and_ba.h Normal file
View File

@@ -0,0 +1,364 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotary_position_embedding_reg_aba_and_ba.h
* \brief
*/
#ifndef ROTARY_POSITION_EMBEDDING_REG_ABA_AND_BA_H
#define ROTARY_POSITION_EMBEDDING_REG_ABA_AND_BA_H
// #include "op_kernel/math_util.h"
#include "apply_rotary_pos_emb_common.h"
namespace InplacePartialRotaryMul {
using namespace AscendC;
template <typename T, bool IsBBoardcast>
class RotaryPositionEmbeddingABAAndBA
{
public:
__aicore__ inline RotaryPositionEmbeddingABAAndBA(){};
__aicore__ inline ~RotaryPositionEmbeddingABAAndBA(){};
__aicore__ inline void Init(
GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData,
TPipe* pipe);
__aicore__ inline void Process();
private:
// Init过程中使用的内部函数
__aicore__ inline void InitAllGlobalBuffer(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut);
__aicore__ inline void InitAllBuffer();
__aicore__ inline void InitLoopParams();
// 各个层级的Process函数
__aicore__ inline void ProcessInSLoop(
int64_t sUbStart,
int64_t sUbLength); // 第一重循环体，给定S范围，沿B轴进行遍历处理
__aicore__ inline void ProcessInSBLoop(
int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, LocalTensor<T>& cos,
LocalTensor<T>& sin); // 第二重循环体，给定BS范围，沿Q和K的N轴进行遍历处理
__aicore__ inline void ProcessInSBNLoop(
int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, int64_t nUbStart, int64_t nUbLength,
int64_t nTotalSize, LocalTensor<T>& cos, LocalTensor<T>& sin, GlobalTensor<T>& in,
GlobalTensor<T>& out); // 第三重循环体，给定BSN范围，计算其中数据的rope
// 拷入拷出函数
__aicore__ inline void CopyInCosAndSin(int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength);
__aicore__ inline void CopyInQ(
GlobalTensor<T>& source, int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart,
int64_t nLength, int64_t nTotalSize);
__aicore__ inline void CopyOutQ(
GlobalTensor<T>& target, int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart,
int64_t nLength, int64_t nTotalSize);
// 计算函数
__aicore__ inline void Compute(
LocalTensor<T>& cos, LocalTensor<T>& sin, int64_t sLength, int64_t bLength, int64_t nLength);
private:
TPipe* pipe_;
// GlobalMemory
GlobalTensor<T> qGm_;
GlobalTensor<T> cosGm_;
GlobalTensor<T> sinGm_;
GlobalTensor<T> qOutGm_;
// UB
TQue<QuePosition::VECIN, DOUBLE_BUFFER> qInQueue_;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> cosInQueue_;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> sinInQueue_;
TQue<QuePosition::VECOUT, DOUBLE_BUFFER> qOutQueue_;
// Split core info
int64_t blockIdx_ = 0;
int64_t bBlockStart_ = 0;
int64_t bBlockLength_ = 0;
int64_t sBlockStart_ = 0;
int64_t sBlockLength_ = 0;
// TilingData
const RopeRegbaseTilingData* tilingData_;
int64_t ubFactorB_ = 0;
int64_t ubFactorS_ = 0;
int64_t ubFactorN_ = 0;
int64_t D_ = 0;
int64_t dAlign_ = 0;
// 拷贝参数
uint8_t dSplitCoef_ = 1;
uint8_t copyInQSplitCoef_ = 1; // 拷贝q时使用的splitCoef
uint64_t ubCopyInStride = 0; // 输入在ub中的stride，deepseek_interleave中不为0
};
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::Init(
GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData,
TPipe* pipe)
{
this->tilingData_ = tilingData;
this->blockIdx_ = GetBlockIdx();
this->pipe_ = pipe;
this->InitAllGlobalBuffer(q, cos, sin, qOut);
this->InitAllBuffer();
this->InitLoopParams();
}
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::InitAllGlobalBuffer(
GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut)
{
this->qGm_.SetGlobalBuffer((__gm__ T*)q);
this->cosGm_.SetGlobalBuffer((__gm__ T*)cos);
this->sinGm_.SetGlobalBuffer((__gm__ T*)sin);
this->qOutGm_.SetGlobalBuffer((__gm__ T*)qOut);
}
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::InitAllBuffer()
{
this->ubFactorB_ = this->tilingData_->ubFactorB;
this->ubFactorS_ = this->tilingData_->ubFactorS;
this->ubFactorN_ = this->tilingData_->ubFactorN;
this->D_ = this->tilingData_->D;
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF) ||
tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
this->dSplitCoef_ = HALF_INTERLEAVE_COEF;
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
this->dSplitCoef_ = QUARTER_MODE_COEF;
}
this->copyInQSplitCoef_ = dSplitCoef_;
this->dAlign_ = ops::CeilAlign<int64_t>(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(T)) * dSplitCoef_;
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
this->copyInQSplitCoef_ = 1;
this->ubCopyInStride =
(this->dAlign_ * sizeof(T) - ops::CeilAlign<int64_t>(tilingData_->sliceLength * sizeof(T), BLOCK_TYPE_SIZE)) / BLOCK_TYPE_SIZE;
}
this->pipe_->InitBuffer(this->qInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(T));
this->pipe_->InitBuffer(
this->qOutQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(T));
if constexpr (IsBBoardcast) {
this->pipe_->InitBuffer(this->cosInQueue_, DOUBLE_BUFFER, ubFactorS_ * dAlign_ * sizeof(T));
this->pipe_->InitBuffer(this->sinInQueue_, DOUBLE_BUFFER, ubFactorS_ * dAlign_ * sizeof(T));
} else {
this->pipe_->InitBuffer(this->cosInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * dAlign_ * sizeof(T));
this->pipe_->InitBuffer(this->sinInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * dAlign_ * sizeof(T));
}
}
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::InitLoopParams()
{
int64_t bIdx = blockIdx_ % tilingData_->blockNumB;
int64_t sIdx = blockIdx_ / tilingData_->blockNumB;
this->bBlockLength_ = tilingData_->blockFactorB;
this->sBlockLength_ = tilingData_->blockFactorS;
if (bIdx == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) {
this->bBlockLength_ = tilingData_->B % tilingData_->blockFactorB;
}
if (sIdx == tilingData_->blockNumS - 1 && tilingData_->S % tilingData_->blockFactorS != 0) {
this->sBlockLength_ = tilingData_->S % tilingData_->blockFactorS;
}
this->bBlockStart_ = bIdx * tilingData_->blockFactorB;
this->sBlockStart_ = sIdx * tilingData_->blockFactorS;
}
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::Process()
{
// 在S轴进行循环
int64_t ubLoopCount = ops::CeilDiv(sBlockLength_, ubFactorS_);
for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) {
this->ProcessInSLoop(
sBlockStart_ + ubLoopIdx * ubFactorS_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorS_ : sBlockLength_ - ubLoopIdx * ubFactorS_);
}
}
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::ProcessInSLoop(
int64_t sUbStart, int64_t sUbLength)
{
// 在B轴进行循环
int64_t ubLoopCount = ops::CeilDiv(bBlockLength_, ubFactorB_);
if constexpr (IsBBoardcast) {
// cos和sin需要在B轴广播的情况
this->CopyInCosAndSin(sUbStart, sUbLength, 0, 1);
LocalTensor<T> cosUb = this->cosInQueue_.template DeQue<T>();
LocalTensor<T> sinUb = this->sinInQueue_.template DeQue<T>();
for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) {
this->ProcessInSBLoop(
sUbStart, sUbLength, bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_, cosUb, sinUb);
}
this->sinInQueue_.FreeTensor(cosUb);
this->cosInQueue_.FreeTensor(sinUb);
} else {
// sin和cos无需在B轴广播的情况
for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) {
this->CopyInCosAndSin(
sUbStart, sUbLength, bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_);
LocalTensor<T> cosUb = this->cosInQueue_.template DeQue<T>();
LocalTensor<T> sinUb = this->sinInQueue_.template DeQue<T>();
this->ProcessInSBLoop(
sUbStart, sUbLength, bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_, cosUb, sinUb);
this->cosInQueue_.FreeTensor(cosUb);
this->sinInQueue_.FreeTensor(sinUb);
}
}
}
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::ProcessInSBLoop(
int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, LocalTensor<T>& cos, LocalTensor<T>& sin)
{
// 循环处理Q
int64_t qUbLoopCount = ops::CeilDiv(tilingData_->N, ubFactorN_);
for (int64_t ubLoopIdx = 0; ubLoopIdx < qUbLoopCount; ubLoopIdx++) {
this->ProcessInSBNLoop(
sUbStart, sUbLength, bUbStart, bUbLength, ubLoopIdx * ubFactorN_,
ubLoopIdx != qUbLoopCount - 1 ? ubFactorN_ : tilingData_->N - ubLoopIdx * ubFactorN_, tilingData_->N, cos,
sin, qGm_, qOutGm_);
}
}
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::ProcessInSBNLoop(
int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, int64_t nUbStart, int64_t nUbLength,
int64_t nTotalSize, LocalTensor<T>& cos, LocalTensor<T>& sin, GlobalTensor<T>& in, GlobalTensor<T>& out)
{
CopyInQ(in, sUbStart, sUbLength, bUbStart, bUbLength, nUbStart, nUbLength, nTotalSize);
Compute(cos, sin, sUbLength, bUbLength, nUbLength);
CopyOutQ(out, sUbStart, sUbLength, bUbStart, bUbLength, nUbStart, nUbLength, nTotalSize);
}
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::CopyInCosAndSin(
int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength)
{
LocalTensor<T> cosUb = this->cosInQueue_.template AllocTensor<T>();
LocalTensor<T> sinUb = this->sinInQueue_.template AllocTensor<T>();
LoopModeParams loopParams;
loopParams.loop2Size = 1;
loopParams.loop1Size = bLength;
loopParams.loop2SrcStride = 0;
loopParams.loop2DstStride = 0;
loopParams.loop1SrcStride = tilingData_->S * tilingData_->sliceLength * sizeof(T);
loopParams.loop1DstStride = ubFactorS_ * dAlign_ * sizeof(T);
SetLoopModePara(loopParams, DataCopyMVType::OUT_TO_UB);
DataCopyPadExtParams<T> copyPadExtparams;
copyPadExtparams.isPad = false;
copyPadExtparams.leftPadding = 0;
copyPadExtparams.rightPadding = 0;
copyPadExtparams.paddingValue = 0;
DataCopyExtParams copyExtParams;
copyExtParams.blockCount = sLength * dSplitCoef_;
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / dSplitCoef_;
copyExtParams.srcStride = 0;
copyExtParams.dstStride = 0;
DataCopyPad(cosUb, this->cosGm_[bStart * tilingData_->S * tilingData_->sliceLength + sStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams);
DataCopyPad(sinUb, this->sinGm_[bStart * tilingData_->S * tilingData_->sliceLength + sStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams);
ResetLoopModePara(DataCopyMVType::OUT_TO_UB);
this->cosInQueue_.template EnQue(cosUb);
this->sinInQueue_.template EnQue(sinUb);
}
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::CopyInQ(
GlobalTensor<T>& source, int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart,
int64_t nLength, int64_t nTotalSize)
{
LocalTensor<T> target = this->qInQueue_.template AllocTensor<T>();
// 数据格式为BNSD，B->N->S->D
LoopModeParams loopParams;
loopParams.loop2Size = bLength;
loopParams.loop1Size = nLength;
loopParams.loop2SrcStride = nTotalSize * tilingData_->S * D_ * sizeof(T);
loopParams.loop2DstStride = ubFactorN_ * ubFactorS_ * dAlign_ * sizeof(T);
loopParams.loop1SrcStride = tilingData_->S * D_ * sizeof(T);
loopParams.loop1DstStride = ubFactorS_ * dAlign_ * sizeof(T);
SetLoopModePara(loopParams, DataCopyMVType::OUT_TO_UB);
DataCopyExtParams copyExtParams;
copyExtParams.blockCount = sLength * copyInQSplitCoef_;
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / copyInQSplitCoef_;
copyExtParams.srcStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(T);
copyExtParams.dstStride = ubCopyInStride;
DataCopyPadExtParams<T> copyPadExtparams;
copyPadExtparams.isPad = false;
copyPadExtparams.leftPadding = 0;
copyPadExtparams.rightPadding = 0;
copyPadExtparams.paddingValue = 0;
DataCopyPad(
target, source[bStart * nTotalSize * tilingData_->S * D_ + nStart * tilingData_->S * D_ + sStart * D_ + tilingData_->sliceStart],
copyExtParams, copyPadExtparams);
ResetLoopModePara(DataCopyMVType::OUT_TO_UB);
this->qInQueue_.template EnQue(target);
}
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::CopyOutQ(
GlobalTensor<T>& target, int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart,
int64_t nLength, int64_t nTotalSize)
{
LocalTensor<T> source = this->qOutQueue_.template DeQue<T>();
// 数据格式为BNSD，B->N->S->D
LoopModeParams loopParams;
loopParams.loop2Size = bLength;
loopParams.loop1Size = nLength;
loopParams.loop2DstStride = nTotalSize * tilingData_->S * D_ * sizeof(T);
loopParams.loop2SrcStride = ubFactorN_ * ubFactorS_ * dAlign_ * sizeof(T);
loopParams.loop1DstStride = tilingData_->S * D_ * sizeof(T);
loopParams.loop1SrcStride = ubFactorS_ * dAlign_ * sizeof(T);
SetLoopModePara(loopParams, DataCopyMVType::UB_TO_OUT);
DataCopyExtParams copyExtParams;
copyExtParams.blockCount = sLength * dSplitCoef_;
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / dSplitCoef_;
copyExtParams.srcStride = 0;
copyExtParams.dstStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(T);
DataCopyPad(
target[bStart * nTotalSize * tilingData_->S * D_ + nStart * tilingData_->S * D_ + sStart * D_ + tilingData_->sliceStart], source,
copyExtParams);
ResetLoopModePara(DataCopyMVType::UB_TO_OUT);
this->qOutQueue_.FreeTensor(source);
}
template <typename T, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBA<T, IsBBoardcast>::Compute(
LocalTensor<T>& cos, LocalTensor<T>& sin, int64_t sLength, int64_t bLength, int64_t nLength)
{
LocalTensor<T> inUb = this->qInQueue_.template DeQue<T>();
LocalTensor<T> outUb = this->qOutQueue_.template AllocTensor<T>();
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF)) {
BatchHalfAlignVF<T, IsBBoardcast>(
(__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), (__local_mem__ T*)sin.GetPhyAddr(),
(__local_mem__ T*)outUb.GetPhyAddr(), sLength, bLength, nLength, tilingData_->sliceLength, dAlign_, ubFactorS_, ubFactorN_);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::INTERLEAVE)) {
BatchInterleaveModeVF<T, IsBBoardcast>(
(__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), (__local_mem__ T*)sin.GetPhyAddr(),
(__local_mem__ T*)outUb.GetPhyAddr(), sLength, bLength, nLength, tilingData_->sliceLength, dAlign_, ubFactorS_, ubFactorN_);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
BatchQuarterAlignVF<T, IsBBoardcast>(
(__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), (__local_mem__ T*)sin.GetPhyAddr(),
(__local_mem__ T*)outUb.GetPhyAddr(), sLength, bLength, nLength, tilingData_->sliceLength, dAlign_, ubFactorS_, ubFactorN_);
} else {
BatchDeepSeekInterleaveModeVF<T, IsBBoardcast>(
(__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), (__local_mem__ T*)sin.GetPhyAddr(),
(__local_mem__ T*)outUb.GetPhyAddr(), sLength, bLength, nLength, tilingData_->sliceLength, dAlign_, ubFactorS_, ubFactorN_);
}
this->qInQueue_.FreeTensor(inUb);
this->qOutQueue_.template EnQue(outUb);
}
} // namespace InplacePartialRotaryMul
#endif

364
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_aba_and_ba_mixed.h Normal file
View File

@@ -0,0 +1,364 @@
/**
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
* CANN Open Software License Agreement Version 2.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
*/
/*!
* \file rotary_position_embedding_reg_aba_and_ba_mixed.h
* \brief Mixed precision kernel for ABA/BA layout: x is half/bfloat16, cos/sin are float
*/
#ifndef ROTARY_POSITION_EMBEDDING_REG_ABA_AND_BA_MIXED_H
#define ROTARY_POSITION_EMBEDDING_REG_ABA_AND_BA_MIXED_H
#include "apply_rotary_pos_emb_common.h"
namespace InplacePartialRotaryMul {
using namespace AscendC;
template <typename TX, bool IsBBoardcast>
class RotaryPositionEmbeddingABAAndBAMixed {
public:
__aicore__ inline RotaryPositionEmbeddingABAAndBAMixed(){};
__aicore__ inline ~RotaryPositionEmbeddingABAAndBAMixed(){};
__aicore__ inline void Init(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace,
const RopeRegbaseTilingData *tilingData, TPipe *pipe);
__aicore__ inline void Process();
private:
__aicore__ inline void InitAllGlobalBuffer(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut);
__aicore__ inline void InitAllBuffer();
__aicore__ inline void InitLoopParams();
__aicore__ inline void ProcessInSLoop(int64_t sUbStart, int64_t sUbLength);
__aicore__ inline void ProcessInSBLoop(int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength,
LocalTensor<TX> &cos, LocalTensor<TX> &sin);
__aicore__ inline void ProcessInSBNLoop(int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength,
int64_t nUbStart, int64_t nUbLength, int64_t nTotalSize, LocalTensor<float> &cosFloat,
LocalTensor<float> &sinFloat, GlobalTensor<TX> &in, GlobalTensor<TX> &out);
__aicore__ inline void CopyInCosAndSin(int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength);
__aicore__ inline void CopyInQ(GlobalTensor<TX> &source, int64_t sStart, int64_t sLength, int64_t bStart,
int64_t bLength, int64_t nStart, int64_t nLength, int64_t nTotalSize);
__aicore__ inline void CopyOutQ(GlobalTensor<TX> &target, int64_t sStart, int64_t sLength, int64_t bStart,
int64_t bLength, int64_t nStart, int64_t nLength, int64_t nTotalSize);
__aicore__ inline void Compute(
LocalTensor<float> &cos, LocalTensor<float> &sin, int64_t sLength, int64_t bLength, int64_t nLength);
private:
TPipe *pipe_;
GlobalTensor<TX> qGm_;
GlobalTensor<float> cosGm_;
GlobalTensor<float> sinGm_;
GlobalTensor<TX> qOutGm_;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> qInQueue_;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> cosInQueue_;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> sinInQueue_;
TQue<QuePosition::VECOUT, DOUBLE_BUFFER> qOutQueue_;
int64_t blockIdx_ = 0;
int64_t bBlockStart_ = 0;
int64_t bBlockLength_ = 0;
int64_t sBlockStart_ = 0;
int64_t sBlockLength_ = 0;
const RopeRegbaseTilingData *tilingData_;
int64_t ubFactorB_ = 0;
int64_t ubFactorS_ = 0;
int64_t ubFactorN_ = 0;
int64_t D_ = 0;
int64_t dAlign_ = 0;
int64_t dAlignFloat_ = 0;
uint8_t dSplitCoef_ = 1;
uint8_t copyInQSplitCoef_ = 1;
uint64_t ubCopyInStride = 0;
};
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::Init(GM_ADDR q, GM_ADDR cos, GM_ADDR sin,
GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData *tilingData, TPipe *pipe)
{
this->tilingData_ = tilingData;
this->blockIdx_ = GetBlockIdx();
this->pipe_ = pipe;
this->InitAllGlobalBuffer(q, cos, sin, qOut);
this->InitAllBuffer();
this->InitLoopParams();
}
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::InitAllGlobalBuffer(
GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut)
{
this->qGm_.SetGlobalBuffer((__gm__ TX *)q);
this->cosGm_.SetGlobalBuffer((__gm__ float *)cos);
this->sinGm_.SetGlobalBuffer((__gm__ float *)sin);
this->qOutGm_.SetGlobalBuffer((__gm__ TX *)qOut);
}
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::InitAllBuffer()
{
this->ubFactorB_ = this->tilingData_->ubFactorB;
this->ubFactorS_ = this->tilingData_->ubFactorS;
this->ubFactorN_ = this->tilingData_->ubFactorN;
this->D_ = this->tilingData_->D;
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF) ||
tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
this->dSplitCoef_ = HALF_INTERLEAVE_COEF;
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
this->dSplitCoef_ = QUARTER_MODE_COEF;
}
this->copyInQSplitCoef_ = dSplitCoef_;
this->dAlign_ =
ops::CeilAlign<int64_t>(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(TX)) * dSplitCoef_;
this->dAlignFloat_ =
ops::CeilAlign<int64_t>(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(float)) * dSplitCoef_;
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
this->copyInQSplitCoef_ = 1;
this->ubCopyInStride = (this->dAlign_ * sizeof(TX) -
ops::CeilAlign<int64_t>(tilingData_->sliceLength * sizeof(TX), BLOCK_TYPE_SIZE)) /
BLOCK_TYPE_SIZE;
}
this->pipe_->InitBuffer(
this->qInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(TX));
this->pipe_->InitBuffer(
this->qOutQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(TX));
if constexpr (IsBBoardcast) {
this->pipe_->InitBuffer(this->cosInQueue_, DOUBLE_BUFFER, ubFactorS_ * dAlignFloat_ * sizeof(float));
this->pipe_->InitBuffer(this->sinInQueue_, DOUBLE_BUFFER, ubFactorS_ * dAlignFloat_ * sizeof(float));
} else {
this->pipe_->InitBuffer(
this->cosInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * dAlignFloat_ * sizeof(float));
this->pipe_->InitBuffer(
this->sinInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * dAlignFloat_ * sizeof(float));
}
}
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::InitLoopParams()
{
int64_t bIdx = blockIdx_ % tilingData_->blockNumB;
int64_t sIdx = blockIdx_ / tilingData_->blockNumB;
this->bBlockLength_ = tilingData_->blockFactorB;
this->sBlockLength_ = tilingData_->blockFactorS;
if (bIdx == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) {
this->bBlockLength_ = tilingData_->B % tilingData_->blockFactorB;
}
if (sIdx == tilingData_->blockNumS - 1 && tilingData_->S % tilingData_->blockFactorS != 0) {
this->sBlockLength_ = tilingData_->S % tilingData_->blockFactorS;
}
this->bBlockStart_ = bIdx * tilingData_->blockFactorB;
this->sBlockStart_ = sIdx * tilingData_->blockFactorS;
}
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::Process()
{
int64_t ubLoopCount = ops::CeilDiv(sBlockLength_, ubFactorS_);
for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) {
this->ProcessInSLoop(sBlockStart_ + ubLoopIdx * ubFactorS_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorS_ : sBlockLength_ - ubLoopIdx * ubFactorS_);
}
}
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::ProcessInSLoop(
int64_t sUbStart, int64_t sUbLength)
{
int64_t ubLoopCount = ops::CeilDiv(bBlockLength_, ubFactorB_);
if constexpr (IsBBoardcast) {
this->CopyInCosAndSin(sUbStart, sUbLength, 0, 1);
LocalTensor<float> cosUbFloat = this->cosInQueue_.template DeQue<float>();
LocalTensor<float> sinUbFloat = this->sinInQueue_.template DeQue<float>();
LocalTensor<TX> cosUb = cosUbFloat.template ReinterpretCast<TX>();
LocalTensor<TX> sinUb = sinUbFloat.template ReinterpretCast<TX>();
for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) {
this->ProcessInSBLoop(sUbStart,
sUbLength,
bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_,
cosUb,
sinUb);
}
this->sinInQueue_.FreeTensor(sinUbFloat);
this->cosInQueue_.FreeTensor(cosUbFloat);
} else {
for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) {
this->CopyInCosAndSin(sUbStart,
sUbLength,
bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_);
LocalTensor<float> cosUbFloat = this->cosInQueue_.template DeQue<float>();
LocalTensor<float> sinUbFloat = this->sinInQueue_.template DeQue<float>();
LocalTensor<TX> cosUb = cosUbFloat.template ReinterpretCast<TX>();
LocalTensor<TX> sinUb = sinUbFloat.template ReinterpretCast<TX>();
this->ProcessInSBLoop(sUbStart,
sUbLength,
bBlockStart_ + ubLoopIdx * ubFactorB_,
ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_,
cosUb,
sinUb);
this->cosInQueue_.FreeTensor(cosUbFloat);
this->sinInQueue_.FreeTensor(sinUbFloat);
}
}
}
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::ProcessInSBLoop(int64_t sUbStart,
int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, LocalTensor<TX> &cos, LocalTensor<TX> &sin)
{
int64_t qUbLoopCount = ops::CeilDiv(tilingData_->N, ubFactorN_);
LocalTensor<float> cosFloat = cos.template ReinterpretCast<float>();
LocalTensor<float> sinFloat = sin.template ReinterpretCast<float>();
for (int64_t ubLoopIdx = 0; ubLoopIdx < qUbLoopCount; ubLoopIdx++) {
this->ProcessInSBNLoop(sUbStart,
sUbLength,
bUbStart,
bUbLength,
ubLoopIdx * ubFactorN_,
ubLoopIdx != qUbLoopCount - 1 ? ubFactorN_ : tilingData_->N - ubLoopIdx * ubFactorN_,
tilingData_->N,
cosFloat,
sinFloat,
qGm_,
qOutGm_);
}
}
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::ProcessInSBNLoop(int64_t sUbStart,
int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, int64_t nUbStart, int64_t nUbLength, int64_t nTotalSize,
LocalTensor<float> &cosFloat, LocalTensor<float> &sinFloat, GlobalTensor<TX> &in, GlobalTensor<TX> &out)
{
CopyInQ(in, sUbStart, sUbLength, bUbStart, bUbLength, nUbStart, nUbLength, nTotalSize);
Compute(cosFloat, sinFloat, sUbLength, bUbLength, nUbLength);
CopyOutQ(out, sUbStart, sUbLength, bUbStart, bUbLength, nUbStart, nUbLength, nTotalSize);
}
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::CopyInCosAndSin(
int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength)
{
LocalTensor<float> cosUb = this->cosInQueue_.template AllocTensor<float>();
LocalTensor<float> sinUb = this->sinInQueue_.template AllocTensor<float>();
LoopModeParams loopParams;
loopParams.loop2Size = 1;
loopParams.loop1Size = bLength;
loopParams.loop2SrcStride = 0;
loopParams.loop2DstStride = 0;
loopParams.loop1SrcStride = tilingData_->S * tilingData_->sliceLength * sizeof(float);
loopParams.loop1DstStride = ubFactorS_ * dAlignFloat_ * sizeof(float);
SetLoopModePara(loopParams, DataCopyMVType::OUT_TO_UB);
DataCopyPadExtParams<float> copyPadExtparams;
copyPadExtparams.isPad = false;
copyPadExtparams.leftPadding = 0;
copyPadExtparams.rightPadding = 0;
copyPadExtparams.paddingValue = 0;
DataCopyExtParams copyExtParams;
copyExtParams.blockCount = sLength * dSplitCoef_;
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(float) / dSplitCoef_;
copyExtParams.srcStride = 0;
copyExtParams.dstStride = 0;
DataCopyPad(cosUb,
this->cosGm_[bStart * tilingData_->S * tilingData_->sliceLength + sStart * tilingData_->sliceLength],
copyExtParams,
copyPadExtparams);
DataCopyPad(sinUb,
this->sinGm_[bStart * tilingData_->S * tilingData_->sliceLength + sStart * tilingData_->sliceLength],
copyExtParams,
copyPadExtparams);
ResetLoopModePara(DataCopyMVType::OUT_TO_UB);
this->cosInQueue_.template EnQue(cosUb);
this->sinInQueue_.template EnQue(sinUb);
}
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::CopyInQ(GlobalTensor<TX> &source,
int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart, int64_t nLength,
int64_t nTotalSize)
{
LocalTensor<TX> target = this->qInQueue_.template AllocTensor<TX>();
LoopModeParams loopParams;
loopParams.loop2Size = bLength;
loopParams.loop1Size = nLength;
loopParams.loop2SrcStride = nTotalSize * tilingData_->S * D_ * sizeof(TX);
loopParams.loop2DstStride = ubFactorN_ * ubFactorS_ * dAlign_ * sizeof(TX);
loopParams.loop1SrcStride = tilingData_->S * D_ * sizeof(TX);
loopParams.loop1DstStride = ubFactorS_ * dAlign_ * sizeof(TX);
SetLoopModePara(loopParams, DataCopyMVType::OUT_TO_UB);
DataCopyExtParams copyExtParams;
copyExtParams.blockCount = sLength * copyInQSplitCoef_;
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(TX) / copyInQSplitCoef_;
copyExtParams.srcStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(TX);
copyExtParams.dstStride = ubCopyInStride;
DataCopyPadExtParams<TX> copyPadExtparams;
copyPadExtparams.isPad = false;
int64_t offset = bStart * nTotalSize * tilingData_->S * D_ + nStart * tilingData_->S * D_ + sStart * D_ +
tilingData_->sliceStart;
DataCopyPad(target, source[offset], copyExtParams, copyPadExtparams);
ResetLoopModePara(DataCopyMVType::OUT_TO_UB);
this->qInQueue_.template EnQue(target);
}
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::CopyOutQ(GlobalTensor<TX> &target,
int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart, int64_t nLength,
int64_t nTotalSize)
{
LocalTensor<TX> source = this->qOutQueue_.template DeQue<TX>();
LoopModeParams loopParams;
loopParams.loop2Size = bLength;
loopParams.loop1Size = nLength;
loopParams.loop2SrcStride = ubFactorN_ * ubFactorS_ * dAlign_ * sizeof(TX);
loopParams.loop2DstStride = nTotalSize * tilingData_->S * D_ * sizeof(TX);
loopParams.loop1SrcStride = ubFactorS_ * dAlign_ * sizeof(TX);
loopParams.loop1DstStride = tilingData_->S * D_ * sizeof(TX);
SetLoopModePara(loopParams, DataCopyMVType::UB_TO_OUT);
DataCopyExtParams copyExtParams;
copyExtParams.blockCount = sLength * dSplitCoef_;
copyExtParams.blockLen = tilingData_->sliceLength * sizeof(TX) / dSplitCoef_;
copyExtParams.srcStride = 0;
copyExtParams.dstStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(TX);
int64_t offset = bStart * nTotalSize * tilingData_->S * D_ + nStart * tilingData_->S * D_ + sStart * D_ +
tilingData_->sliceStart;
DataCopyPad(target[offset], source, copyExtParams);
ResetLoopModePara(DataCopyMVType::UB_TO_OUT);
this->qOutQueue_.FreeTensor(source);
}
template <typename TX, bool IsBBoardcast>
__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed<TX, IsBBoardcast>::Compute(
LocalTensor<float> &cos, LocalTensor<float> &sin, int64_t sLength, int64_t bLength, int64_t nLength)
{
LocalTensor<TX> inUb = this->qInQueue_.template DeQue<TX>();
LocalTensor<TX> outUb = this->qOutQueue_.template AllocTensor<TX>();
int64_t totalLength = sLength * bLength * nLength * tilingData_->sliceLength;
BatchInterleaveModeVFMixed<TX, IsBBoardcast>((__local_mem__ TX *)inUb.GetPhyAddr(),
(__local_mem__ float *)cos.GetPhyAddr(),
(__local_mem__ float *)sin.GetPhyAddr(),
(__local_mem__ TX *)outUb.GetPhyAddr(),
sLength,
bLength,
nLength,
tilingData_->sliceLength,
dAlign_,
dAlignFloat_,
ubFactorS_,
ubFactorN_);
this->qInQueue_.FreeTensor(inUb);
this->qOutQueue_.template EnQue(outUb);
}
} // namespace InplacePartialRotaryMul
#endif // ROTARY_POSITION_EMBEDDING_REG_ABA_AND_BA_MIXED_H

192
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_bab.h Normal file
View File

@@ -0,0 +1,192 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotary_position_embedding_reg_bab.h
* \brief
*/
#ifndef ROTARY_POSITION_EMBEDDING_REG_BAB_H
#define ROTARY_POSITION_EMBEDDING_REG_BAB_H
// #include "op_kernel/math_util.h"
#include "apply_rotary_pos_emb_common.h"
namespace InplacePartialRotaryMul {
using namespace AscendC;
template <typename T>
class RotaryPositionEmbeddingBAB
{
public:
__aicore__ inline RotaryPositionEmbeddingBAB(TPipe* pipe, const RopeRegbaseTilingData* tiling)
: pipe_(pipe), tilingData_(tiling){};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y);
__aicore__ inline void Process();
private:
constexpr static int32_t bufferNum = 2;
const RopeRegbaseTilingData* tilingData_;
TPipe* pipe_;
int64_t blockIdx_ = 0;
int64_t dSplitCoef_ = 1; // 切分系数初始化为1
uint32_t dSplitSize_ = 0;
int64_t dAlign_ = 0;
int64_t bIdx_ = 0;
int64_t sIdx_ = 0;
int64_t bNum_ = 0;
int64_t sNum_ = 0;
int64_t ubFactorS_ = 0;
int64_t ubFactorN_ = 0;
GlobalTensor<T> xGm_;
GlobalTensor<T> cosGm_;
GlobalTensor<T> sinGm_;
GlobalTensor<T> yOutGm_;
TQue<QuePosition::VECIN, bufferNum> xInQue_;
TQue<QuePosition::VECIN, bufferNum> cosInQue_;
TQue<QuePosition::VECIN, bufferNum> sinInQue_;
TQue<QuePosition::VECOUT, bufferNum> yOutQue_;
private:
__aicore__ inline void PrePareParams();
__aicore__ inline void ProcessNLoop(const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum);
__aicore__ inline void Compute(
const LocalTensor<T>& sinTensor, const LocalTensor<T>& cosTensor, const LocalTensor<T>& inTensor,
const LocalTensor<T>& outTensor, const uint32_t currSNum, const uint32_t currDNum);
__aicore__ inline void ProcessN(
const LocalTensor<T>& sinTensor, const LocalTensor<T>& cosTensor, const uint32_t bIdx, const uint32_t sIdx,
const uint32_t currSNum);
};
template <typename T>
__aicore__ inline void RotaryPositionEmbeddingBAB<T>::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y)
{
this->blockIdx_ = GetBlockIdx();
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF) ||
tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
this->dSplitCoef_ = HALF_INTERLEAVE_COEF;
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
this->dSplitCoef_ = QUARTER_MODE_COEF;
}
this->dSplitSize_ = tilingData_->sliceLength / dSplitCoef_ * sizeof(T);
this->dAlign_ = ops::CeilAlign<int64_t>(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(T)) * dSplitCoef_;
ubFactorN_ = tilingData_->ubFactorN;
ubFactorS_ = tilingData_->ubFactorS;
this->xGm_.SetGlobalBuffer((__gm__ T*)x);
this->cosGm_.SetGlobalBuffer((__gm__ T*)cos);
this->sinGm_.SetGlobalBuffer((__gm__ T*)sin);
this->yOutGm_.SetGlobalBuffer((__gm__ T*)y);
this->pipe_->InitBuffer(xInQue_, bufferNum, ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(T));
this->pipe_->InitBuffer(cosInQue_, bufferNum, ubFactorS_ * dAlign_ * sizeof(T));
this->pipe_->InitBuffer(sinInQue_, bufferNum, ubFactorS_ * dAlign_ * sizeof(T));
this->pipe_->InitBuffer(yOutQue_, bufferNum, ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(T));
}
template <typename T>
__aicore__ inline void RotaryPositionEmbeddingBAB<T>::PrePareParams()
{
bIdx_ = blockIdx_ % tilingData_->blockNumB;
sIdx_ = blockIdx_ / tilingData_->blockNumB;
bNum_ = tilingData_->blockFactorB;
sNum_ = tilingData_->blockFactorS;
if (bIdx_ == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) {
bNum_ = tilingData_->B % tilingData_->blockFactorB;
}
if (sIdx_ == tilingData_->blockNumS - 1 && tilingData_->S % tilingData_->blockFactorS != 0) {
sNum_ = tilingData_->S % tilingData_->blockFactorS;
}
}
template <typename T>
__aicore__ inline void RotaryPositionEmbeddingBAB<T>::Process()
{
PrePareParams();
uint32_t bIdxStart = bIdx_ * tilingData_->blockFactorB;
for (uint32_t bIdx = bIdxStart; bIdx < bIdxStart + bNum_; bIdx++) {
uint32_t sIdxStart = sIdx_ * tilingData_->blockFactorS;
uint32_t sLoopCnt = ops::CeilDiv(sNum_, ubFactorS_);
for (uint32_t loopIdx = 0; loopIdx < sLoopCnt; loopIdx++) {
uint32_t currSNum = (loopIdx != sLoopCnt - 1) ? ubFactorS_ : sNum_ - loopIdx * ubFactorS_;
ProcessNLoop(bIdx, sIdxStart + loopIdx * ubFactorS_, currSNum);
}
}
}
template <typename T>
__aicore__ inline void RotaryPositionEmbeddingBAB<T>::ProcessNLoop(
const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum)
{
LocalTensor<T> sinTensor = sinInQue_.AllocTensor<T>();
LocalTensor<T> cosTensor = cosInQue_.AllocTensor<T>();
int64_t offset = sIdx * tilingData_->sliceLength;
DataCopyExtParams copyParams{static_cast<uint16_t>(currSNum * dSplitCoef_), dSplitSize_, 0, 0, 0};
DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
DataCopyPad(sinTensor, sinGm_[offset], copyParams, padParams);
DataCopyPad(cosTensor, cosGm_[offset], copyParams, padParams);
sinInQue_.EnQue(sinTensor);
cosInQue_.EnQue(cosTensor);
sinTensor = sinInQue_.DeQue<T>();
cosTensor = cosInQue_.DeQue<T>();
ProcessN(sinTensor, cosTensor, bIdx, sIdx, currSNum);
sinInQue_.FreeTensor(sinTensor);
cosInQue_.FreeTensor(cosTensor);
}
template <typename T>
__aicore__ inline void RotaryPositionEmbeddingBAB<T>::ProcessN(
const LocalTensor<T>& sinTensor, const LocalTensor<T>& cosTensor, const uint32_t bIdx, const uint32_t sIdx,
const uint32_t currSNum)
{
LocalTensor<T> xTensor;
LocalTensor<T> yTensor;
int64_t baseOffset = (bIdx * tilingData_->S + sIdx) * tilingData_->N * tilingData_->D + tilingData_->sliceStart;
for (uint32_t idxN = 0; idxN < tilingData_->ubLoopNumN; idxN++) {
int64_t currDNum = (idxN == tilingData_->ubLoopNumN - 1) ? tilingData_->ubTailFactorN : ubFactorN_;
int64_t offset = baseOffset + idxN * ubFactorN_ * tilingData_->D;
xTensor = xInQue_.AllocTensor<T>();
DataCopyExtParams copyInParams{static_cast<uint16_t>(currSNum * currDNum * dSplitCoef_), dSplitSize_, static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0, 0};
DataCopyExtParams copyOutParams{static_cast<uint16_t>(currSNum * currDNum * dSplitCoef_), dSplitSize_, 0, static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0};
// deepSeekInterleave 场景，copyIn时按照D对齐，copyOut 时，按照 d/2 对齐
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
copyInParams = {static_cast<uint16_t>(currSNum * currDNum), tilingData_->sliceLength * sizeof(T), static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0, 0};
}
DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
DataCopyPad(xTensor, xGm_[offset], copyInParams, padParams);
xInQue_.EnQue(xTensor);
xTensor = xInQue_.DeQue<T>();
yTensor = yOutQue_.AllocTensor<T>();
Compute(sinTensor, cosTensor, xTensor, yTensor, currSNum, currDNum);
xInQue_.FreeTensor(xTensor);
yOutQue_.EnQue(yTensor);
yTensor = yOutQue_.DeQue<T>();
DataCopyPad(yOutGm_[offset], yTensor, copyOutParams);
yOutQue_.FreeTensor(yTensor);
}
}
template <typename T>
__aicore__ inline void RotaryPositionEmbeddingBAB<T>::Compute(
const LocalTensor<T>& sinTensor, const LocalTensor<T>& cosTensor, const LocalTensor<T>& inTensor,
const LocalTensor<T>& outTensor, const uint32_t currSNum, const uint32_t currDNum)
{
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF)) {
HalfAlignVF<T>(sinTensor, cosTensor, inTensor, outTensor, tilingData_->sliceLength, dAlign_, currSNum, currDNum);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::INTERLEAVE)) {
InterleaveModeVF<T>(sinTensor, cosTensor, inTensor, outTensor, tilingData_->sliceLength, currSNum, currDNum);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
QuarterAlignVF<T>(sinTensor, cosTensor, inTensor, outTensor, tilingData_->sliceLength, dAlign_, currSNum, currDNum);
} else {
DeepSeekInterleaveModeVF<T>(sinTensor, cosTensor, inTensor, outTensor, tilingData_->sliceLength, currSNum, currDNum);
}
}
} // namespace InplacePartialRotaryMul
#endif // ROTARY_POSITION_EMBEDDING_REG_BAB_H

207
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_bab_mixed.h Normal file
View File

@@ -0,0 +1,207 @@
/**
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
* CANN Open Software License Agreement Version 2.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
*/
/*!
* \file rotary_position_embedding_reg_bab_mixed.h
* \brief Mixed precision kernel: x is half/bfloat16, cos/sin are float
*/
#ifndef ROTARY_POSITION_EMBEDDING_REG_BAB_MIXED_H
#define ROTARY_POSITION_EMBEDDING_REG_BAB_MIXED_H
#include "apply_rotary_pos_emb_common.h"
namespace InplacePartialRotaryMul {
using namespace AscendC;
template <typename TX>
class RotaryPositionEmbeddingBABMixed {
public:
__aicore__ inline RotaryPositionEmbeddingBABMixed(TPipe *pipe, const RopeRegbaseTilingData *tiling)
: pipe_(pipe), tilingData_(tiling){};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y);
__aicore__ inline void Process();
private:
constexpr static int32_t bufferNum = 2;
const RopeRegbaseTilingData *tilingData_;
TPipe *pipe_;
int64_t blockIdx_ = 0;
int64_t dSplitCoef_ = 1;
uint32_t dSplitSize_ = 0;
int64_t dAlign_ = 0;
int64_t dAlignFloat_ = 0;
int64_t bIdx_ = 0;
int64_t sIdx_ = 0;
int64_t bNum_ = 0;
int64_t sNum_ = 0;
int64_t ubFactorS_ = 0;
int64_t ubFactorN_ = 0;
GlobalTensor<TX> xGm_;
GlobalTensor<float> cosGm_;
GlobalTensor<float> sinGm_;
GlobalTensor<TX> yOutGm_;
TQue<QuePosition::VECIN, bufferNum> xInQue_;
TQue<QuePosition::VECIN, bufferNum> cosInQue_;
TQue<QuePosition::VECIN, bufferNum> sinInQue_;
TQue<QuePosition::VECOUT, bufferNum> yOutQue_;
private:
__aicore__ inline void PrePareParams();
__aicore__ inline void ProcessNLoop(const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum);
__aicore__ inline void Compute(const LocalTensor<float> &sinTensor, const LocalTensor<float> &cosTensor,
const LocalTensor<TX> &inTensor, const LocalTensor<TX> &outTensor, const uint32_t currSNum,
const uint32_t currDNum);
__aicore__ inline void ProcessN(const LocalTensor<float> &sinTensor, const LocalTensor<float> &cosTensor,
const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum);
};
template <typename TX>
__aicore__ inline void RotaryPositionEmbeddingBABMixed<TX>::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y)
{
this->blockIdx_ = GetBlockIdx();
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF) ||
tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
this->dSplitCoef_ = HALF_INTERLEAVE_COEF;
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
this->dSplitCoef_ = QUARTER_MODE_COEF;
}
this->dSplitSize_ = tilingData_->sliceLength / dSplitCoef_ * sizeof(TX);
this->dAlign_ =
ops::CeilAlign<int64_t>(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(TX)) * dSplitCoef_;
this->dAlignFloat_ =
ops::CeilAlign<int64_t>(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(float)) * dSplitCoef_;
ubFactorN_ = tilingData_->ubFactorN;
ubFactorS_ = tilingData_->ubFactorS;
this->xGm_.SetGlobalBuffer((__gm__ TX *)x);
this->cosGm_.SetGlobalBuffer((__gm__ float *)cos);
this->sinGm_.SetGlobalBuffer((__gm__ float *)sin);
this->yOutGm_.SetGlobalBuffer((__gm__ TX *)y);
this->pipe_->InitBuffer(xInQue_, bufferNum, ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(TX));
this->pipe_->InitBuffer(cosInQue_, bufferNum, ubFactorS_ * dAlignFloat_ * sizeof(float));
this->pipe_->InitBuffer(sinInQue_, bufferNum, ubFactorS_ * dAlignFloat_ * sizeof(float));
this->pipe_->InitBuffer(yOutQue_, bufferNum, ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(TX));
}
template <typename TX>
__aicore__ inline void RotaryPositionEmbeddingBABMixed<TX>::PrePareParams()
{
bIdx_ = blockIdx_ % tilingData_->blockNumB;
sIdx_ = blockIdx_ / tilingData_->blockNumB;
bNum_ = tilingData_->blockFactorB;
sNum_ = tilingData_->blockFactorS;
if (bIdx_ == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) {
bNum_ = tilingData_->B % tilingData_->blockFactorB;
}
if (sIdx_ == tilingData_->blockNumS - 1 && tilingData_->S % tilingData_->blockFactorS != 0) {
sNum_ = tilingData_->S % tilingData_->blockFactorS;
}
}
template <typename TX>
__aicore__ inline void RotaryPositionEmbeddingBABMixed<TX>::Process()
{
PrePareParams();
uint32_t bIdxStart = bIdx_ * tilingData_->blockFactorB;
for (uint32_t bIdx = bIdxStart; bIdx < bIdxStart + bNum_; bIdx++) {
uint32_t sIdxStart = sIdx_ * tilingData_->blockFactorS;
uint32_t sLoopCnt = ops::CeilDiv(sNum_, ubFactorS_);
for (uint32_t loopIdx = 0; loopIdx < sLoopCnt; loopIdx++) {
uint32_t currSNum = (loopIdx != sLoopCnt - 1) ? ubFactorS_ : sNum_ - loopIdx * ubFactorS_;
ProcessNLoop(bIdx, sIdxStart + loopIdx * ubFactorS_, currSNum);
}
}
}
template <typename TX>
__aicore__ inline void RotaryPositionEmbeddingBABMixed<TX>::ProcessNLoop(
const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum)
{
LocalTensor<float> sinTensor = sinInQue_.AllocTensor<float>();
LocalTensor<float> cosTensor = cosInQue_.AllocTensor<float>();
int64_t offset = sIdx * tilingData_->sliceLength;
uint32_t dSplitSizeFloat = tilingData_->sliceLength / dSplitCoef_ * sizeof(float);
DataCopyExtParams copyParams{static_cast<uint16_t>(currSNum * dSplitCoef_), dSplitSizeFloat, 0, 0, 0};
DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
DataCopyPad(sinTensor, sinGm_[offset], copyParams, padParams);
DataCopyPad(cosTensor, cosGm_[offset], copyParams, padParams);
sinInQue_.EnQue(sinTensor);
cosInQue_.EnQue(cosTensor);
sinTensor = sinInQue_.DeQue<float>();
cosTensor = cosInQue_.DeQue<float>();
ProcessN(sinTensor, cosTensor, bIdx, sIdx, currSNum);
sinInQue_.FreeTensor(sinTensor);
cosInQue_.FreeTensor(cosTensor);
}
template <typename TX>
__aicore__ inline void RotaryPositionEmbeddingBABMixed<TX>::ProcessN(const LocalTensor<float> &sinTensor,
const LocalTensor<float> &cosTensor, const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum)
{
LocalTensor<TX> xTensor;
LocalTensor<TX> yTensor;
int64_t baseOffset = (bIdx * tilingData_->S + sIdx) * tilingData_->N * tilingData_->D + tilingData_->sliceStart;
for (uint32_t idxN = 0; idxN < tilingData_->ubLoopNumN; idxN++) {
int64_t currDNum = (idxN == tilingData_->ubLoopNumN - 1) ? tilingData_->ubTailFactorN : ubFactorN_;
int64_t offset = baseOffset + idxN * ubFactorN_ * tilingData_->D;
xTensor = xInQue_.AllocTensor<TX>();
DataCopyExtParams copyInParams{static_cast<uint16_t>(currSNum * currDNum * dSplitCoef_),
dSplitSize_,
static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)),
0,
0};
DataCopyExtParams copyOutParams{static_cast<uint16_t>(currSNum * currDNum * dSplitCoef_),
dSplitSize_,
0,
static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)),
0};
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) {
copyInParams = {static_cast<uint16_t>(currSNum * currDNum),
tilingData_->sliceLength * sizeof(TX),
static_cast<uint32_t>((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)),
0,
0};
}
DataCopyPadExtParams<TX> padParams{false, 0, 0, 0};
DataCopyPad(xTensor, xGm_[offset], copyInParams, padParams);
xInQue_.EnQue(xTensor);
xTensor = xInQue_.DeQue<TX>();
yTensor = yOutQue_.AllocTensor<TX>();
Compute(sinTensor, cosTensor, xTensor, yTensor, currSNum, currDNum);
xInQue_.FreeTensor(xTensor);
yOutQue_.EnQue(yTensor);
yTensor = yOutQue_.DeQue<TX>();
DataCopyPad(yOutGm_[offset], yTensor, copyOutParams);
yOutQue_.FreeTensor(yTensor);
}
}
template <typename TX>
__aicore__ inline void RotaryPositionEmbeddingBABMixed<TX>::Compute(const LocalTensor<float> &sinTensor,
const LocalTensor<float> &cosTensor, const LocalTensor<TX> &inTensor, const LocalTensor<TX> &outTensor,
const uint32_t currSNum, const uint32_t currDNum)
{
if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::INTERLEAVE)) {
InterleaveModeVFMixed<TX>(
inTensor, cosTensor, sinTensor, outTensor, tilingData_->sliceLength, currSNum, currDNum);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::HALF)) {
// For HALF mode, need to implement HalfAlignVFMixed
InterleaveModeVFMixed<TX>(
inTensor, cosTensor, sinTensor, outTensor, tilingData_->sliceLength, currSNum, currDNum);
} else if (tilingData_->rotaryMode == static_cast<int64_t>(RotaryPosEmbeddingMode::QUARTER)) {
// For QUARTER mode, need to implement QuarterAlignVFMixed
InterleaveModeVFMixed<TX>(
inTensor, cosTensor, sinTensor, outTensor, tilingData_->sliceLength, currSNum, currDNum);
}
}
} // namespace InplacePartialRotaryMul
#endif // ROTARY_POSITION_EMBEDDING_REG_BAB_MIXED_H

350
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotate_half.h Normal file
View File

@@ -0,0 +1,350 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotate_half.h
* \brief
*/
#ifndef ROTATE_HALF_H
#define ROTATE_HALF_H
#include "rotate_half_base.h"
namespace RotateHalfN {
using namespace AscendC;
template <typename T>
class RotateHalf : public RotateHalfBase<T, T> {
public:
__aicore__ inline RotateHalf(){};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RotaryPositionEmbeddingTilingData &tilingData);
__aicore__ inline void Process();
protected:
TPipe pipe;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> inQueueX;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> inQueueCos;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> inQueueSin;
TQue<QuePosition::VECOUT, DOUBLE_BUFFER> outQueueY;
GlobalTensor<T> xGm;
GlobalTensor<T> cosGm;
GlobalTensor<T> sinGm;
GlobalTensor<T> yGm;
__aicore__ inline void NormalProcess();
__aicore__ inline void RB1sdProcess();
__aicore__ inline void BndProcess();
__aicore__ inline void SingleStepProcess(uint32_t progress, uint32_t sLines, uint64_t copyLength,
uint64_t calcLength);
__aicore__ inline void RB1sdSingleStepProcess(uint32_t progress, uint32_t sLines, uint64_t xBatchStartOffset,
uint64_t rBatchStartOffset, uint64_t copyLength, uint64_t calcLength);
__aicore__ inline void Compute(LocalTensor<T> &cos, LocalTensor<T> &sin, uint32_t sLines, uint32_t calcLength);
__aicore__ inline void CopyInR(uint64_t rStartOffset, uint16_t sLines, uint32_t copyLength);
__aicore__ inline void CopyInX(uint64_t xStartOffset, uint16_t sLines, uint32_t copyLength);
__aicore__ inline void CopyOut(uint64_t yStartOffset, uint16_t sLines, uint32_t copyLength);
};
template <typename T>
__aicore__ inline void RotateHalf<T>::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RotaryPositionEmbeddingTilingData &tilingData)
{
this->BaseMemberInit(tilingData);
xGm.SetGlobalBuffer((__gm__ T *)x + this->xOffset + this->xCoreOffset * this->coreRelativeIdx, this->xAllocLength);
yGm.SetGlobalBuffer((__gm__ T *)y + this->xOffset + this->xCoreOffset * this->coreRelativeIdx, this->xAllocLength);
cosGm.SetGlobalBuffer((__gm__ T *)cos + this->rOffset + this->rCoreOffset * this->coreRelativeIdx,
this->rAllocLength);
sinGm.SetGlobalBuffer((__gm__ T *)sin + this->rOffset + this->rCoreOffset * this->coreRelativeIdx,
this->rAllocLength);
pipe.InitBuffer(inQueueX, DOUBLE_BUFFER, this->storePadDataLength * sizeof(T));
pipe.InitBuffer(outQueueY, DOUBLE_BUFFER, this->storePadDataLength * sizeof(T));
pipe.InitBuffer(inQueueCos, DOUBLE_BUFFER, this->storePadDataLength * sizeof(T));
pipe.InitBuffer(inQueueSin, DOUBLE_BUFFER, this->storePadDataLength * sizeof(T));
}
template <typename T>
__aicore__ inline void RotateHalf<T>::Process()
{
if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_BSND || this->layout == LAYOUT_SBND ||
this->layout == LAYOUT_NO_BROADCAST) {
NormalProcess();
} else if (this->layout == LAYOUT_R_B1SD) {
RB1sdProcess();
} else if (this->layout == LAYOUT_BND) {
BndProcess();
}
}
template <typename T>
__aicore__ inline void RotateHalf<T>::NormalProcess()
{
for (uint32_t progress = 0; progress < this->ubLoop; progress++) {
SingleStepProcess(progress, this->storeSLines, this->storeDataLength, this->storePadDataLength);
}
if (this->ubLast > 0) {
SingleStepProcess(this->ubLoop, this->ubLast, this->ubLastDataLength, this->ubLastPadDataLength);
}
}
template <typename T>
__aicore__ inline void RotateHalf<T>::RB1sdProcess()
{
uint64_t totalSdLength = this->totalSLines * this->dLength;
uint64_t totalNsdLength = totalSdLength * this->bcSecondDim;
uint64_t xBatchOffset, rBatchOffset;
for (uint32_t batchLoop = 0; batchLoop < this->bcFirstDim; batchLoop++) {
xBatchOffset = batchLoop * totalNsdLength;
rBatchOffset = batchLoop * totalSdLength;
for (uint32_t progress = 0; progress < this->ubLoop; progress++) {
RB1sdSingleStepProcess(progress, this->storeSLines, xBatchOffset, rBatchOffset, this->storeDataLength,
this->storePadDataLength);
}
if (this->ubLast > 0) {
RB1sdSingleStepProcess(this->ubLoop, this->ubLast, xBatchOffset, rBatchOffset, this->ubLastDataLength,
this->ubLastPadDataLength);
}
}
}
template <typename T>
__aicore__ inline void RotateHalf<T>::BndProcess()
{
CopyInR(0, 1, this->dLength);
LocalTensor<T> cosLocal = inQueueCos.DeQue<T>();
LocalTensor<T> sinLocal = inQueueSin.DeQue<T>();
Muls(sinLocal, sinLocal, (T)(-1.0), this->halfDPadLength);
uint32_t broadcastLines = this->ubLoop > 0 ? this->storeSLines - 1 : this->ubLast - 1;
if (broadcastLines > 0) {
this->RBroadCast(cosLocal, sinLocal, broadcastLines);
}
uint64_t xOffset;
for (uint32_t progress = 0; progress < this->ubLoop; progress++) {
xOffset = progress * this->storeDataLength;
CopyInX(xOffset, this->storeSLines, this->storeDataLength);
Compute(cosLocal, sinLocal, this->storeSLines, this->storePadDataLength);
CopyOut(xOffset, this->storeSLines, this->storeDataLength);
}
if (this->ubLast > 0) {
xOffset = this->ubLoop * this->storeDataLength;
CopyInX(xOffset, this->ubLast, this->ubLastDataLength);
Compute(cosLocal, sinLocal, this->ubLast, this->ubLastPadDataLength);
CopyOut(xOffset, this->ubLast, this->ubLastDataLength);
}
inQueueCos.FreeTensor(cosLocal);
inQueueSin.FreeTensor(sinLocal);
}
template <typename T>
__aicore__ inline void RotateHalf<T>::SingleStepProcess(uint32_t progress, uint32_t sLines, uint64_t copyLength,
uint64_t calcLength)
{
uint64_t xOffset, rOffset, bnLoopXStartOffset, progressOffset, batchOffset;
rOffset = progress * this->storeDataLength;
CopyInR(rOffset, sLines, copyLength);
LocalTensor<T> cosLocal = inQueueCos.DeQue<T>();
LocalTensor<T> sinLocal = inQueueSin.DeQue<T>();
this->SinCompute(sinLocal, sLines);
if (this->layout == LAYOUT_BNSD) {
uint64_t totalSdSize = this->totalSLines * this->dLength;
bnLoopXStartOffset = progress * this->storeDataLength;
for (uint32_t bnLoop = 0; bnLoop < this->bnSize; bnLoop++) {
xOffset = bnLoopXStartOffset + bnLoop * totalSdSize;
CopyInX(xOffset, sLines, copyLength);
Compute(cosLocal, sinLocal, sLines, calcLength);
CopyOut(xOffset, sLines, copyLength);
}
} else if (this->layout == LAYOUT_BSND) {
uint64_t totalSndSize = this->totalSLines * this->ndSize;
progressOffset = progress * this->bcSecondDim * this->storeDataLength;
for (uint32_t bLoop = 0; bLoop < this->bcFirstDim; bLoop++) {
batchOffset = bLoop * totalSndSize;
for (uint32_t nLoop = 0; nLoop < this->bcSecondDim; nLoop++) {
xOffset = nLoop * this->dLength + batchOffset + progressOffset;
CopyInX(xOffset, sLines, copyLength);
Compute(cosLocal, sinLocal, sLines, calcLength);
CopyOut(xOffset, sLines, copyLength);
}
}
} else if (this->layout == LAYOUT_SBND) {
bnLoopXStartOffset = progress * this->storeDataLength * this->bnSize;
for (uint32_t bnLoop = 0; bnLoop < this->bnSize; bnLoop++) {
xOffset = bnLoopXStartOffset + bnLoop * this->dLength;
CopyInX(xOffset, sLines, copyLength);
Compute(cosLocal, sinLocal, sLines, calcLength);
CopyOut(xOffset, sLines, copyLength);
}
} else if (this->layout == LAYOUT_NO_BROADCAST) {
CopyInX(rOffset, sLines, copyLength);
Compute(cosLocal, sinLocal, sLines, calcLength);
CopyOut(rOffset, sLines, copyLength);
}
inQueueCos.FreeTensor<T>(cosLocal);
inQueueSin.FreeTensor<T>(sinLocal);
}
template <typename T>
__aicore__ inline void RotateHalf<T>::RB1sdSingleStepProcess(uint32_t progress, uint32_t sLines,
uint64_t xBatchStartOffset, uint64_t rBatchStartOffset,
uint64_t copyLength, uint64_t calcLength)
{
CopyInR(progress * this->storeDataLength + rBatchStartOffset, sLines, copyLength);
LocalTensor<T> cosLocal = inQueueCos.DeQue<T>();
LocalTensor<T> sinLocal = inQueueSin.DeQue<T>();
this->SinCompute(sinLocal, sLines);
uint64_t xOffset, progressXOffset;
progressXOffset = progress * this->storeDataLength + xBatchStartOffset;
for (uint32_t nLoop = 0; nLoop < this->bcSecondDim; nLoop++) {
xOffset = nLoop * this->totalSLines * this->dLength + progressXOffset;
CopyInX(xOffset, sLines, copyLength);
Compute(cosLocal, sinLocal, sLines, calcLength);
CopyOut(xOffset, sLines, copyLength);
}
inQueueCos.FreeTensor<T>(cosLocal);
inQueueSin.FreeTensor<T>(sinLocal);
}
template <typename T>
__aicore__ inline void RotateHalf<T>::CopyInR(uint64_t rStartOffset, uint16_t sLines, uint32_t copyLength)
{
LocalTensor<T> cosLocal = inQueueCos.AllocTensor<T>();
LocalTensor<T> sinLocal = inQueueSin.AllocTensor<T>();
if (this->isAligned == true) {
DataCopy(cosLocal, cosGm[rStartOffset], copyLength);
DataCopy(sinLocal, sinGm[rStartOffset], copyLength);
} else {
DataCopyExtParams copyParams{(uint16_t)(2 * sLines), // blockCount
this->halfDBytes, // blockLen
0, // srcStride(bytes)
0, // dstStride(block)
0};
DataCopyPad(cosLocal, cosGm[rStartOffset], copyParams, this->noPadParams);
DataCopyPad(sinLocal, sinGm[rStartOffset], copyParams, this->noPadParams);
}
inQueueCos.EnQue(cosLocal);
inQueueSin.EnQue(sinLocal);
}
template <typename T>
__aicore__ inline void RotateHalf<T>::CopyInX(uint64_t xStartOffset, uint16_t sLines, uint32_t copyLength)
{
LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
DataCopyExtParams copyParams;
if (this->isAligned == true) {
if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND ||
this->layout == LAYOUT_R_B1SD) {
DataCopy(xLocal, xGm[xStartOffset], copyLength);
} else if (this->layout == LAYOUT_BSND) {
copyParams.blockCount = sLines;
copyParams.blockLen = this->dBytes;
copyParams.srcStride = (this->bcSecondDim - 1) * this->dBytes;
copyParams.dstStride = 0;
DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams);
} else if (this->layout == LAYOUT_SBND) {
copyParams.blockCount = sLines;
copyParams.blockLen = this->dBytes;
copyParams.srcStride = (this->bnSize - 1) * this->dBytes;
copyParams.dstStride = 0;
DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams);
}
} else {
if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND ||
this->layout == LAYOUT_R_B1SD) {
copyParams.blockCount = (uint16_t)(2 * sLines);
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = 0;
copyParams.dstStride = 0;
DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams);
} else if (this->layout == LAYOUT_BSND) {
copyParams.blockCount = sLines;
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = (2 * this->bcSecondDim - 1) * this->halfDBytes;
copyParams.dstStride = this->halfDPadBlocks;
DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams);
DataCopyPad(xLocal[this->halfDPadLength], xGm[xStartOffset + this->halfDLength], copyParams,
this->noPadParams);
} else if (this->layout == LAYOUT_SBND) {
copyParams.blockCount = sLines;
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = (2 * this->bnSize - 1) * this->halfDBytes;
copyParams.dstStride = this->halfDPadBlocks;
DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams);
DataCopyPad(xLocal[this->halfDPadLength], xGm[xStartOffset + this->halfDLength], copyParams,
this->noPadParams);
}
}
inQueueX.EnQue(xLocal);
}
template <typename T>
__aicore__ inline void RotateHalf<T>::CopyOut(uint64_t yStartOffset, uint16_t sLines, uint32_t copyLength)
{
LocalTensor<T> yLocal = outQueueY.DeQue<T>();
DataCopyExtParams copyParams;
if (this->isAligned == true) {
copyParams.blockCount = sLines;
copyParams.blockLen = this->dBytes;
copyParams.srcStride = 0;
if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND ||
this->layout == LAYOUT_R_B1SD) {
DataCopy(yGm[yStartOffset], yLocal, copyLength);
} else if (this->layout == LAYOUT_BSND) {
copyParams.dstStride = (this->bcSecondDim - 1) * this->dBytes;
DataCopyPad(yGm[yStartOffset], yLocal, copyParams);
} else if (this->layout == LAYOUT_SBND) {
copyParams.dstStride = (this->bnSize - 1) * this->dBytes;
DataCopyPad(yGm[yStartOffset], yLocal, copyParams);
}
} else {
if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND ||
this->layout == LAYOUT_R_B1SD) {
copyParams.blockCount = (uint16_t)(2 * sLines);
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = 0;
copyParams.dstStride = 0;
DataCopyPad(yGm[yStartOffset], yLocal, copyParams);
} else if (this->layout == LAYOUT_BSND) {
copyParams.blockCount = sLines;
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = this->halfDPadBlocks;
copyParams.dstStride = (2 * this->bcSecondDim - 1) * this->halfDBytes;
DataCopyPad(yGm[yStartOffset], yLocal, copyParams);
DataCopyPad(yGm[yStartOffset + this->halfDLength], yLocal[this->halfDPadLength], copyParams);
} else if (this->layout == LAYOUT_SBND) {
copyParams.blockCount = sLines;
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = this->halfDPadBlocks;
copyParams.dstStride = (2 * this->bnSize - 1) * this->halfDBytes;
DataCopyPad(yGm[yStartOffset], yLocal, copyParams);
DataCopyPad(yGm[yStartOffset + this->halfDLength], yLocal[this->halfDPadLength], copyParams);
}
}
outQueueY.FreeTensor(yLocal);
}
template <typename T>
__aicore__ inline void RotateHalf<T>::Compute(LocalTensor<T> &cos, LocalTensor<T> &sin, uint32_t sLines,
uint32_t calcLength)
{
LocalTensor<T> xLocal = inQueueX.DeQue<T>();
LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();
this->XNewCopy(xLocal, yLocal, sLines);
this->ComputeInner(xLocal, yLocal, cos, sin, calcLength);
outQueueY.EnQue(yLocal);
inQueueX.FreeTensor<T>(xLocal);
}
} // namespace RotateHalfN
#endif // ROTATE_HALF_H

283
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotate_half_base.h Normal file
View File

@@ -0,0 +1,283 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotate_half_base.h
* \brief
*/
#ifndef ROTATE_HALF_BASE_H
#define ROTATE_HALF_BASE_H
#include "kernel_operator.h"
namespace RotateHalfN {
using namespace AscendC;
constexpr uint8_t REPEAT_MAX = 255;
constexpr int32_t SINGLE_BUFFER = 1;
constexpr int32_t DOUBLE_BUFFER = 2;
constexpr uint32_t BYTE_OF_BLOCK = 32;
constexpr uint32_t BYTE_OF_REPEAT = 256;
constexpr uint16_t LAYOUT_BNSD = 1;
constexpr uint16_t LAYOUT_BSND = 2;
constexpr uint16_t LAYOUT_SBND = 3;
constexpr uint16_t LAYOUT_NO_BROADCAST = 4;
constexpr uint16_t LAYOUT_BND = 5;
constexpr uint16_t LAYOUT_R_B1SD = 6;
template <typename OriT, typename CmpT>
class RotateHalfBase {
public:
__aicore__ inline RotateHalfBase(){};
__aicore__ inline void BaseMemberInit(const RotaryPositionEmbeddingTilingData &tilingData);
protected:
__aicore__ inline void GetTilingData(const RotateHalfParams &tiling);
__aicore__ inline void SinCompute(LocalTensor<CmpT> &sin, uint32_t sLines);
__aicore__ inline void ComputeInner(LocalTensor<CmpT> &x, LocalTensor<CmpT> &xNew, LocalTensor<CmpT> &cos,
LocalTensor<CmpT> &sin, uint32_t calcLength);
__aicore__ inline void RBroadCast(LocalTensor<CmpT> &cos, LocalTensor<CmpT> &sin, uint32_t broadcastLines);
__aicore__ inline void XNewCopy(LocalTensor<CmpT> &x, LocalTensor<CmpT> &xNew, uint16_t sLines);
bool isAligned;
uint16_t layout;
uint64_t gmLength;
uint64_t bcFirstDim;
uint64_t bcSecondDim;
uint64_t dLength;
uint64_t dPadLength;
uint64_t halfDLength;
uint64_t halfDPadLength;
uint64_t totalSLines;
uint64_t storeSLines;
uint64_t storeDataLength;
uint64_t storePadDataLength;
uint64_t ubLoop;
uint64_t ubLast;
uint64_t formerCoreNum;
uint64_t tailCoreNum;
uint64_t formerSLines;
uint64_t tailSLines;
uint64_t coreSLines;
uint64_t xDataLength;
uint64_t rDataLength;
uint64_t ubLastDataLength;
uint64_t ubLastPadDataLength;
uint64_t xOffset;
uint64_t rOffset;
uint8_t repeatStride;
uint32_t dataEachRepeat;
uint32_t bnSize;
uint32_t ndSize;
uint32_t bndSize;
uint32_t dBytes;
uint32_t halfDBytes;
uint32_t halfDPadBlocks;
uint32_t innerHalfLoop;
uint32_t innerHalfLast;
uint64_t xAllocLength;
uint64_t rAllocLength;
uint64_t xCoreOffset;
uint64_t rCoreOffset;
uint64_t coreRelativeIdx;
DataCopyPadExtParams<OriT> noPadParams{false, 0, 0, 0};
};
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBase<OriT, CmpT>::BaseMemberInit(const RotaryPositionEmbeddingTilingData &tilingData)
{
const RotateHalfParams &tiling = tilingData.rotateHalfParams;
GetTilingData(tiling);
// intermediate variable
repeatStride = dPadLength * sizeof(CmpT) / BYTE_OF_BLOCK;
dataEachRepeat = BYTE_OF_REPEAT / sizeof(CmpT);
innerHalfLoop = halfDLength / dataEachRepeat;
innerHalfLast = halfDLength % dataEachRepeat;
bnSize = bcFirstDim * bcSecondDim;
ndSize = bcSecondDim * dLength;
bndSize = bnSize * dLength;
dBytes = dLength * sizeof(OriT);
halfDBytes = halfDLength * sizeof(OriT);
halfDPadBlocks = halfDPadLength * sizeof(OriT) / BYTE_OF_BLOCK;
// gm and ub space params
xAllocLength = gmLength;
rAllocLength = rDataLength;
rCoreOffset = rDataLength;
if (layout == LAYOUT_BNSD || layout == LAYOUT_NO_BROADCAST) {
xCoreOffset = coreSLines * dLength;
} else if (layout == LAYOUT_BSND) {
xCoreOffset = coreSLines * ndSize;
} else if (layout == LAYOUT_SBND) {
xCoreOffset = coreSLines * bndSize;
} else if (layout == LAYOUT_BND) {
xCoreOffset = xDataLength;
xAllocLength = xDataLength;
rCoreOffset = 0;
rOffset = 0;
} else if (layout == LAYOUT_R_B1SD) {
xCoreOffset = coreSLines * dLength;
rCoreOffset = xCoreOffset;
rAllocLength = bcFirstDim * totalSLines * dLength;
}
}
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBase<OriT, CmpT>::GetTilingData(const RotateHalfParams &tiling)
{
isAligned = tiling.isAligned == 1;
layout = tiling.tilingMode;
gmLength = tiling.gmLength;
bcFirstDim = tiling.broadcastFirstDim;
bcSecondDim = tiling.broadcastSecondDim;
dLength = tiling.dLength;
dPadLength = tiling.dPadLength;
halfDLength = tiling.halfDLength;
halfDPadLength = tiling.halfDPadLength;
totalSLines = tiling.totalSLines;
storeSLines = tiling.storeSLines;
storeDataLength = tiling.storeDataLength;
storePadDataLength = tiling.storePadDataLength;
formerCoreNum = tiling.formerCoreNum;
tailCoreNum = tiling.tailCoreNum;
formerSLines = tiling.formerSLines;
tailSLines = tiling.tailSLines;
if (GetBlockIdx() < formerCoreNum) {
coreRelativeIdx = GetBlockIdx();
coreSLines = formerSLines;
ubLoop = tiling.formerUbLoop;
ubLast = tiling.formerUbLast;
xDataLength = tiling.formerXDataLength;
rDataLength = tiling.formerRDataLength;
ubLastDataLength = tiling.formerUbLastDataLength;
ubLastPadDataLength = tiling.formerUbLastPadDataLength;
xOffset = 0;
rOffset = 0;
} else {
coreRelativeIdx = GetBlockIdx() - formerCoreNum;
coreSLines = tailSLines;
ubLoop = tiling.tailUbLoop;
ubLast = tiling.tailUbLast;
xDataLength = tiling.tailXDataLength;
rDataLength = tiling.tailRDataLength;
ubLastDataLength = tiling.tailUbLastDataLength;
ubLastPadDataLength = tiling.tailUbLastPadDataLength;
xOffset = tiling.formerXCoreOffset;
rOffset = tiling.formerRCoreOffset;
}
}
/* sin_l = -1 * sin_l */
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBase<OriT, CmpT>::SinCompute(LocalTensor<CmpT> &sin, uint32_t sLines)
{
uint32_t repeatOffset, innerOffset;
uint32_t repeatLoop = sLines / REPEAT_MAX;
uint8_t repeatLast = sLines % REPEAT_MAX;
UnaryRepeatParams repeatParams{1, 1, repeatStride, repeatStride};
for (uint32_t i = 0; i < repeatLoop; i++) {
repeatOffset = i * REPEAT_MAX * dPadLength;
for (uint32_t j = 0; j < innerHalfLoop; j++) {
innerOffset = j * dataEachRepeat + repeatOffset;
Muls(sin[innerOffset], sin[innerOffset], (CmpT)(-1.0), dataEachRepeat, REPEAT_MAX, repeatParams);
}
if (innerHalfLast > 0) {
innerOffset = innerHalfLoop * dataEachRepeat + repeatOffset;
Muls(sin[innerOffset], sin[innerOffset], (CmpT)(-1.0), innerHalfLast, REPEAT_MAX, repeatParams);
}
}
if (repeatLast > 0) {
repeatOffset = repeatLoop * REPEAT_MAX * dPadLength;
for (uint32_t j = 0; j < innerHalfLoop; j++) {
innerOffset = j * dataEachRepeat + repeatOffset;
Muls(sin[innerOffset], sin[innerOffset], (CmpT)(-1.0), dataEachRepeat, repeatLast, repeatParams);
}
if (innerHalfLast > 0) {
innerOffset = innerHalfLoop * dataEachRepeat + repeatOffset;
Muls(sin[innerOffset], sin[innerOffset], (CmpT)(-1.0), innerHalfLast, repeatLast, repeatParams);
}
}
}
/* x = x * cos, xNew = xNew * sin, y = x + xNew */
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBase<OriT, CmpT>::ComputeInner(LocalTensor<CmpT> &x, LocalTensor<CmpT> &xNew,
LocalTensor<CmpT> &cos, LocalTensor<CmpT> &sin,
uint32_t calcLength)
{
Mul(x, x, cos, calcLength);
Mul(xNew, xNew, sin, calcLength);
Add(xNew, xNew, x, calcLength);
}
/* broadcast cos, sin from (1, D) to (storeSLines, D) or (ubLast, D) shape */
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBase<OriT, CmpT>::RBroadCast(LocalTensor<CmpT> &cos, LocalTensor<CmpT> &sin,
uint32_t broadcastLines)
{
uint32_t repeatOffset, copySrcOffset, copyDstOffset;
uint32_t innerLoop = dPadLength / dataEachRepeat;
uint32_t innerLast = dPadLength % dataEachRepeat;
uint32_t repeatLoop = broadcastLines / REPEAT_MAX;
uint8_t repeatLast = broadcastLines % REPEAT_MAX;
CopyRepeatParams repParams{1, 1, repeatStride, 0};
for (uint32_t i = 0; i < repeatLoop; i++) {
repeatOffset = i * REPEAT_MAX * dPadLength;
for (uint32_t j = 0; j < innerLoop; j++) {
copySrcOffset = j * dataEachRepeat;
copyDstOffset = copySrcOffset + repeatOffset + dPadLength;
Copy(cos[copyDstOffset], cos[copySrcOffset], dataEachRepeat, REPEAT_MAX, repParams);
Copy(sin[copyDstOffset], sin[copySrcOffset], dataEachRepeat, REPEAT_MAX, repParams);
}
if (innerLast > 0) {
copySrcOffset = innerLoop * dataEachRepeat;
copyDstOffset = copySrcOffset + repeatOffset + dPadLength;
Copy(cos[copyDstOffset], cos[copySrcOffset], innerLast, REPEAT_MAX, repParams);
Copy(sin[copyDstOffset], sin[copySrcOffset], innerLast, REPEAT_MAX, repParams);
}
}
if (repeatLast > 0) {
repeatOffset = repeatLoop * REPEAT_MAX * dPadLength;
for (uint32_t j = 0; j < innerLoop; j++) {
copySrcOffset = j * dataEachRepeat;
copyDstOffset = copySrcOffset + repeatOffset + dPadLength;
Copy(cos[copyDstOffset], cos[copySrcOffset], dataEachRepeat, repeatLast, repParams);
Copy(sin[copyDstOffset], sin[copySrcOffset], dataEachRepeat, repeatLast, repParams);
}
if (innerLast > 0) {
copySrcOffset = innerLoop * dataEachRepeat;
copyDstOffset = copySrcOffset + repeatOffset + dPadLength;
Copy(cos[copyDstOffset], cos[copySrcOffset], innerLast, repeatLast, repParams);
Copy(sin[copyDstOffset], sin[copySrcOffset], innerLast, repeatLast, repParams);
}
}
}
/* copy x to xNew: x_l --> xNew_r, x_r --> xNew_l */
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBase<OriT, CmpT>::XNewCopy(LocalTensor<CmpT> &x, LocalTensor<CmpT> &xNew,
uint16_t sLines)
{
uint16_t stride = this->halfDPadLength * sizeof(CmpT) / BYTE_OF_BLOCK;
DataCopyParams copyParams{sLines, stride, stride, stride};
DataCopy(xNew, x[this->halfDPadLength], copyParams);
DataCopy(xNew[this->halfDPadLength], x, copyParams);
}
} // namespace RotateHalfN
#endif // ROTATE_HALF_BASE_H

385
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotate_half_bf16.h Normal file
View File

@@ -0,0 +1,385 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotate_half_bf16.h
* \brief
*/
#ifndef ROTATE_HALF_BF16_H
#define ROTATE_HALF_BF16_H
#include "rotate_half_base.h"
namespace RotateHalfN {
using namespace AscendC;
template <typename OriT, typename CmpT>
class RotateHalfBf16 : public RotateHalfBase<OriT, CmpT> {
public:
__aicore__ inline RotateHalfBf16(){};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RotaryPositionEmbeddingTilingData &tilingData);
__aicore__ inline void Process();
protected:
TPipe pipe;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> inQueueX;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> inQueueCos;
TQue<QuePosition::VECIN, DOUBLE_BUFFER> inQueueSin;
TQue<QuePosition::VECOUT, DOUBLE_BUFFER> outQueueY;
TBuf<TPosition::VECCALC> xBuf;
TBuf<TPosition::VECCALC> xNewBuf;
TBuf<TPosition::VECCALC> cosBuf;
TBuf<TPosition::VECCALC> sinBuf;
GlobalTensor<OriT> xGm;
GlobalTensor<OriT> cosGm;
GlobalTensor<OriT> sinGm;
GlobalTensor<OriT> yGm;
__aicore__ inline void NormalProcessBf16();
__aicore__ inline void RB1sdProcessBf16();
__aicore__ inline void BndProcessBf16();
__aicore__ inline void SingleStepProcess(uint32_t progress, uint32_t sLines, uint64_t copyLength,
uint64_t calcLength);
__aicore__ inline void RB1sdSingleStepProcessBf16(uint32_t progress, uint64_t sLines, uint64_t xBatchStartOffset,
uint64_t rBatchStartOffset, uint64_t copyLength,
uint64_t calcLength);
__aicore__ inline void Compute(LocalTensor<CmpT> &cos, LocalTensor<CmpT> &sin, uint32_t sLines,
uint32_t calcLength);
__aicore__ inline void CopyInR(uint64_t rStartOffset, uint16_t sLines, uint32_t copyLength);
__aicore__ inline void CopyInX(uint64_t xStartOffset, uint16_t sLines, uint32_t copyLength);
__aicore__ inline void CopyOut(uint64_t yOffset, uint16_t sLines, uint32_t copyLength);
};
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBf16<OriT, CmpT>::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RotaryPositionEmbeddingTilingData &tilingData)
{
this->BaseMemberInit(tilingData);
xGm.SetGlobalBuffer((__gm__ OriT *)x + this->xOffset + this->xCoreOffset * this->coreRelativeIdx,
this->xAllocLength);
yGm.SetGlobalBuffer((__gm__ OriT *)y + this->xOffset + this->xCoreOffset * this->coreRelativeIdx,
this->xAllocLength);
cosGm.SetGlobalBuffer((__gm__ OriT *)cos + this->rOffset + this->rCoreOffset * this->coreRelativeIdx,
this->rAllocLength);
sinGm.SetGlobalBuffer((__gm__ OriT *)sin + this->rOffset + this->rCoreOffset * this->coreRelativeIdx,
this->rAllocLength);
pipe.InitBuffer(inQueueX, DOUBLE_BUFFER, this->storePadDataLength * sizeof(OriT));
pipe.InitBuffer(outQueueY, DOUBLE_BUFFER, this->storePadDataLength * sizeof(OriT));
pipe.InitBuffer(inQueueCos, DOUBLE_BUFFER, this->storePadDataLength * sizeof(OriT));
pipe.InitBuffer(inQueueSin, DOUBLE_BUFFER, this->storePadDataLength * sizeof(OriT));
pipe.InitBuffer(xBuf, this->storePadDataLength * sizeof(CmpT));
pipe.InitBuffer(cosBuf, this->storePadDataLength * sizeof(CmpT));
pipe.InitBuffer(sinBuf, this->storePadDataLength * sizeof(CmpT));
pipe.InitBuffer(xNewBuf, this->storePadDataLength * sizeof(CmpT));
}
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBf16<OriT, CmpT>::Process()
{
if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_BSND || this->layout == LAYOUT_SBND ||
this->layout == LAYOUT_NO_BROADCAST) {
NormalProcessBf16();
} else if (this->layout == LAYOUT_R_B1SD) {
RB1sdProcessBf16();
} else if (this->layout == LAYOUT_BND) {
BndProcessBf16();
}
}
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBf16<OriT, CmpT>::NormalProcessBf16()
{
for (uint32_t progress = 0; progress < this->ubLoop; progress++) {
SingleStepProcess(progress, this->storeSLines, this->storeDataLength, this->storePadDataLength);
}
if (this->ubLast > 0) {
SingleStepProcess(this->ubLoop, this->ubLast, this->ubLastDataLength, this->ubLastPadDataLength);
}
}
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBf16<OriT, CmpT>::RB1sdProcessBf16()
{
uint64_t totalSdLength = this->totalSLines * this->dLength;
uint64_t totalNsdLength = totalSdLength * this->bcSecondDim;
uint64_t xBatchOffset, rBatchOffset;
for (uint32_t loop = 0; loop < this->bcFirstDim; loop++) {
xBatchOffset = loop * totalNsdLength;
rBatchOffset = loop * totalSdLength;
for (uint32_t progress = 0; progress < this->ubLoop; progress++) {
RB1sdSingleStepProcessBf16(progress, this->storeSLines, xBatchOffset, rBatchOffset, this->storeDataLength,
this->storePadDataLength);
}
if (this->ubLast > 0) {
RB1sdSingleStepProcessBf16(this->ubLoop, this->ubLast, xBatchOffset, rBatchOffset, this->ubLastDataLength,
this->ubLastPadDataLength);
}
}
}
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBf16<OriT, CmpT>::BndProcessBf16()
{
CopyInR(0, 1, this->dLength);
LocalTensor<OriT> cosLocal = inQueueCos.DeQue<OriT>();
LocalTensor<OriT> sinLocal = inQueueSin.DeQue<OriT>();
LocalTensor<CmpT> cosFp32 = cosBuf.Get<CmpT>();
LocalTensor<CmpT> sinFp32 = sinBuf.Get<CmpT>();
Cast(cosFp32, cosLocal, RoundMode::CAST_NONE, this->dPadLength);
Cast(sinFp32, sinLocal, RoundMode::CAST_NONE, this->dPadLength);
inQueueCos.FreeTensor<OriT>(cosLocal);
inQueueSin.FreeTensor<OriT>(sinLocal);
Muls(sinFp32, sinFp32, (CmpT)(-1.0), this->halfDPadLength);
uint32_t broadcastLines = this->ubLoop > 0 ? this->storeSLines - 1 : this->ubLast - 1;
if (broadcastLines > 0) {
this->RBroadCast(cosFp32, sinFp32, broadcastLines);
}
uint64_t xOffset;
for (uint32_t progress = 0; progress < this->ubLoop; progress++) {
xOffset = progress * this->storeDataLength;
CopyInX(xOffset, this->storeSLines, this->storeDataLength);
Compute(cosFp32, sinFp32, this->storeSLines, this->storePadDataLength);
CopyOut(xOffset, this->storeSLines, this->storeDataLength);
}
if (this->ubLast > 0) {
xOffset = this->ubLoop * this->storeDataLength;
CopyInX(xOffset, this->ubLast, this->ubLastDataLength);
Compute(cosFp32, sinFp32, this->ubLast, this->ubLastPadDataLength);
CopyOut(xOffset, this->ubLast, this->ubLastDataLength);
}
}
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBf16<OriT, CmpT>::SingleStepProcess(uint32_t progress, uint32_t sLines,
uint64_t copyLength, uint64_t calcLength)
{
uint64_t xOffset, rOffset, bnLoopXStartOffset, progressOffset, batchOffset;
rOffset = progress * this->storeDataLength;
CopyInR(rOffset, sLines, copyLength);
LocalTensor<OriT> cosLocal = inQueueCos.DeQue<OriT>();
LocalTensor<OriT> sinLocal = inQueueSin.DeQue<OriT>();
LocalTensor<CmpT> cosFp32 = cosBuf.Get<CmpT>();
LocalTensor<CmpT> sinFp32 = sinBuf.Get<CmpT>();
Cast(cosFp32, cosLocal, RoundMode::CAST_NONE, calcLength);
Cast(sinFp32, sinLocal, RoundMode::CAST_NONE, calcLength);
inQueueCos.FreeTensor<OriT>(cosLocal);
inQueueSin.FreeTensor<OriT>(sinLocal);
this->SinCompute(sinFp32, sLines);
if (this->layout == LAYOUT_BNSD) {
uint64_t totalSdSize = this->totalSLines * this->dLength;
bnLoopXStartOffset = progress * this->storeDataLength;
for (uint32_t bnLoop = 0; bnLoop < this->bnSize; bnLoop++) {
xOffset = bnLoopXStartOffset + bnLoop * totalSdSize;
CopyInX(xOffset, sLines, copyLength);
Compute(cosFp32, sinFp32, sLines, calcLength);
CopyOut(xOffset, sLines, copyLength);
}
} else if (this->layout == LAYOUT_BSND) {
uint64_t totalSndSize = this->totalSLines * this->ndSize;
progressOffset = progress * this->bcSecondDim * this->storeDataLength;
for (uint32_t batchLoop = 0; batchLoop < this->bcFirstDim; batchLoop++) {
batchOffset = batchLoop * totalSndSize;
for (uint32_t nLoop = 0; nLoop < this->bcSecondDim; nLoop++) {
xOffset = nLoop * this->dLength + batchOffset + progressOffset;
CopyInX(xOffset, sLines, copyLength);
Compute(cosFp32, sinFp32, sLines, calcLength);
CopyOut(xOffset, sLines, copyLength);
}
}
} else if (this->layout == LAYOUT_SBND) {
bnLoopXStartOffset = progress * this->storeDataLength * this->bnSize;
for (uint32_t bnLoop = 0; bnLoop < this->bnSize; bnLoop++) {
xOffset = bnLoopXStartOffset + bnLoop * this->dLength;
CopyInX(xOffset, sLines, copyLength);
Compute(cosFp32, sinFp32, sLines, calcLength);
CopyOut(xOffset, sLines, copyLength);
}
} else if (this->layout == LAYOUT_NO_BROADCAST) {
CopyInX(rOffset, sLines, copyLength);
Compute(cosFp32, sinFp32, sLines, calcLength);
CopyOut(rOffset, sLines, copyLength);
}
}
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBf16<OriT, CmpT>::RB1sdSingleStepProcessBf16(uint32_t progress, uint64_t sLines,
uint64_t xBatchStartOffset,
uint64_t rBatchStartOffset,
uint64_t copyLength, uint64_t calcLength)
{
CopyInR(progress * this->storeDataLength + rBatchStartOffset, sLines, copyLength);
LocalTensor<OriT> cosLocal = inQueueCos.DeQue<OriT>();
LocalTensor<OriT> sinLocal = inQueueSin.DeQue<OriT>();
LocalTensor<CmpT> cosFp32 = cosBuf.Get<CmpT>();
LocalTensor<CmpT> sinFp32 = sinBuf.Get<CmpT>();
Cast(cosFp32, cosLocal, RoundMode::CAST_NONE, calcLength);
Cast(sinFp32, sinLocal, RoundMode::CAST_NONE, calcLength);
inQueueCos.FreeTensor<OriT>(cosLocal);
inQueueSin.FreeTensor<OriT>(sinLocal);
this->SinCompute(sinFp32, sLines);
uint64_t xOffset, progressXOffset;
progressXOffset = progress * this->storeDataLength + xBatchStartOffset;
for (uint32_t nLoop = 0; nLoop < this->bcSecondDim; nLoop++) {
xOffset = nLoop * this->totalSLines * this->dLength + progressXOffset;
CopyInX(xOffset, sLines, copyLength);
Compute(cosFp32, sinFp32, sLines, calcLength);
CopyOut(xOffset, sLines, copyLength);
}
}
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBf16<OriT, CmpT>::CopyInR(uint64_t rStartOffset, uint16_t sLines, uint32_t copyLength)
{
LocalTensor<OriT> sinLocal = inQueueSin.AllocTensor<OriT>();
LocalTensor<OriT> cosLocal = inQueueCos.AllocTensor<OriT>();
if (this->isAligned == true) {
DataCopy(sinLocal, sinGm[rStartOffset], copyLength);
DataCopy(cosLocal, cosGm[rStartOffset], copyLength);
} else {
DataCopyExtParams copyParams{(uint16_t)(2 * sLines), // blockCount
this->halfDBytes, // blockLen
0, // srcStride(bytes)
0, // dstStride(block)
0};
DataCopyPad(sinLocal, sinGm[rStartOffset], copyParams, this->noPadParams);
DataCopyPad(cosLocal, cosGm[rStartOffset], copyParams, this->noPadParams);
}
inQueueSin.EnQue(sinLocal);
inQueueCos.EnQue(cosLocal);
}
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBf16<OriT, CmpT>::CopyInX(uint64_t xStartOffset, uint16_t storeSLines,
uint32_t copyLength)
{
LocalTensor<OriT> xLocal = inQueueX.AllocTensor<OriT>();
DataCopyExtParams copyParams;
if (this->isAligned == true) {
if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND ||
this->layout == LAYOUT_R_B1SD) {
DataCopy(xLocal, xGm[xStartOffset], copyLength);
} else if (this->layout == LAYOUT_BSND) {
copyParams.blockCount = storeSLines;
copyParams.blockLen = this->dBytes;
copyParams.srcStride = (this->bcSecondDim - 1) * this->dBytes;
copyParams.dstStride = 0;
DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams);
} else if (this->layout == LAYOUT_SBND) {
copyParams.blockCount = storeSLines;
copyParams.blockLen = this->dBytes;
copyParams.srcStride = (this->bnSize - 1) * this->dBytes;
copyParams.dstStride = 0;
DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams);
}
} else {
if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND ||
this->layout == LAYOUT_R_B1SD) {
copyParams.blockCount = (uint16_t)(2 * storeSLines);
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = 0;
copyParams.dstStride = 0;
DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams);
} else if (this->layout == LAYOUT_BSND) {
copyParams.blockCount = storeSLines;
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = (2 * this->bcSecondDim - 1) * this->halfDBytes;
copyParams.dstStride = this->halfDPadBlocks;
DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams);
DataCopyPad(xLocal[this->halfDPadLength], xGm[xStartOffset + this->halfDLength], copyParams,
this->noPadParams);
} else if (this->layout == LAYOUT_SBND) {
copyParams.blockCount = storeSLines;
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = (this->bnSize * 2 - 1) * this->halfDBytes;
copyParams.dstStride = this->halfDPadBlocks;
DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams);
DataCopyPad(xLocal[this->halfDPadLength], xGm[xStartOffset + this->halfDLength], copyParams,
this->noPadParams);
}
}
inQueueX.EnQue(xLocal);
}
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBf16<OriT, CmpT>::CopyOut(uint64_t yOffset, uint16_t storeSLines, uint32_t copyLength)
{
LocalTensor<OriT> yLocal = outQueueY.DeQue<OriT>();
DataCopyExtParams copyParams;
if (this->isAligned == true) {
copyParams.blockCount = storeSLines;
copyParams.blockLen = this->dBytes;
copyParams.srcStride = 0;
if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND ||
this->layout == LAYOUT_R_B1SD) {
DataCopy(yGm[yOffset], yLocal, copyLength);
} else if (this->layout == LAYOUT_BSND) {
copyParams.dstStride = (this->bcSecondDim - 1) * this->dBytes;
DataCopyPad(yGm[yOffset], yLocal, copyParams);
} else if (this->layout == LAYOUT_SBND) {
copyParams.dstStride = (this->bnSize - 1) * this->dBytes;
DataCopyPad(yGm[yOffset], yLocal, copyParams);
}
} else {
if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND ||
this->layout == LAYOUT_R_B1SD) {
copyParams.blockCount = (uint16_t)(2 * storeSLines);
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = 0;
copyParams.dstStride = 0;
DataCopyPad(yGm[yOffset], yLocal, copyParams);
} else if (this->layout == LAYOUT_BSND) {
copyParams.blockCount = storeSLines;
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = this->halfDPadBlocks;
copyParams.dstStride = (2 * this->bcSecondDim - 1) * this->halfDBytes;
DataCopyPad(yGm[yOffset], yLocal, copyParams);
DataCopyPad(yGm[yOffset + this->halfDLength], yLocal[this->halfDPadLength], copyParams);
} else if (this->layout == LAYOUT_SBND) {
copyParams.blockCount = storeSLines;
copyParams.blockLen = this->halfDBytes;
copyParams.srcStride = this->halfDPadBlocks;
copyParams.dstStride = (2 * this->bnSize - 1) * this->halfDBytes;
DataCopyPad(yGm[yOffset], yLocal, copyParams);
DataCopyPad(yGm[yOffset + this->halfDLength], yLocal[this->halfDPadLength], copyParams);
}
}
outQueueY.FreeTensor(yLocal);
}
template <typename OriT, typename CmpT>
__aicore__ inline void RotateHalfBf16<OriT, CmpT>::Compute(LocalTensor<CmpT> &cos, LocalTensor<CmpT> &sin,
uint32_t sLines, uint32_t calcLength)
{
LocalTensor<OriT> xLocal = inQueueX.DeQue<OriT>();
LocalTensor<OriT> yLocal = outQueueY.AllocTensor<OriT>();
LocalTensor<CmpT> xFp32 = xBuf.Get<CmpT>();
LocalTensor<CmpT> xNewFp32 = xNewBuf.Get<CmpT>();
Cast(xFp32, xLocal, RoundMode::CAST_NONE, calcLength);
inQueueX.FreeTensor(xLocal);
this->XNewCopy(xFp32, xNewFp32, sLines);
this->ComputeInner(xFp32, xNewFp32, cos, sin, calcLength);
Cast(yLocal, xNewFp32, RoundMode::CAST_RINT, calcLength);
outQueueY.EnQue(yLocal);
}
} // namespace RotateHalfN
#endif // ROTATE_HALF_BF16_H

145
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_common.h Normal file
View File

@@ -0,0 +1,145 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotate_interleaved_common.h
* \brief
*/
#ifndef ROTATE_INTERLEAVED_COMMON_H
#define ROTATE_INTERLEAVED_COMMON_H
#include "kernel_operator.h"
#include "impl/dav_c220/kernel_operator_reg_others_impl.h"
namespace RotateInterleavedN {
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 1;
constexpr int32_t NUM_8 = 8;
constexpr uint8_t REPEAT_MAX = 255;
constexpr int32_t BLOCK_SIZE = 32;
constexpr int32_t MASK_FP16 = 128;
constexpr int32_t MASK_FP32 = 64;
constexpr int32_t ALIGN_16 = 16;
constexpr int32_t ALIGN_32 = 8;
// SD -> BSND
template <typename T>
__aicore__ inline void BroadCastTriToBsnd(LocalTensor<T> &tri, uint32_t batchSize, uint32_t calcLen, uint32_t numHeads,
uint32_t headDimAlign)
{
DataCopyParams intriParams;
intriParams.blockCount = static_cast<uint16_t>(calcLen);
intriParams.blockLen = static_cast<uint16_t>(headDimAlign * sizeof(T) / BLOCK_SIZE);
intriParams.srcStride = static_cast<uint16_t>((numHeads - 1) * headDimAlign * sizeof(T) / BLOCK_SIZE);
intriParams.dstStride = static_cast<uint16_t>((numHeads - 1) * headDimAlign * sizeof(T) / BLOCK_SIZE);
// SD -> SND
for (uint32_t numHeadsIdx = 1; numHeadsIdx < numHeads; ++numHeadsIdx) {
DataCopy(tri[numHeadsIdx * headDimAlign], tri, intriParams);
}
intriParams.blockCount = 1;
intriParams.blockLen = static_cast<uint16_t>(calcLen * numHeads * headDimAlign * sizeof(T) / BLOCK_SIZE);
intriParams.srcStride = 0;
intriParams.dstStride = 0;
// SND -> BSND
for (uint32_t batchIdx = 1; batchIdx < batchSize; ++batchIdx) {
DataCopy(tri[batchIdx * calcLen * numHeads * headDimAlign], tri, intriParams);
}
}
// D -> BND
template <typename T>
__aicore__ inline void BroadCastTriToB1nd(LocalTensor<T> &tri, uint32_t calcLen, uint32_t numHeads,
uint32_t headDimAlign)
{
using ElementType =
typename std::conditional<std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value, half, T>::type;
LocalTensor<ElementType> triNew = tri.template ReinterpretCast<ElementType>();
const int32_t mask = (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) ? MASK_FP16 : MASK_FP32;
const int32_t count = headDimAlign / mask;
const int32_t remain = headDimAlign % mask;
const int32_t repeatTimes = calcLen * numHeads - 1;
const int32_t repeatTimesLoop = repeatTimes / REPEAT_MAX;
const int32_t repeatTimesRemain = repeatTimes % REPEAT_MAX;
CopyRepeatParams repeatParams;
repeatParams.dstStride = 1;
repeatParams.srcStride = 1;
repeatParams.dstRepeatSize = headDimAlign * sizeof(ElementType) / BLOCK_SIZE;
repeatParams.srcRepeatSize = 0;
for (uint32_t loopIdx = 0; loopIdx < count; ++loopIdx) {
for (uint32_t i = 0; i < repeatTimesLoop; ++i) {
Copy(triNew[headDimAlign + loopIdx * mask + i * REPEAT_MAX * headDimAlign], triNew[loopIdx * mask], mask,
REPEAT_MAX, repeatParams);
}
Copy(triNew[headDimAlign * (repeatTimesLoop * REPEAT_MAX + 1) + loopIdx * mask], triNew[loopIdx * mask], mask,
repeatTimesRemain, repeatParams);
}
if (remain != 0) {
for (uint32_t i = 0; i < repeatTimesLoop; ++i) {
Copy(triNew[headDimAlign + count * mask + i * REPEAT_MAX * headDimAlign], triNew[count * mask], remain,
REPEAT_MAX, repeatParams);
}
Copy(triNew[headDimAlign * (repeatTimesLoop * REPEAT_MAX + 1) + count * mask], triNew[count * mask], remain,
repeatTimesRemain, repeatParams);
}
}
// The minimum amount of data set by offset is 8
__aicore__ inline void SetGatherSrcOffset(LocalTensor<int32_t> &gatherOffset, int32_t count, int32_t srcSizeof)
{
for (int32_t i = 0; i < NUM_8; ++i) {
gatherOffset.SetValue(i, i ^ 1); // XOR with 1 to swap even and odd indices
}
int32_t scalarValue = 8;
while (scalarValue < count) {
int32_t nextValue = scalarValue * 2;
if (nextValue < count) {
Adds(gatherOffset[scalarValue], gatherOffset, scalarValue, scalarValue);
} else {
Adds(gatherOffset[scalarValue], gatherOffset, scalarValue, count - scalarValue);
break;
}
scalarValue = nextValue;
}
Muls(gatherOffset, gatherOffset, srcSizeof, count);
}
// count < 256 * 64 and count % 8 == 0
__aicore__ inline void InterleavedInversion(LocalTensor<float> &srcInversion, int32_t count, bool isOffset = false)
{
SetMaskNorm();
const int32_t mask = MASK_FP32;
const int32_t repeatTimes = count / mask;
const int32_t remainder = count % mask;
// Define masks based on the 'isOffset' flag
const uint64_t fullMask = isOffset ? 0xAAAAAAAAAAAAAAAA : 0x5555555555555555;
const uint64_t partialMask = isOffset ? 0xAA : 0x55;
// Apply the mask and multiplication for the full
SetVectorMask<float, MaskMode::NORMAL>(0, fullMask);
Muls<float, false>(srcInversion, srcInversion, float(-1), MASK_PLACEHOLDER, repeatTimes, {1, 1, 8, 8});
// Apply the mask and multiplication for the remainder if needed
if (remainder) {
SetVectorMask<float, MaskMode::NORMAL>(0, partialMask);
Muls<float, false>(srcInversion[repeatTimes * MASK_FP32], srcInversion[repeatTimes * MASK_FP32], float(-1),
MASK_PLACEHOLDER, count % MASK_FP32 / NUM_8, {1, 1, 1, 1});
}
ResetMask();
}
} // namespace RotateInterleavedN
#endif // ROTATE_INTERLEAVED_COMMON_H

316
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bs.h Normal file
View File

@@ -0,0 +1,316 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotate_interleaved_split_bs.h
* \brief
*/
#ifndef ROTATE_INTERLEAVED_SPLIT_BS_H
#define ROTATE_INTERLEAVED_SPLIT_BS_H
#include "rotate_interleaved_common.h"
namespace RotateInterleavedN {
using namespace AscendC;
template <typename T>
class InterleavedSplitBS {
public:
__aicore__ inline InterleavedSplitBS(){};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe);
__aicore__ inline void Process();
protected:
GlobalTensor<T> xGm;
GlobalTensor<T> cosGm;
GlobalTensor<T> sinGm;
GlobalTensor<T> yGm;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueX;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueCos;
TQue<QuePosition::VECOUT, BUFFER_NUM> outQueY;
TBuf<TPosition::VECCALC> tmpFp32Buf1;
TBuf<TPosition::VECCALC> tmpFp32Buf2;
TBuf<TPosition::VECCALC> tmpFp32Buf3;
TBuf<TPosition::VECCALC> gatherOffsetBuf;
const RopeRegbaseTilingData* tiling_;
// tilingdata
uint64_t batchSize;
uint64_t seqLen;
uint64_t numHeads;
uint64_t headDim;
uint64_t frontCoreNum;
uint64_t tailCoreNum;
uint64_t coreCalcNum;
uint64_t coreCalcTail;
uint64_t ubCalcBNum;
uint64_t ubCalcBLoop;
uint64_t ubCalcBTail;
uint64_t allHeadDim;
uint64_t start;
uint64_t ioOffsetAll;
uint64_t bufferNdSizeAll;
// init tmp data
uint32_t blockIdx;
uint32_t ubCalcSeqLoop;
uint64_t ioOffset;
uint64_t triOffset;
uint64_t bufferBsndSize;
uint64_t bufferSdSize;
uint64_t bufferNdSize;
uint64_t bufferLenSize;
uint64_t gatherOffsetLenSize;
__aicore__ inline void InitData(const RopeRegbaseTilingData *tiling);
__aicore__ inline void CopyInX(LocalTensor<T> &x, uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen);
__aicore__ inline void CopyInCos(LocalTensor<T> &cos, uint32_t seqIdx, uint32_t calcLen);
__aicore__ inline void CopyInSin(LocalTensor<T> &sin, uint32_t seqIdx, uint32_t calcLen);
__aicore__ inline void CopyOut(uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen);
__aicore__ inline void Compute(uint32_t seqIdx, uint32_t batchIdx, LocalTensor<uint32_t> &gatherOffsetCast,
uint32_t calcLen);
__aicore__ inline void ComputeCastFp32(uint32_t seqIdx, uint32_t batchIdx, LocalTensor<uint32_t> &gatherOffsetCast,
uint32_t calcLen);
};
template <typename T>
__aicore__ inline void InterleavedSplitBS<T>::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe)
{
InitData(tiling);
blockIdx = GetBlockIdx();
bufferSdSize = seqLen * headDim;
bufferNdSize = numHeads * headDim;
bufferNdSizeAll = numHeads * allHeadDim;
if (blockIdx < frontCoreNum) {
ubCalcSeqLoop = coreCalcNum;
ioOffset = blockIdx * coreCalcNum * bufferNdSize;
ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll;
triOffset = blockIdx * coreCalcNum * headDim;
} else if (coreCalcTail != 0) {
ubCalcSeqLoop = coreCalcTail;
ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize;
ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll;
triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim;
}
bufferBsndSize = batchSize * seqLen * bufferNdSizeAll;
xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSize);
yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSize);
cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize);
sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize);
bufferLenSize = ubCalcBNum * bufferNdSize * sizeof(T);
pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize);
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
bufferLenSize = ubCalcBNum * bufferNdSize * sizeof(float);
pipe->InitBuffer(tmpFp32Buf1, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf2, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf3, bufferLenSize);
}
gatherOffsetLenSize = bufferNdSize * sizeof(int32_t);
pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize);
}
template <typename T>
__aicore__ inline void InterleavedSplitBS<T>::InitData(const RopeRegbaseTilingData *tiling)
{
tiling_ = tiling;
batchSize = tiling_->batchSize;
seqLen = tiling_->seqLen;
numHeads = tiling_->numHeads;
headDim = tiling_->headDim;
frontCoreNum = tiling_->frontCoreNum;
tailCoreNum = tiling_->tailCoreNum;
coreCalcNum = tiling_->coreCalcNum;
coreCalcTail = tiling_->coreCalcTail;
ubCalcBNum = tiling_->ubCalcBNum;
ubCalcBLoop = tiling_->ubCalcBLoop;
ubCalcBTail = tiling_->ubCalcBTail;
allHeadDim = tiling_->allHeadDim;
start = tiling_->start;
}
template <typename T>
__aicore__ inline void InterleavedSplitBS<T>::CopyInX(LocalTensor<T> &x, uint32_t seqIdx, uint32_t batchIdx,
uint32_t calcLen)
{
uint64_t startOffset = batchIdx * ubCalcBNum * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll;
DataCopyExtParams dataCopyParams;
for (uint32_t loopIdx = 0; loopIdx < calcLen; ++loopIdx) {
dataCopyParams.blockCount = numHeads;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = (allHeadDim - headDim) * sizeof(T);
dataCopyParams.dstStride = 0;
DataCopyPad(x[loopIdx * numHeads * headDim], xGm[startOffset + loopIdx * seqLen * bufferNdSizeAll + start],
dataCopyParams, {false, 0, 0, 0});
}
event_t eventIdMTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
}
template <typename T>
__aicore__ inline void InterleavedSplitBS<T>::CopyInCos(LocalTensor<T> &cos, uint32_t seqIdx, uint32_t calcLen)
{
DataCopyExtParams bsDataCopyTriParams;
bsDataCopyTriParams.blockCount = 1;
bsDataCopyTriParams.blockLen = headDim * sizeof(T);
bsDataCopyTriParams.srcStride = 0;
bsDataCopyTriParams.dstStride = 0;
DataCopyPad(cos, cosGm[seqIdx * headDim], bsDataCopyTriParams, {false, 0, 0, 0});
event_t eventId2MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
BroadCastTriToB1nd(cos, calcLen, numHeads, headDim);
}
template <typename T>
__aicore__ inline void InterleavedSplitBS<T>::CopyInSin(LocalTensor<T> &sin, uint32_t seqIdx, uint32_t calcLen)
{
event_t eventIdVToMTE2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2));
SetFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
WaitFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
DataCopyExtParams bsDataCopyTriParams;
bsDataCopyTriParams.blockCount = 1;
bsDataCopyTriParams.blockLen = headDim * sizeof(T);
bsDataCopyTriParams.srcStride = 0;
bsDataCopyTriParams.dstStride = 0;
DataCopyPad(sin, sinGm[seqIdx * headDim], bsDataCopyTriParams, {false, 0, 0, 0});
event_t eventId3MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
BroadCastTriToB1nd(sin, calcLen, numHeads, headDim);
}
template <typename T>
__aicore__ inline void InterleavedSplitBS<T>::CopyOut(uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen)
{
DataCopyExtParams dataCopyParams;
LocalTensor<T> y = outQueY.DeQue<T>();
uint64_t startOffset = batchIdx * ubCalcBNum * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll;
for (uint32_t loopIdx = 0; loopIdx < calcLen; ++loopIdx) {
dataCopyParams.blockCount = numHeads;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = 0;
dataCopyParams.dstStride = (allHeadDim - headDim) * sizeof(T);
DataCopyPad(yGm[startOffset + loopIdx * seqLen * bufferNdSizeAll + start], y[loopIdx * numHeads * headDim],
dataCopyParams);
}
outQueY.FreeTensor(y);
}
template <typename T>
__aicore__ inline void InterleavedSplitBS<T>::Process()
{
LocalTensor<int32_t> gatherOffset = gatherOffsetBuf.Get<int32_t>();
SetGatherSrcOffset(gatherOffset, headDim * numHeads, static_cast<int32_t>(sizeof(float)));
LocalTensor<uint32_t> gatherOffsetCast = gatherOffset.ReinterpretCast<uint32_t>();
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
for (uint32_t i = 0; i < ubCalcSeqLoop; ++i) {
for (uint32_t j = 0; j < (ubCalcBTail == 0 ? ubCalcBLoop : ubCalcBLoop - 1); ++j) {
ComputeCastFp32(i, j, gatherOffsetCast, ubCalcBNum);
CopyOut(i, j, ubCalcBNum);
}
if (ubCalcBTail != 0) {
ComputeCastFp32(i, ubCalcBLoop - 1, gatherOffsetCast, ubCalcBTail);
CopyOut(i, ubCalcBLoop - 1, ubCalcBTail);
}
}
} else {
for (uint32_t i = 0; i < ubCalcSeqLoop; ++i) {
for (uint32_t j = 0; j < (ubCalcBTail == 0 ? ubCalcBLoop : ubCalcBLoop - 1); ++j) {
Compute(i, j, gatherOffsetCast, ubCalcBNum);
CopyOut(i, j, ubCalcBNum);
}
if (ubCalcBTail != 0) {
Compute(i, ubCalcBLoop - 1, gatherOffsetCast, ubCalcBTail);
CopyOut(i, ubCalcBLoop - 1, ubCalcBTail);
}
}
}
}
template <typename T>
__aicore__ inline void InterleavedSplitBS<T>::Compute(uint32_t seqIdx, uint32_t batchIdx,
LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen)
{
uint64_t calcTotalNum = calcLen * bufferNdSize;
LocalTensor<T> x = inQueX.AllocTensor<T>();
CopyInX(x, seqIdx, batchIdx, calcLen);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, seqIdx, calcLen);
LocalTensor<T> y = outQueY.AllocTensor<T>();
Mul(y, x, cos, calcTotalNum);
for (uint32_t i = 0; i < calcLen; ++i) {
Gather(x[i * bufferNdSize], x[i * bufferNdSize], gatherOffsetCast, 0, bufferNdSize);
}
CopyInSin(cos, seqIdx, calcLen);
Mul(x, x, cos, calcTotalNum);
inQueCos.FreeTensor(cos);
InterleavedInversion(x, calcTotalNum);
Add(y, y, x, calcTotalNum);
inQueX.FreeTensor(x);
outQueY.EnQue(y);
}
template <typename T>
__aicore__ inline void InterleavedSplitBS<T>::ComputeCastFp32(uint32_t seqIdx, uint32_t batchIdx,
LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen)
{
uint64_t calcTotalNum = calcLen * bufferNdSize;
LocalTensor<T> x = inQueX.AllocTensor<T>();
CopyInX(x, seqIdx, batchIdx, calcLen);
LocalTensor<float> tmp32BSBuf1 = tmpFp32Buf1.Get<float>();
Cast(tmp32BSBuf1, x, RoundMode::CAST_NONE, calcTotalNum);
inQueX.FreeTensor(x);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, seqIdx, calcLen);
LocalTensor<float> tmp32Buf2 = tmpFp32Buf2.Get<float>();
Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum);
LocalTensor<float> tmp32Buf3 = tmpFp32Buf3.Get<float>();
Mul(tmp32Buf3, tmp32BSBuf1, tmp32Buf2, calcTotalNum);
for (uint32_t i = 0; i < calcLen; ++i) {
Gather(tmp32BSBuf1[i * bufferNdSize], tmp32BSBuf1[i * bufferNdSize], gatherOffsetCast, 0, bufferNdSize);
}
CopyInSin(cos, seqIdx, calcLen);
Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum);
inQueCos.FreeTensor(cos);
Mul(tmp32BSBuf1, tmp32BSBuf1, tmp32Buf2, calcTotalNum);
InterleavedInversion(tmp32BSBuf1, calcTotalNum);
Add(tmp32Buf3, tmp32Buf3, tmp32BSBuf1, calcTotalNum);
LocalTensor<T> y = outQueY.AllocTensor<T>();
Cast(y, tmp32Buf3, RoundMode::CAST_RINT, calcTotalNum);
outQueY.EnQue(y);
}
} // namespace RotateInterleavedN
#endif // ROTATE_INTERLEAVED_SPLIT_BS_H

328
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bs_pad.h Normal file
View File

@@ -0,0 +1,328 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotate_interleaved_split_bs_pad.h
* \brief
*/
#ifndef ROTATE_INTERLEAVED_SPLIT_BS_PAD_H
#define ROTATE_INTERLEAVED_SPLIT_BS_PAD_H
#include "rotate_interleaved_common.h"
namespace RotateInterleavedN {
using namespace AscendC;
template <typename T>
class InterleavedSplitBSPad {
public:
__aicore__ inline InterleavedSplitBSPad(){};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe);
__aicore__ inline void Process();
protected:
GlobalTensor<T> xGm;
GlobalTensor<T> cosGm;
GlobalTensor<T> sinGm;
GlobalTensor<T> yGm;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueX;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueCos;
TQue<QuePosition::VECOUT, BUFFER_NUM> outQueY;
TBuf<TPosition::VECCALC> tmpFp32Buf1;
TBuf<TPosition::VECCALC> tmpFp32Buf2;
TBuf<TPosition::VECCALC> tmpFp32Buf3;
TBuf<TPosition::VECCALC> gatherOffsetBuf;
const RopeRegbaseTilingData* tiling_;
uint64_t allHeadDim;
uint64_t start;
uint64_t ioOffsetAll;
uint64_t bufferNdSizeAll;
// tilingdata
uint64_t batchSize;
uint64_t seqLen;
uint64_t numHeads;
uint64_t headDim;
uint64_t frontCoreNum;
uint64_t tailCoreNum;
uint64_t coreCalcNum;
uint64_t coreCalcTail;
uint64_t ubCalcBNum;
uint64_t ubCalcBLoop;
uint64_t ubCalcBTail;
// init tmp data
uint32_t alignLen;
uint32_t headDimAlign;
uint32_t blockIdx;
uint32_t ubCalcSeq;
uint32_t ubCalcSeqTail;
uint32_t ubCalcSeqLoop;
uint64_t ioOffset;
uint64_t triOffset;
uint64_t bufferBsndSize;
uint64_t bufferSdSize;
uint64_t bufferNdSize;
uint64_t bufferLenSize;
uint64_t gatherOffsetLenSize;
uint32_t blockNum = BLOCK_SIZE / sizeof(T);
__aicore__ inline void InitData(const RopeRegbaseTilingData *tiling);
__aicore__ inline void CopyInX(LocalTensor<T> &x, uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen);
__aicore__ inline void CopyInCos(LocalTensor<T> &cos, uint32_t seqIdx, uint32_t calcLen);
__aicore__ inline void CopyInSin(LocalTensor<T> &sin, uint32_t seqIdx, uint32_t calcLen);
__aicore__ inline void CopyOut(uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen);
__aicore__ inline void Compute(uint32_t seqIdx, uint32_t batchIdx, LocalTensor<uint32_t> &gatherOffsetCast,
uint32_t calcLen);
__aicore__ inline void ComputeCastFp32(uint32_t seqIdx, uint32_t batchIdx, LocalTensor<uint32_t> &gatherOffsetCast,
uint32_t calcLen);
};
template <typename T>
__aicore__ inline void InterleavedSplitBSPad<T>::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe)
{
InitData(tiling);
blockIdx = GetBlockIdx();
bufferSdSize = seqLen * headDim;
bufferNdSize = numHeads * headDim;
bufferNdSizeAll = numHeads * allHeadDim;
if (blockIdx < frontCoreNum) {
ubCalcSeqLoop = coreCalcNum;
ioOffset = blockIdx * coreCalcNum * bufferNdSize;
ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll;
triOffset = blockIdx * coreCalcNum * headDim;
} else if (coreCalcTail != 0) {
ubCalcSeqLoop = coreCalcTail;
ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize;
ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll;
triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim;
}
bufferBsndSize = batchSize * seqLen * bufferNdSizeAll;
xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSize);
yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSize);
cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize);
sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize);
bufferLenSize = ubCalcBNum * numHeads * headDimAlign * sizeof(T);
pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize);
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
bufferLenSize = ubCalcBNum * numHeads * headDimAlign * sizeof(float);
pipe->InitBuffer(tmpFp32Buf1, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf2, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf3, bufferLenSize);
}
gatherOffsetLenSize = numHeads * headDimAlign * sizeof(int32_t);
pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSPad<T>::InitData(const RopeRegbaseTilingData *tiling)
{
tiling_ = tiling;
batchSize = tiling_->batchSize;
seqLen = tiling_->seqLen;
numHeads = tiling_->numHeads;
headDim = tiling_->headDim;
frontCoreNum = tiling_->frontCoreNum;
tailCoreNum = tiling_->tailCoreNum;
coreCalcNum = tiling_->coreCalcNum;
coreCalcTail = tiling_->coreCalcTail;
ubCalcBNum = tiling_->ubCalcBNum;
ubCalcBLoop = tiling_->ubCalcBLoop;
ubCalcBTail = tiling_->ubCalcBTail;
allHeadDim = tiling_->allHeadDim;
start = tiling_->start;
alignLen = (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) ? ALIGN_16 : ALIGN_32;
headDimAlign = (headDim + alignLen - 1) / alignLen * alignLen;
}
template <typename T>
__aicore__ inline void InterleavedSplitBSPad<T>::CopyInX(LocalTensor<T> &x, uint32_t seqIdx, uint32_t batchIdx,
uint32_t calcLen)
{
DataCopyExtParams dataCopyParams;
uint64_t startOffset = batchIdx * ubCalcBNum * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll;
for (uint32_t loopIdx = 0; loopIdx < calcLen; ++loopIdx) {
dataCopyParams.blockCount = numHeads;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = (allHeadDim - headDim) * sizeof(T);
dataCopyParams.dstStride = 0;
DataCopyPad(x[loopIdx * numHeads * headDimAlign], xGm[startOffset + loopIdx * seqLen * bufferNdSizeAll + start],
dataCopyParams, {false, 0, 0, 0});
}
event_t eventIdMTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSPad<T>::CopyInCos(LocalTensor<T> &cos, uint32_t seqIdx, uint32_t calcLen)
{
DataCopyExtParams bsPadDataCopyTriParams;
bsPadDataCopyTriParams.blockCount = 1;
bsPadDataCopyTriParams.blockLen = headDim * sizeof(T);
bsPadDataCopyTriParams.srcStride = 0;
bsPadDataCopyTriParams.dstStride = 0;
DataCopyPad(cos, cosGm[seqIdx * headDim], bsPadDataCopyTriParams, {false, 0, 0, 0});
event_t eventId2MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
BroadCastTriToB1nd(cos, calcLen, numHeads, headDimAlign);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSPad<T>::CopyInSin(LocalTensor<T> &sin, uint32_t seqIdx, uint32_t calcLen)
{
event_t eventIdVToMTE2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2));
SetFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
WaitFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
DataCopyExtParams bsPadDataCopyTriParams;
bsPadDataCopyTriParams.blockCount = 1;
bsPadDataCopyTriParams.blockLen = headDim * sizeof(T);
bsPadDataCopyTriParams.srcStride = 0;
bsPadDataCopyTriParams.dstStride = 0;
DataCopyPad(sin, sinGm[seqIdx * headDim], bsPadDataCopyTriParams, {false, 0, 0, 0});
event_t eventId3MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
BroadCastTriToB1nd(sin, calcLen, numHeads, headDimAlign);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSPad<T>::CopyOut(uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen)
{
LocalTensor<T> y = outQueY.DeQue<T>();
DataCopyExtParams dataCopyParams;
uint64_t startOffset = batchIdx * ubCalcBNum * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll;
for (uint32_t loopIdx = 0; loopIdx < calcLen; ++loopIdx) {
dataCopyParams.blockCount = numHeads;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = 0;
dataCopyParams.dstStride =(allHeadDim - headDim) * sizeof(T);
DataCopyPad(yGm[startOffset + loopIdx * seqLen * bufferNdSizeAll + start], y[loopIdx * numHeads * headDimAlign],
dataCopyParams);
}
outQueY.FreeTensor(y);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSPad<T>::Process()
{
LocalTensor<int32_t> gatherOffset = gatherOffsetBuf.Get<int32_t>();
SetGatherSrcOffset(gatherOffset, headDimAlign * numHeads, static_cast<int32_t>(sizeof(float)));
LocalTensor<uint32_t> gatherOffsetCast = gatherOffset.ReinterpretCast<uint32_t>();
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
for (uint32_t loop = 0; loop < ubCalcSeqLoop; ++loop) {
for (uint32_t j = 0; j < (ubCalcBTail == 0 ? ubCalcBLoop : ubCalcBLoop - 1); ++j) {
ComputeCastFp32(loop, j, gatherOffsetCast, ubCalcBNum);
CopyOut(loop, j, ubCalcBNum);
}
if (ubCalcBTail != 0) {
ComputeCastFp32(loop, ubCalcBLoop - 1, gatherOffsetCast, ubCalcBTail);
CopyOut(loop, ubCalcBLoop - 1, ubCalcBTail);
}
}
} else {
for (uint32_t loop = 0; loop < ubCalcSeqLoop; ++loop) {
for (uint32_t j = 0; j < (ubCalcBTail == 0 ? ubCalcBLoop : ubCalcBLoop - 1); ++j) {
Compute(loop, j, gatherOffsetCast, ubCalcBNum);
CopyOut(loop, j, ubCalcBNum);
}
if (ubCalcBTail != 0) {
Compute(loop, ubCalcBLoop - 1, gatherOffsetCast, ubCalcBTail);
CopyOut(loop, ubCalcBLoop - 1, ubCalcBTail);
}
}
}
}
template <typename T>
__aicore__ inline void InterleavedSplitBSPad<T>::Compute(uint32_t seqIdx, uint32_t batchIdx,
LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen)
{
uint64_t calcTotalNum = calcLen * numHeads * headDimAlign;
LocalTensor<T> x = inQueX.AllocTensor<T>();
CopyInX(x, seqIdx, batchIdx, calcLen);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, seqIdx, calcLen);
LocalTensor<T> yTensor = outQueY.AllocTensor<T>();
Mul(yTensor, x, cos, calcTotalNum);
for (uint32_t i = 0; i < calcLen; ++i) {
Gather(x[i * numHeads * headDimAlign], x[i * numHeads * headDimAlign], gatherOffsetCast, 0,
numHeads * headDimAlign);
}
CopyInSin(cos, seqIdx, calcLen);
Mul(x, x, cos, calcTotalNum);
inQueCos.FreeTensor(cos);
InterleavedInversion(x, calcTotalNum);
Add(yTensor, yTensor, x, calcTotalNum);
inQueX.FreeTensor(x);
outQueY.EnQue(yTensor);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSPad<T>::ComputeCastFp32(uint32_t seqIdx, uint32_t batchIdx,
LocalTensor<uint32_t> &gatherOffsetCast,
uint32_t calcLen)
{
uint64_t calcTotalNum = calcLen * numHeads * headDimAlign;
LocalTensor<T> xTensor = inQueX.AllocTensor<T>();
CopyInX(xTensor, seqIdx, batchIdx, calcLen);
LocalTensor<float> tmp32BsPadBuf1 = tmpFp32Buf1.Get<float>();
Cast(tmp32BsPadBuf1, xTensor, RoundMode::CAST_NONE, calcTotalNum);
inQueX.FreeTensor(xTensor);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, seqIdx, calcLen);
LocalTensor<float> tmp32Buf2 = tmpFp32Buf2.Get<float>();
Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum);
LocalTensor<float> tmp32Buf3 = tmpFp32Buf3.Get<float>();
Mul(tmp32Buf3, tmp32BsPadBuf1, tmp32Buf2, calcTotalNum);
for (uint32_t i = 0; i < calcLen; ++i) {
Gather(tmp32BsPadBuf1[i * numHeads * headDimAlign], tmp32BsPadBuf1[i * numHeads * headDimAlign],
gatherOffsetCast, 0, numHeads * headDimAlign);
}
CopyInSin(cos, seqIdx, calcLen);
Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum);
inQueCos.FreeTensor(cos);
Mul(tmp32BsPadBuf1, tmp32BsPadBuf1, tmp32Buf2, calcTotalNum);
InterleavedInversion(tmp32BsPadBuf1, calcTotalNum);
Add(tmp32Buf3, tmp32Buf3, tmp32BsPadBuf1, calcTotalNum);
LocalTensor<T> y = outQueY.AllocTensor<T>();
Cast(y, tmp32Buf3, RoundMode::CAST_RINT, calcTotalNum);
outQueY.EnQue(y);
}
} // namespace RotateInterleavedN
#endif // ROTATE_INTERLEAVED_SPLIT_BS_PAD_H

314
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bsn.h Normal file
View File

@@ -0,0 +1,314 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotate_interleaved_split_bsn.h
* \brief
*/
#ifndef ROTATE_INTERLEAVED_SPLIT_BSN_H
#define ROTATE_INTERLEAVED_SPLIT_BSN_H
#include "rotate_interleaved_common.h"
namespace RotateInterleavedN {
using namespace AscendC;
template <typename T>
class InterleavedSplitBSN {
public:
__aicore__ inline InterleavedSplitBSN(){};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe);
__aicore__ inline void Process();
protected:
GlobalTensor<T> xGm;
GlobalTensor<T> cosGm;
GlobalTensor<T> sinGm;
GlobalTensor<T> yGm;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueX;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueCos;
TQue<QuePosition::VECOUT, BUFFER_NUM> outQueY;
TBuf<TPosition::VECCALC> tmpFp32Buf1;
TBuf<TPosition::VECCALC> tmpFp32Buf2;
TBuf<TPosition::VECCALC> tmpFp32Buf3;
TBuf<TPosition::VECCALC> gatherOffsetBuf;
const RopeRegbaseTilingData* tiling_;
// tilingdata
uint64_t batchSize;
uint64_t seqLen;
uint64_t numHeads;
uint64_t headDim;
uint64_t frontCoreNum;
uint64_t tailCoreNum;
uint64_t coreCalcNum;
uint64_t coreCalcTail;
uint64_t ubCalcNNum;
uint64_t ubCalcNLoop;
uint64_t ubCalcNTail;
uint64_t allHeadDim;
uint64_t start;
uint64_t ioOffsetAll;
uint64_t bufferNdSizeAll;
// init tmp data
uint32_t blockIdx;
uint32_t ubCalcSeqLoop;
uint64_t ioOffset;
uint64_t triOffset;
uint64_t bufferBsndSize;
uint64_t bufferSdSize;
uint64_t bufferNdSize;
uint64_t bufferLenSize;
uint64_t gatherOffsetLenSize;
__aicore__ inline void InitData(const RopeRegbaseTilingData *tiling);
__aicore__ inline void CopyInX(LocalTensor<T> &x, uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
uint32_t calcLen);
__aicore__ inline void CopyInCos(LocalTensor<T> &cos, uint32_t seqIdx, uint32_t calcLen);
__aicore__ inline void CopyInSin(LocalTensor<T> &sin, uint32_t seqIdx, uint32_t calcLen);
__aicore__ inline void CopyOut(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, uint32_t calcLen);
__aicore__ inline void Compute(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen);
__aicore__ inline void ComputeCastFp32(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen);
};
template <typename T>
__aicore__ inline void InterleavedSplitBSN<T>::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe)
{
InitData(tiling);
blockIdx = GetBlockIdx();
bufferSdSize = seqLen * headDim;
bufferNdSize = numHeads * headDim;
bufferNdSizeAll = numHeads * allHeadDim;
if (blockIdx < frontCoreNum) {
ubCalcSeqLoop = coreCalcNum;
ioOffset = blockIdx * coreCalcNum * bufferNdSize;
ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll;
triOffset = blockIdx * coreCalcNum * headDim;
} else if (coreCalcTail != 0) {
ubCalcSeqLoop = coreCalcTail;
ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize;
ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll;
triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim;
}
bufferBsndSize = batchSize * seqLen * bufferNdSizeAll;
xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSize);
yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSize);
cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize);
sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize);
bufferLenSize = ubCalcNNum * headDim * sizeof(T);
pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize);
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
bufferLenSize = ubCalcNNum * headDim * sizeof(float);
pipe->InitBuffer(tmpFp32Buf1, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf2, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf3, bufferLenSize);
}
gatherOffsetLenSize = ubCalcNNum * headDim * sizeof(int32_t);
pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSN<T>::InitData(const RopeRegbaseTilingData *tiling)
{
tiling_ = tiling;
batchSize = tiling_->batchSize;
seqLen = tiling_->seqLen;
numHeads = tiling_->numHeads;
headDim = tiling_->headDim;
frontCoreNum = tiling_->frontCoreNum;
tailCoreNum = tiling_->tailCoreNum;
coreCalcNum = tiling_->coreCalcNum;
coreCalcTail = tiling_->coreCalcTail;
ubCalcNNum = tiling_->ubCalcNNum;
ubCalcNLoop = tiling_->ubCalcNLoop;
ubCalcNTail = tiling_->ubCalcNTail;
allHeadDim = tiling_->allHeadDim;
start = tiling_->start;
}
template <typename T>
__aicore__ inline void InterleavedSplitBSN<T>::CopyInX(LocalTensor<T> &x, uint32_t batchIdx, uint32_t seqIdx,
uint32_t numHeadsIdx, uint32_t calcLen)
{
DataCopyExtParams dataCopyParams;
dataCopyParams.blockCount = calcLen;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = (allHeadDim - headDim)* sizeof(T);
dataCopyParams.dstStride = 0;
DataCopyPad(x, xGm[batchIdx * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll + numHeadsIdx * ubCalcNNum * allHeadDim + start],
dataCopyParams, {false, 0, 0, 0});
event_t eventIdMTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSN<T>::CopyInCos(LocalTensor<T> &cos, uint32_t seqIdx, uint32_t calcLen)
{
DataCopyExtParams bsnDataCopyTriParams;
bsnDataCopyTriParams.blockCount = 1;
bsnDataCopyTriParams.blockLen = headDim * sizeof(T);
bsnDataCopyTriParams.srcStride = 0;
bsnDataCopyTriParams.dstStride = 0;
DataCopyPad(cos, cosGm[seqIdx * headDim], bsnDataCopyTriParams, {false, 0, 0, 0});
event_t eventId2MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
BroadCastTriToB1nd(cos, 1, calcLen, headDim);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSN<T>::CopyInSin(LocalTensor<T> &sin, uint32_t seqIdx, uint32_t calcLen)
{
event_t eventIdVToMTE2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2));
SetFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
WaitFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
DataCopyExtParams bsnDataCopyTriParams;
bsnDataCopyTriParams.blockCount = 1;
bsnDataCopyTriParams.blockLen = headDim * sizeof(T);
bsnDataCopyTriParams.srcStride = 0;
bsnDataCopyTriParams.dstStride = 0;
DataCopyPad(sin, sinGm[seqIdx * headDim], bsnDataCopyTriParams, {false, 0, 0, 0});
event_t eventId3MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
BroadCastTriToB1nd(sin, 1, calcLen, headDim);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSN<T>::CopyOut(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
uint32_t calcLen)
{
DataCopyExtParams dataCopyParams;
dataCopyParams.blockCount = calcLen;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = 0;
dataCopyParams.dstStride = (allHeadDim - headDim)* sizeof(T);
LocalTensor<T> y = outQueY.DeQue<T>();
DataCopyPad(yGm[batchIdx * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll + numHeadsIdx * ubCalcNNum * allHeadDim + start], y,
dataCopyParams);
outQueY.FreeTensor(y);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSN<T>::Process()
{
LocalTensor<int32_t> gatherOffset = gatherOffsetBuf.Get<int32_t>();
SetGatherSrcOffset(gatherOffset, ubCalcNNum * headDim, static_cast<int32_t>(sizeof(float)));
LocalTensor<uint32_t> gatherOffsetCast = gatherOffset.ReinterpretCast<uint32_t>();
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
for (uint32_t i = 0; i < batchSize; ++i) {
for (uint32_t j = 0; j < ubCalcSeqLoop; ++j) {
for (uint32_t z = 0; z < (ubCalcNTail == 0 ? ubCalcNLoop : ubCalcNLoop - 1); ++z) {
ComputeCastFp32(i, j, z, gatherOffsetCast, ubCalcNNum);
CopyOut(i, j, z, ubCalcNNum);
}
if (ubCalcNTail != 0) {
ComputeCastFp32(i, j, ubCalcNLoop - 1, gatherOffsetCast, ubCalcNTail);
CopyOut(i, j, ubCalcNLoop - 1, ubCalcNTail);
}
}
}
} else {
for (uint32_t i = 0; i < batchSize; ++i) {
for (uint32_t j = 0; j < ubCalcSeqLoop; ++j) {
for (uint32_t z = 0; z < (ubCalcNTail == 0 ? ubCalcNLoop : ubCalcNLoop - 1); ++z) {
Compute(i, j, z, gatherOffsetCast, ubCalcNNum);
CopyOut(i, j, z, ubCalcNNum);
}
if (ubCalcNTail != 0) {
Compute(i, j, ubCalcNLoop - 1, gatherOffsetCast, ubCalcNTail);
CopyOut(i, j, ubCalcNLoop - 1, ubCalcNTail);
}
}
}
}
}
template <typename T>
__aicore__ inline void InterleavedSplitBSN<T>::Compute(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen)
{
uint64_t calcTotalNum = calcLen * headDim;
LocalTensor<T> x = inQueX.AllocTensor<T>();
CopyInX(x, batchIdx, seqIdx, numHeadsIdx, calcLen);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, seqIdx, calcLen);
LocalTensor<T> y = outQueY.AllocTensor<T>();
Mul(y, x, cos, calcTotalNum);
Gather(x, x, gatherOffsetCast, 0, calcTotalNum);
CopyInSin(cos, seqIdx, calcLen);
Mul(x, x, cos, calcTotalNum);
inQueCos.FreeTensor(cos);
InterleavedInversion(x, calcTotalNum);
Add(y, y, x, calcTotalNum);
inQueX.FreeTensor(x);
outQueY.EnQue(y);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSN<T>::ComputeCastFp32(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
LocalTensor<uint32_t> &gatherOffsetCast,
uint32_t calcLen)
{
uint64_t calcTotalNum = calcLen * headDim;
LocalTensor<T> x = inQueX.AllocTensor<T>();
CopyInX(x, batchIdx, seqIdx, numHeadsIdx, calcLen);
LocalTensor<float> tmp32BSNBuf1 = tmpFp32Buf1.Get<float>();
Cast(tmp32BSNBuf1, x, RoundMode::CAST_NONE, calcTotalNum);
inQueX.FreeTensor(x);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, seqIdx, calcLen);
LocalTensor<float> tmp32Buf2 = tmpFp32Buf2.Get<float>();
Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum);
LocalTensor<float> tmp32Buf3 = tmpFp32Buf3.Get<float>();
Mul(tmp32Buf3, tmp32BSNBuf1, tmp32Buf2, calcTotalNum);
Gather(tmp32BSNBuf1, tmp32BSNBuf1, gatherOffsetCast, 0, calcTotalNum);
CopyInSin(cos, seqIdx, calcLen);
Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum);
inQueCos.FreeTensor(cos);
Mul(tmp32BSNBuf1, tmp32BSNBuf1, tmp32Buf2, calcTotalNum);
InterleavedInversion(tmp32BSNBuf1, calcTotalNum);
Add(tmp32Buf3, tmp32Buf3, tmp32BSNBuf1, calcTotalNum);
LocalTensor<T> y = outQueY.AllocTensor<T>();
Cast(y, tmp32Buf3, RoundMode::CAST_RINT, calcTotalNum);
outQueY.EnQue(y);
}
} // namespace RotateInterleavedN
#endif // ROTATE_INTERLEAVED_SPLIT_BSN_H

321
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bsn_pad.h Normal file
View File

@@ -0,0 +1,321 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotate_interleaved_split_bsn_pad.h
* \brief
*/
#ifndef ROTATE_INTERLEAVED_SPLIT_BSN_PAD_H
#define ROTATE_INTERLEAVED_SPLIT_BSN_PAD_H
#include "rotate_interleaved_common.h"
namespace RotateInterleavedN {
using namespace AscendC;
template <typename T>
class InterleavedSplitBSNPad {
public:
__aicore__ inline InterleavedSplitBSNPad(){};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe);
__aicore__ inline void Process();
protected:
GlobalTensor<T> xGm;
GlobalTensor<T> cosGm;
GlobalTensor<T> sinGm;
GlobalTensor<T> yGm;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueX;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueCos;
TQue<QuePosition::VECOUT, BUFFER_NUM> outQueY;
TBuf<TPosition::VECCALC> tmpFp32Buf1;
TBuf<TPosition::VECCALC> tmpFp32Buf2;
TBuf<TPosition::VECCALC> tmpFp32Buf3;
TBuf<TPosition::VECCALC> gatherOffsetBuf;
const RopeRegbaseTilingData* tiling_;
// tilingdata
uint64_t batchSize;
uint64_t seqLen;
uint64_t numHeads;
uint64_t headDim;
uint64_t frontCoreNum;
uint64_t tailCoreNum;
uint64_t coreCalcNum;
uint64_t coreCalcTail;
uint64_t ubCalcNNum;
uint64_t ubCalcNLoop;
uint64_t ubCalcNTail;
uint64_t allHeadDim;
uint64_t start;
uint64_t ioOffsetAll;
uint64_t bufferNdSizeAll;
// init tmp data
uint32_t alignLen;
uint32_t headDimAlign;
uint32_t blockIdx;
uint32_t ubCalcSeqLoop;
uint64_t ioOffset;
uint64_t triOffset;
uint64_t bufferBsndSize;
uint64_t bufferSdSize;
uint64_t bufferNdSize;
uint64_t bufferLenSize;
uint64_t gatherOffsetLenSize;
__aicore__ inline void InitData(const RopeRegbaseTilingData *tiling);
__aicore__ inline void CopyInX(LocalTensor<T> &x, uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
uint32_t calcLen);
__aicore__ inline void CopyInCos(LocalTensor<T> &cos, uint32_t seqIdx, uint32_t calcLen);
__aicore__ inline void CopyInSin(LocalTensor<T> &sin, uint32_t seqIdx, uint32_t calcLen);
__aicore__ inline void CopyOut(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, uint32_t calcLen);
__aicore__ inline void Compute(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen);
__aicore__ inline void ComputeCastFp32(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen);
};
template <typename T>
__aicore__ inline void InterleavedSplitBSNPad<T>::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe)
{
InitData(tiling);
blockIdx = GetBlockIdx();
bufferSdSize = seqLen * headDim;
bufferNdSize = numHeads * headDim;
bufferNdSizeAll = numHeads * allHeadDim;
if (blockIdx < frontCoreNum) {
ubCalcSeqLoop = coreCalcNum;
ioOffset = blockIdx * coreCalcNum * bufferNdSize;
ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll;
triOffset = blockIdx * coreCalcNum * headDim;
} else if (coreCalcTail != 0) {
ubCalcSeqLoop = coreCalcTail;
ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize;
ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll;
triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim;
}
bufferBsndSize = batchSize * seqLen * bufferNdSizeAll;
xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSize);
yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSize);
cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize);
sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize);
bufferLenSize = ubCalcNNum * headDimAlign * sizeof(T);
pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize);
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
bufferLenSize = ubCalcNNum * headDimAlign * sizeof(float);
pipe->InitBuffer(tmpFp32Buf1, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf2, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf3, bufferLenSize);
}
gatherOffsetLenSize = ubCalcNNum * headDimAlign * sizeof(int32_t);
pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSNPad<T>::InitData(const RopeRegbaseTilingData *tiling)
{
tiling_ = tiling;
batchSize = tiling_->batchSize;
seqLen = tiling_->seqLen;
numHeads = tiling_->numHeads;
headDim = tiling_->headDim;
frontCoreNum = tiling_->frontCoreNum;
tailCoreNum = tiling_->tailCoreNum;
coreCalcNum = tiling_->coreCalcNum;
coreCalcTail = tiling_->coreCalcTail;
ubCalcNNum = tiling_->ubCalcNNum;
ubCalcNLoop = tiling_->ubCalcNLoop;
ubCalcNTail = tiling_->ubCalcNTail;
allHeadDim = tiling_->allHeadDim;
start = tiling_->start;
alignLen = (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) ? ALIGN_16 : ALIGN_32;
headDimAlign = (headDim + alignLen - 1) / alignLen * alignLen;
}
template <typename T>
__aicore__ inline void InterleavedSplitBSNPad<T>::CopyInX(LocalTensor<T> &xTensor, uint32_t batchIdx, uint32_t seqIdx,
uint32_t numHeadsIdx, uint32_t calcLen)
{
DataCopyExtParams dataCopyParams;
dataCopyParams.blockCount = calcLen;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = (allHeadDim - headDim) * sizeof(T);
dataCopyParams.dstStride = 0;
DataCopyPad(xTensor,
xGm[batchIdx * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll + numHeadsIdx * ubCalcNNum * allHeadDim + start],
dataCopyParams, {false, 0, 0, 0});
event_t eventIdMTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSNPad<T>::CopyInCos(LocalTensor<T> &cos, uint32_t seqIdx, uint32_t calcLen)
{
DataCopyExtParams bsnPadDataCopyTriParams;
bsnPadDataCopyTriParams.blockCount = 1;
bsnPadDataCopyTriParams.blockLen = headDim * sizeof(T);
bsnPadDataCopyTriParams.srcStride = 0;
bsnPadDataCopyTriParams.dstStride = 0;
DataCopyPad(cos, cosGm[seqIdx * headDim], bsnPadDataCopyTriParams, {false, 0, 0, 0});
event_t eventId2MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
BroadCastTriToB1nd(cos, 1, calcLen, headDimAlign);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSNPad<T>::CopyInSin(LocalTensor<T> &sinTensor, uint32_t seqIdx,
uint32_t calcLen)
{
event_t eventIdVToMTE2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2));
SetFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
WaitFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
DataCopyExtParams bsnPadDataCopyTriParams;
bsnPadDataCopyTriParams.blockCount = 1;
bsnPadDataCopyTriParams.blockLen = headDim * sizeof(T);
bsnPadDataCopyTriParams.srcStride = 0;
bsnPadDataCopyTriParams.dstStride = 0;
DataCopyPad(sinTensor, sinGm[seqIdx * headDim], bsnPadDataCopyTriParams, {false, 0, 0, 0});
event_t eventId3MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
BroadCastTriToB1nd(sinTensor, 1, calcLen, headDimAlign);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSNPad<T>::CopyOut(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
uint32_t calcLen)
{
DataCopyExtParams dataCopyParams;
dataCopyParams.blockCount = calcLen;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = 0;
dataCopyParams.dstStride = (allHeadDim - headDim) * sizeof(T);
LocalTensor<T> y = outQueY.DeQue<T>();
DataCopyPad(yGm[batchIdx * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll + numHeadsIdx * ubCalcNNum * allHeadDim + start], y,
dataCopyParams);
outQueY.FreeTensor(y);
}
template <typename T>
__aicore__ inline void InterleavedSplitBSNPad<T>::Process()
{
LocalTensor<int32_t> gatherOffset = gatherOffsetBuf.Get<int32_t>();
SetGatherSrcOffset(gatherOffset, ubCalcNNum * headDimAlign, static_cast<int32_t>(sizeof(float)));
LocalTensor<uint32_t> gatherOffsetCast = gatherOffset.ReinterpretCast<uint32_t>();
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
for (uint32_t batch = 0; batch < batchSize; ++batch) {
for (uint32_t j = 0; j < ubCalcSeqLoop; ++j) {
for (uint32_t z = 0; z < (ubCalcNTail == 0 ? ubCalcNLoop : ubCalcNLoop - 1); ++z) {
ComputeCastFp32(batch, j, z, gatherOffsetCast, ubCalcNNum);
CopyOut(batch, j, z, ubCalcNNum);
}
if (ubCalcNTail != 0) {
ComputeCastFp32(batch, j, ubCalcNLoop - 1, gatherOffsetCast, ubCalcNTail);
CopyOut(batch, j, ubCalcNLoop - 1, ubCalcNTail);
}
}
}
} else {
for (uint32_t batch = 0; batch < batchSize; ++batch) {
for (uint32_t j = 0; j < ubCalcSeqLoop; ++j) {
for (uint32_t z = 0; z < (ubCalcNTail == 0 ? ubCalcNLoop : ubCalcNLoop - 1); ++z) {
Compute(batch, j, z, gatherOffsetCast, ubCalcNNum);
CopyOut(batch, j, z, ubCalcNNum);
}
if (ubCalcNTail != 0) {
Compute(batch, j, ubCalcNLoop - 1, gatherOffsetCast, ubCalcNTail);
CopyOut(batch, j, ubCalcNLoop - 1, ubCalcNTail);
}
}
}
}
}
template <typename T>
__aicore__ inline void InterleavedSplitBSNPad<T>::Compute(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen)
{
uint64_t calcTotalNum = calcLen * headDimAlign;
LocalTensor<T> xTensor = inQueX.AllocTensor<T>();
CopyInX(xTensor, batchIdx, seqIdx, numHeadsIdx, calcLen);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, seqIdx, calcLen);
LocalTensor<T> yTensor = outQueY.AllocTensor<T>();
Mul(yTensor, xTensor, cos, calcTotalNum);
Gather(xTensor, xTensor, gatherOffsetCast, 0, calcTotalNum);
CopyInSin(cos, seqIdx, calcLen);
Mul(xTensor, xTensor, cos, calcTotalNum);
inQueCos.FreeTensor(cos);
InterleavedInversion(xTensor, calcTotalNum);
Add(yTensor, yTensor, xTensor, calcTotalNum);
inQueX.FreeTensor(xTensor);
outQueY.EnQue(yTensor);
}
template <typename T>
__aicore__ inline void
InterleavedSplitBSNPad<T>::ComputeCastFp32(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx,
LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen)
{
uint64_t totalCount = calcLen * headDimAlign;
LocalTensor<T> xTensor = inQueX.AllocTensor<T>();
CopyInX(xTensor, batchIdx, seqIdx, numHeadsIdx, calcLen);
LocalTensor<float> tmp32Buf1 = tmpFp32Buf1.Get<float>();
Cast(tmp32Buf1, xTensor, RoundMode::CAST_NONE, totalCount);
inQueX.FreeTensor(xTensor);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, seqIdx, calcLen);
LocalTensor<float> tmp32Buf2 = tmpFp32Buf2.Get<float>();
Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, totalCount);
LocalTensor<float> tmp32Buf3 = tmpFp32Buf3.Get<float>();
Mul(tmp32Buf3, tmp32Buf1, tmp32Buf2, totalCount);
Gather(tmp32Buf1, tmp32Buf1, gatherOffsetCast, 0, totalCount);
CopyInSin(cos, seqIdx, calcLen);
Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, totalCount);
inQueCos.FreeTensor(cos);
Mul(tmp32Buf1, tmp32Buf1, tmp32Buf2, totalCount);
InterleavedInversion(tmp32Buf1, totalCount);
Add(tmp32Buf3, tmp32Buf3, tmp32Buf1, totalCount);
LocalTensor<T> y = outQueY.AllocTensor<T>();
Cast(y, tmp32Buf3, RoundMode::CAST_RINT, totalCount);
outQueY.EnQue(y);
}
} // namespace RotateInterleavedN
#endif // ROTATE_INTERLEAVED_SPLIT_BSN_PAD_H

313
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_s.h Normal file
View File

@@ -0,0 +1,313 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotate_interleaved_split_s.h
* \brief
*/
#ifndef ROTATE_INTERLEAVED_SPLIT_S_H
#define ROTATE_INTERLEAVED_SPLIT_S_H
#include "rotate_interleaved_common.h"
namespace RotateInterleavedN {
using namespace AscendC;
template <typename T>
class InterleavedSplitS {
public:
__aicore__ inline InterleavedSplitS(){};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe);
__aicore__ inline void Process();
protected:
GlobalTensor<T> xGm;
GlobalTensor<T> cosGm;
GlobalTensor<T> sinGm;
GlobalTensor<T> yGm;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueX;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueCos;
TQue<QuePosition::VECOUT, BUFFER_NUM> outQueY;
TBuf<TPosition::VECCALC> tmpFp32Buf1;
TBuf<TPosition::VECCALC> tmpFp32Buf2;
TBuf<TPosition::VECCALC> tmpFp32Buf3;
TBuf<TPosition::VECCALC> gatherOffsetBuf;
const RopeRegbaseTilingData* tiling_;
// tilingdata
uint64_t batchSize;
uint64_t seqLen;
uint64_t numHeads;
uint64_t headDim;
uint64_t frontCoreNum;
uint64_t tailCoreNum;
uint64_t coreCalcNum;
uint64_t coreCalcTail;
uint64_t ubCalcNum;
uint64_t ubCalcLoop;
uint64_t ubCalcTail;
uint64_t ubCalcTailNum;
uint64_t ubCalcTailLoop;
uint64_t ubCalcTailTail;
uint64_t allHeadDim;
uint64_t start;
// init tmp data
uint32_t blockIdx;
uint32_t ubCalcSeq;
uint32_t ubCalcSeqTail;
uint32_t ubCalcSeqLoop;
uint64_t ioOffset;
uint64_t ioOffsetAll;
uint64_t triOffset;
uint64_t bufferBsndSize;
uint64_t bufferBsndSizeAll;
uint64_t bufferSdSize;
uint64_t bufferNdSize;
uint64_t bufferNdSizeAll;
uint64_t bufferLenSize;
uint64_t gatherOffsetLenSize;
__aicore__ inline void InitData(const RopeRegbaseTilingData *tiling);
__aicore__ inline void CopyInX(LocalTensor<T> &x, uint32_t loopIdx, uint32_t calcLen);
__aicore__ inline void CopyInCos(LocalTensor<T> &cos, uint32_t loopIdx, uint32_t calcLen);
__aicore__ inline void CopyInSin(LocalTensor<T> &sin, uint32_t loopIdx, uint32_t calcLen);
__aicore__ inline void CopyOut(uint32_t loopIdx, uint32_t calcLen);
__aicore__ inline void Compute(uint32_t loopIdx, LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen);
__aicore__ inline void ComputeCastFp32(uint32_t loopIdx, LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen);
};
template <typename T>
__aicore__ inline void InterleavedSplitS<T>::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe)
{
InitData(tiling);
blockIdx = GetBlockIdx();
bufferSdSize = seqLen * headDim;
bufferNdSize = numHeads * headDim;
bufferNdSizeAll = numHeads * allHeadDim;
if (blockIdx < frontCoreNum) {
ubCalcSeq = ubCalcNum;
ubCalcSeqTail = ubCalcTail;
ubCalcSeqLoop = ubCalcLoop;
ioOffset = blockIdx * coreCalcNum * bufferNdSize;
ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll;
triOffset = blockIdx * coreCalcNum * headDim;
} else if (coreCalcTail != 0) {
ubCalcSeq = ubCalcTailNum;
ubCalcSeqTail = ubCalcTailTail;
ubCalcSeqLoop = ubCalcTailLoop;
ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize;
ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll;
triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim;
}
bufferBsndSize = batchSize * seqLen * bufferNdSize;
bufferBsndSizeAll = batchSize * seqLen * bufferNdSizeAll;
xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSizeAll);
yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSizeAll);
cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize);
sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize);
bufferLenSize = batchSize * ubCalcSeq * bufferNdSize * sizeof(T);
pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize);
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
bufferLenSize = batchSize * ubCalcSeq * bufferNdSize * sizeof(float);
pipe->InitBuffer(tmpFp32Buf1, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf2, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf3, bufferLenSize);
}
gatherOffsetLenSize = bufferNdSize * sizeof(int32_t);
pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize);
}
template <typename T>
__aicore__ inline void InterleavedSplitS<T>::InitData(const RopeRegbaseTilingData *tiling)
{
tiling_ = tiling;
batchSize = tiling_->batchSize; //B
seqLen = tiling_->seqLen; // S
numHeads = tiling_->numHeads; // N
headDim = tiling_->headDim; // D
frontCoreNum = tiling_->frontCoreNum;
tailCoreNum = tiling_->tailCoreNum;
coreCalcNum = tiling_->coreCalcNum;
coreCalcTail = tiling_->coreCalcTail;
ubCalcNum = tiling_->ubCalcNum;
ubCalcLoop = tiling_->ubCalcLoop;
ubCalcTail = tiling_->ubCalcTail;
ubCalcTailNum = tiling_->ubCalcTailNum;
ubCalcTailLoop = tiling_->ubCalcTailLoop;
ubCalcTailTail = tiling_->ubCalcTailTail;
allHeadDim = tiling_->allHeadDim;
start = tiling_->start;
}
template <typename T>
__aicore__ inline void InterleavedSplitS<T>::CopyInX(LocalTensor<T> &x, uint32_t loopIdx, uint32_t calcLen)
{
DataCopyExtParams dataCopyParams;
dataCopyParams.blockCount = calcLen*numHeads;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = (allHeadDim - headDim)* sizeof(T);
dataCopyParams.dstStride = 0;
DataCopyPad(x, xGm[loopIdx * ubCalcSeq * bufferNdSizeAll + start], dataCopyParams, {false, 0, 0, 0});
event_t eventIdMTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
}
template <typename T>
__aicore__ inline void InterleavedSplitS<T>::CopyInCos(LocalTensor<T> &cos, uint32_t loopIdx, uint32_t calcLen)
{
DataCopyExtParams dataCopyTriParams;
dataCopyTriParams.blockCount = calcLen;
dataCopyTriParams.blockLen = headDim * sizeof(T);
dataCopyTriParams.srcStride = 0;
dataCopyTriParams.dstStride = static_cast<uint16_t>((numHeads - 1) * headDim * sizeof(T) / BLOCK_SIZE);
DataCopyPad(cos, cosGm[loopIdx * ubCalcSeq * headDim], dataCopyTriParams, {false, 0, 0, 0});
event_t eventId2MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
BroadCastTriToBsnd(cos, batchSize, calcLen, numHeads, headDim);
}
template <typename T>
__aicore__ inline void InterleavedSplitS<T>::CopyInSin(LocalTensor<T> &sin, uint32_t loopIdx, uint32_t calcLen)
{
event_t eventIdVToMTE2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2));
SetFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
WaitFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
DataCopyExtParams dataCopyTriParams;
dataCopyTriParams.blockCount = calcLen;
dataCopyTriParams.blockLen = headDim * sizeof(T);
dataCopyTriParams.srcStride = 0;
dataCopyTriParams.dstStride = static_cast<uint16_t>((numHeads - 1) * headDim * sizeof(T) / BLOCK_SIZE);
DataCopyPad(sin, sinGm[loopIdx * ubCalcSeq * headDim], dataCopyTriParams, {false, 0, 0, 0});
event_t eventId3MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
BroadCastTriToBsnd(sin, batchSize, calcLen, numHeads, headDim);
}
template <typename T>
__aicore__ inline void InterleavedSplitS<T>::CopyOut(uint32_t loopIdx, uint32_t calcLen)
{
DataCopyExtParams dataCopyParams;
dataCopyParams.blockCount = calcLen*numHeads;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = 0;
dataCopyParams.dstStride = (allHeadDim - headDim) * sizeof(T);
LocalTensor<T> y = outQueY.DeQue<T>();
DataCopyPad(yGm[loopIdx * ubCalcSeq * bufferNdSizeAll + start], y, dataCopyParams);
outQueY.FreeTensor(y);
}
template <typename T>
__aicore__ inline void InterleavedSplitS<T>::Process()
{
LocalTensor<int32_t> gatherOffset = gatherOffsetBuf.Get<int32_t>();
SetGatherSrcOffset(gatherOffset, headDim * numHeads, static_cast<int32_t>(sizeof(float)));
LocalTensor<uint32_t> gatherOffsetCast = gatherOffset.ReinterpretCast<uint32_t>();
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
for (uint32_t i = 0; i < (ubCalcSeqTail == 0 ? ubCalcSeqLoop : ubCalcSeqLoop - 1); ++i) {
ComputeCastFp32(i, gatherOffsetCast, ubCalcSeq);
CopyOut(i, ubCalcSeq);
}
if (ubCalcSeqTail != 0) {
ComputeCastFp32(ubCalcSeqLoop - 1, gatherOffsetCast, ubCalcSeqTail);
CopyOut(ubCalcSeqLoop - 1, ubCalcSeqTail);
}
} else {
for (uint32_t i = 0; i < (ubCalcSeqTail == 0 ? ubCalcSeqLoop : ubCalcSeqLoop - 1); ++i) {
Compute(i, gatherOffsetCast, ubCalcSeq);
CopyOut(i, ubCalcSeq);
}
if (ubCalcSeqTail != 0) {
Compute(ubCalcSeqLoop - 1, gatherOffsetCast, ubCalcSeqTail);
CopyOut(ubCalcSeqLoop - 1, ubCalcSeqTail);
}
}
}
template <typename T>
__aicore__ inline void InterleavedSplitS<T>::Compute(uint32_t loopIdx, LocalTensor<uint32_t> &gatherOffsetCast,
uint32_t calcLen)
{
uint64_t calcTotalNum = calcLen * batchSize * bufferNdSize;
LocalTensor<T> x = inQueX.AllocTensor<T>();
CopyInX(x, loopIdx, calcLen);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, loopIdx, calcLen);
LocalTensor<T> y = outQueY.AllocTensor<T>();
Mul(y, x, cos, calcTotalNum);
for (uint32_t i = 0; i < batchSize * calcLen; ++i) {
Gather(x[i * bufferNdSize], x[i * bufferNdSize], gatherOffsetCast, 0, bufferNdSize);
}
CopyInSin(cos, loopIdx, calcLen);
Mul(x, x, cos, calcTotalNum);
inQueCos.FreeTensor(cos);
InterleavedInversion(x, calcTotalNum);
Add(y, y, x, calcTotalNum);
inQueX.FreeTensor(x);
outQueY.EnQue(y);
}
template <typename T>
__aicore__ inline void InterleavedSplitS<T>::ComputeCastFp32(uint32_t loopIdx, LocalTensor<uint32_t> &gatherOffsetCast,
uint32_t calcLen)
{
uint64_t calcTotalNum = calcLen * batchSize * bufferNdSize;
LocalTensor<T> x = inQueX.AllocTensor<T>();
CopyInX(x, loopIdx, calcLen);
LocalTensor<float> tmp32Buf1 = tmpFp32Buf1.Get<float>();
Cast(tmp32Buf1, x, RoundMode::CAST_NONE, calcTotalNum);
inQueX.FreeTensor(x);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, loopIdx, calcLen);
LocalTensor<float> tmp32Buf2 = tmpFp32Buf2.Get<float>();
Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum);
LocalTensor<float> tmp32Buf3 = tmpFp32Buf3.Get<float>();
Mul(tmp32Buf3, tmp32Buf1, tmp32Buf2, calcTotalNum);
for (uint32_t i = 0; i < batchSize * calcLen; ++i) {
Gather(tmp32Buf1[i * bufferNdSize], tmp32Buf1[i * bufferNdSize], gatherOffsetCast, 0, bufferNdSize);
}
CopyInSin(cos, loopIdx, calcLen);
Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum);
inQueCos.FreeTensor(cos);
Mul(tmp32Buf1, tmp32Buf1, tmp32Buf2, calcTotalNum);
InterleavedInversion(tmp32Buf1, calcTotalNum);
Add(tmp32Buf3, tmp32Buf3, tmp32Buf1, calcTotalNum);
LocalTensor<T> yTensor = outQueY.AllocTensor<T>();
Cast(yTensor, tmp32Buf3, RoundMode::CAST_RINT, calcTotalNum);
outQueY.EnQue(yTensor);
}
} // namespace RotateInterleavedN
#endif // ROTATE_INTERLEAVED_SPLIT_S_H

328
csrc/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_s_pad.h Normal file
View File

@@ -0,0 +1,328 @@
/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file rotate_interleaved_split_s_pad.h
* \brief
*/
#ifndef ROTATE_INTERLEAVED_SPLIT_S_PAD_H
#define ROTATE_INTERLEAVED_SPLIT_S_PAD_H
#include "rotate_interleaved_common.h"
namespace RotateInterleavedN {
using namespace AscendC;
template <typename T>
class InterleavedSplitSPad {
public:
__aicore__ inline InterleavedSplitSPad(){};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe);
__aicore__ inline void Process();
protected:
GlobalTensor<T> xGm;
GlobalTensor<T> cosGm;
GlobalTensor<T> sinGm;
GlobalTensor<T> yGm;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueX;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueCos;
TQue<QuePosition::VECOUT, BUFFER_NUM> outQueY;
TBuf<TPosition::VECCALC> tmpFp32Buf1;
TBuf<TPosition::VECCALC> tmpFp32Buf2;
TBuf<TPosition::VECCALC> tmpFp32Buf3;
TBuf<TPosition::VECCALC> gatherOffsetBuf;
const RopeRegbaseTilingData* tiling_;
// tilingdata
uint64_t batchSize;
uint64_t seqLen;
uint64_t numHeads;
uint64_t headDim;
uint64_t frontCoreNum;
uint64_t tailCoreNum;
uint64_t coreCalcNum;
uint64_t coreCalcTail;
uint64_t ubCalcNum;
uint64_t ubCalcLoop;
uint64_t ubCalcTail;
uint64_t ubCalcTailNum;
uint64_t ubCalcTailLoop;
uint64_t ubCalcTailTail;
uint64_t bufferNdSizeAll;
uint64_t allHeadDim;
uint64_t start;
// init tmp data
uint32_t alignLen;
uint32_t headDimAlign;
uint32_t allHeadDimAlign;
uint32_t blockIdx;
uint32_t ubCalcSeq;
uint32_t ubCalcSeqTail;
uint32_t ubCalcSeqLoop;
uint64_t ioOffset;
uint64_t ioOffsetAll;
uint64_t triOffset;
uint64_t bufferBsndSize;
uint64_t bufferSdSize;
uint64_t bufferNdSize;
uint64_t bufferLenSize;
uint64_t gatherOffsetLenSize;
uint32_t blockNum = BLOCK_SIZE / sizeof(T);
__aicore__ inline void InitData(const RopeRegbaseTilingData *tiling);
__aicore__ inline void CopyInX(LocalTensor<T> &x, uint32_t loopIdx, uint32_t calcLen);
__aicore__ inline void CopyInCos(LocalTensor<T> &cos, uint32_t loopIdx, uint32_t calcLen);
__aicore__ inline void CopyInSin(LocalTensor<T> &sin, uint32_t loopIdx, uint32_t calcLen);
__aicore__ inline void CopyOut(uint32_t loopIdx, uint32_t calcLen);
__aicore__ inline void Compute(uint32_t loopIdx, LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen);
__aicore__ inline void ComputeCastFp32(uint32_t loopIdx, LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen);
};
template <typename T>
__aicore__ inline void InterleavedSplitSPad<T>::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y,
const RopeRegbaseTilingData *tiling, TPipe *pipe)
{
InitData(tiling);
blockIdx = GetBlockIdx();
bufferSdSize = seqLen * headDim;
bufferNdSize = numHeads * headDim;
bufferNdSizeAll = numHeads * allHeadDim;
if (blockIdx < frontCoreNum) {
ubCalcSeq = ubCalcNum;
ubCalcSeqTail = ubCalcTail;
ubCalcSeqLoop = ubCalcLoop;
ioOffset = blockIdx * coreCalcNum * bufferNdSize;
ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll;
triOffset = blockIdx * coreCalcNum * headDim;
} else if (coreCalcTail != 0) {
ubCalcSeq = ubCalcTailNum;
ubCalcSeqTail = ubCalcTailTail;
ubCalcSeqLoop = ubCalcTailLoop;
ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize;
ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll;
triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim;
}
bufferBsndSize = batchSize * seqLen * bufferNdSizeAll;
xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSize);
yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSize);
cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize);
sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize);
bufferLenSize = batchSize * ubCalcSeq * numHeads * headDimAlign * sizeof(T);
pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize);
pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize);
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
bufferLenSize = batchSize * ubCalcSeq * numHeads * headDimAlign * sizeof(float);
pipe->InitBuffer(tmpFp32Buf1, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf2, bufferLenSize);
pipe->InitBuffer(tmpFp32Buf3, bufferLenSize);
}
gatherOffsetLenSize = numHeads * headDimAlign * sizeof(int32_t);
pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize);
}
template <typename T>
__aicore__ inline void InterleavedSplitSPad<T>::InitData(const RopeRegbaseTilingData *tiling)
{
tiling_ = tiling;
batchSize = tiling_->batchSize;
seqLen = tiling_->seqLen;
numHeads = tiling_->numHeads;
headDim = tiling_->headDim;
frontCoreNum = tiling_->frontCoreNum;
tailCoreNum = tiling_->tailCoreNum;
coreCalcNum = tiling_->coreCalcNum;
coreCalcTail = tiling_->coreCalcTail;
ubCalcNum = tiling_->ubCalcNum;
ubCalcLoop = tiling_->ubCalcLoop;
ubCalcTail = tiling_->ubCalcTail;
ubCalcTailNum = tiling_->ubCalcTailNum;
ubCalcTailLoop = tiling_->ubCalcTailLoop;
ubCalcTailTail = tiling_->ubCalcTailTail;
allHeadDim = tiling_->allHeadDim;
start = tiling_->start;
alignLen = (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) ? ALIGN_16 : ALIGN_32;
headDimAlign = (headDim + alignLen - 1) / alignLen * alignLen;
allHeadDimAlign = (allHeadDim + alignLen - 1) / alignLen * alignLen;
}
template <typename T>
__aicore__ inline void InterleavedSplitSPad<T>::CopyInX(LocalTensor<T> &x, uint32_t loopIdx, uint32_t calcLen)
{
DataCopyExtParams dataCopyParams;
for (uint32_t batchIdx = 0; batchIdx < batchSize; ++batchIdx) {
dataCopyParams.blockCount = calcLen * numHeads;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = (allHeadDim - headDim)* sizeof(T);
dataCopyParams.dstStride = 0;
DataCopyPad(x[batchIdx * calcLen * numHeads * headDimAlign],
xGm[batchIdx * seqLen * bufferNdSizeAll + loopIdx * ubCalcSeq * bufferNdSizeAll + start], dataCopyParams,
{false, 0, 0, 0});
}
event_t eventIdMTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
}
template <typename T>
__aicore__ inline void InterleavedSplitSPad<T>::CopyInCos(LocalTensor<T> &cos, uint32_t loopIdx, uint32_t calcLen)
{
DataCopyExtParams dataCopyTriParams;
dataCopyTriParams.blockCount = calcLen;
dataCopyTriParams.blockLen = headDim * sizeof(T);
dataCopyTriParams.srcStride = 0;
dataCopyTriParams.dstStride = static_cast<uint16_t>((numHeads - 1) * headDimAlign / blockNum);
DataCopyPad(cos, cosGm[loopIdx * ubCalcSeq * headDim], dataCopyTriParams, {false, 0, 0, 0});
event_t eventId2MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId2MTE2ToV);
BroadCastTriToBsnd(cos, batchSize, calcLen, numHeads, headDimAlign);
}
template <typename T>
__aicore__ inline void InterleavedSplitSPad<T>::CopyInSin(LocalTensor<T> &sin, uint32_t loopIdx, uint32_t calcLen)
{
event_t eventIdVToMTE2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2));
SetFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
WaitFlag<HardEvent::V_MTE2>(eventIdVToMTE2);
DataCopyExtParams dataCopyTriParams;
dataCopyTriParams.blockCount = calcLen;
dataCopyTriParams.blockLen = headDim * sizeof(T);
dataCopyTriParams.srcStride = 0;
dataCopyTriParams.dstStride = static_cast<uint16_t>((numHeads - 1) * headDimAlign / blockNum);
DataCopyPad(sin, sinGm[loopIdx * ubCalcSeq * headDim], dataCopyTriParams, {false, 0, 0, 0});
event_t eventId3MTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
WaitFlag<HardEvent::MTE2_V>(eventId3MTE2ToV);
BroadCastTriToBsnd(sin, batchSize, calcLen, numHeads, headDimAlign);
}
template <typename T>
__aicore__ inline void InterleavedSplitSPad<T>::CopyOut(uint32_t loopIdx, uint32_t calcLen)
{
LocalTensor<T> y = outQueY.DeQue<T>();
DataCopyExtParams dataCopyParams;
for (uint32_t batchIdx = 0; batchIdx < batchSize; ++batchIdx) {
dataCopyParams.blockCount = calcLen * numHeads;
dataCopyParams.blockLen = headDim * sizeof(T);
dataCopyParams.srcStride = 0;
dataCopyParams.dstStride = (allHeadDim - headDim)* sizeof(T);
DataCopyPad(yGm[batchIdx * seqLen * bufferNdSizeAll + loopIdx * ubCalcSeq * bufferNdSizeAll + start],
y[batchIdx * calcLen * numHeads * headDimAlign], dataCopyParams);
}
outQueY.FreeTensor(y);
}
template <typename T>
__aicore__ inline void InterleavedSplitSPad<T>::Process()
{
LocalTensor<int32_t> gatherOffset = gatherOffsetBuf.Get<int32_t>();
SetGatherSrcOffset(gatherOffset, headDimAlign * numHeads, static_cast<int32_t>(sizeof(float)));
LocalTensor<uint32_t> gatherOffsetCast = gatherOffset.ReinterpretCast<uint32_t>();
if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
for (uint32_t loop = 0; loop < (ubCalcSeqTail == 0 ? ubCalcSeqLoop : ubCalcSeqLoop - 1); ++loop) {
ComputeCastFp32(loop, gatherOffsetCast, ubCalcSeq);
CopyOut(loop, ubCalcSeq);
}
if (ubCalcSeqTail != 0) {
ComputeCastFp32(ubCalcSeqLoop - 1, gatherOffsetCast, ubCalcSeqTail);
CopyOut(ubCalcSeqLoop - 1, ubCalcSeqTail);
}
} else {
for (uint32_t loop = 0; loop < (ubCalcSeqTail == 0 ? ubCalcSeqLoop : ubCalcSeqLoop - 1); ++loop) {
Compute(loop, gatherOffsetCast, ubCalcSeq);
CopyOut(loop, ubCalcSeq);
}
if (ubCalcSeqTail != 0) {
Compute(ubCalcSeqLoop - 1, gatherOffsetCast, ubCalcSeqTail);
CopyOut(ubCalcSeqLoop - 1, ubCalcSeqTail);
}
}
}
template <typename T>
__aicore__ inline void InterleavedSplitSPad<T>::Compute(uint32_t loopIdx, LocalTensor<uint32_t> &gatherOffsetCast,
uint32_t calcLen)
{
uint64_t totalCount = calcLen * batchSize * numHeads * headDimAlign;
LocalTensor<T> x = inQueX.AllocTensor<T>();
CopyInX(x, loopIdx, calcLen);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, loopIdx, calcLen);
LocalTensor<T> y = outQueY.AllocTensor<T>();
Mul(y, x, cos, totalCount);
for (uint32_t i = 0; i < batchSize * calcLen; ++i) {
Gather(x[i * numHeads * headDimAlign], x[i * numHeads * headDimAlign], gatherOffsetCast, 0,
numHeads * headDimAlign);
}
CopyInSin(cos, loopIdx, calcLen);
Mul(x, x, cos, totalCount);
inQueCos.FreeTensor(cos);
InterleavedInversion(x, totalCount);
Add(y, y, x, totalCount);
inQueX.FreeTensor(x);
outQueY.EnQue(y);
}
template <typename T>
__aicore__ inline void
InterleavedSplitSPad<T>::ComputeCastFp32(uint32_t loopIdx, LocalTensor<uint32_t> &gatherOffsetCast, uint32_t calcLen)
{
uint64_t totalCount = calcLen * batchSize * numHeads * headDimAlign;
LocalTensor<T> xTensor = inQueX.AllocTensor<T>();
CopyInX(xTensor, loopIdx, calcLen);
LocalTensor<float> tmp32SPadBuf1 = tmpFp32Buf1.Get<float>();
Cast(tmp32SPadBuf1, xTensor, RoundMode::CAST_NONE, totalCount);
inQueX.FreeTensor(xTensor);
LocalTensor<T> cos = inQueCos.AllocTensor<T>();
CopyInCos(cos, loopIdx, calcLen);
LocalTensor<float> tmp32SPadBuf2 = tmpFp32Buf2.Get<float>();
Cast(tmp32SPadBuf2, cos, RoundMode::CAST_NONE, totalCount);
LocalTensor<float> tmp32Buf3 = tmpFp32Buf3.Get<float>();
Mul(tmp32Buf3, tmp32SPadBuf1, tmp32SPadBuf2, totalCount);
for (uint32_t i = 0; i < batchSize * calcLen; ++i) {
Gather(tmp32SPadBuf1[i * numHeads * headDimAlign], tmp32SPadBuf1[i * numHeads * headDimAlign], gatherOffsetCast,
0, numHeads * headDimAlign);
}
CopyInSin(cos, loopIdx, calcLen);
Cast(tmp32SPadBuf2, cos, RoundMode::CAST_NONE, totalCount);
inQueCos.FreeTensor(cos);
Mul(tmp32SPadBuf1, tmp32SPadBuf1, tmp32SPadBuf2, totalCount);
InterleavedInversion(tmp32SPadBuf1, totalCount);
Add(tmp32Buf3, tmp32Buf3, tmp32SPadBuf1, totalCount);
LocalTensor<T> y = outQueY.AllocTensor<T>();
Cast(y, tmp32Buf3, RoundMode::CAST_RINT, totalCount);
outQueY.EnQue(y);
}
} // namespace RotateInterleavedN
#endif // ROTATE_INTERLEAVED_SPLIT_S_PAD_H