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

130
csrc/attention/compressor/op_kernel/arch32/rope.h Normal file
View File

@@ -0,0 +1,130 @@
/**
 * 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 rope.h
* \brief
*/
#ifndef ROPE_H
#define ROPE_H
#include "compressor_comm.h"
#include "compressor_vector_comm.h"
namespace Compressor {
/**
* @brief SetGatherSrcOffset 计算用于interleave模式的offset
* @param gatherOffsetLocal 输出tensor [count]，数据类型需要为int64_t，使用时要转换
* @param count offset的元素个数，一般为列数
*/
template <typename T>
__aicore__ inline void SetGatherSrcOffset(const LocalTensor<int32_t> &gatherOffsetLocal, uint32_t count)
{
for (uint32_t i = 0; i < 8; i++) {
gatherOffsetLocal.SetValue(i, i ^ 1);
}
event_t eventId_S_V = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
SetFlag<HardEvent::S_V>(eventId_S_V);
WaitFlag<HardEvent::S_V>(eventId_S_V);
int32_t scalarValue = 8;
while (scalarValue < count) {
int32_t nextValue = scalarValue * 2;
PipeBarrier<PIPE_V>();
if (nextValue < count) {
Adds(gatherOffsetLocal[scalarValue], gatherOffsetLocal, scalarValue, scalarValue);
} else {
Adds(gatherOffsetLocal[scalarValue], gatherOffsetLocal, scalarValue, count - scalarValue);
break;
}
scalarValue = nextValue;
}
PipeBarrier<PIPE_V>();
Muls(gatherOffsetLocal, gatherOffsetLocal, static_cast<int32_t>(sizeof(T)), count);
}
/**
* @brief RotaryPosEmb 同时做row行的RotaryPosEmb，每一行的元素为col
* @param dstLocal 输出tensor [row, actualCol]，支持和srcLocal是同一块空间
* @param srcLocal 输入tensor [row, actualCol]
* @param cosLocal cos系数tensor [row, col]
* @param sinLocal sin系数tensor [row, col]
* @param shareTmpUb 临时buffer 内部需要的空间为 [row * col * sizeof(float)]
* @param gatherOffsetcastLocal 用于interleave模式的offset，数据类型需要为uint64_t
* @param row 待处理的行数
* @param col 待处理的列数
* @param actualCol 实际列数
* @param baseAddr 计算基地址
*/
template <ROTARY_MODE MODE>
__aicore__ inline void RotaryPosEmb(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
const LocalTensor<float> &cosLocal, const LocalTensor<float> &sinLocal,
const LocalTensor<float> &shareTmpUb,
const LocalTensor<uint32_t> &gatherOffsetcastLocal, uint32_t row, uint32_t col,
uint32_t actualCol, uint64_t baseAddr)
{
uint64_t cnt = row * col;
uint32_t half_col = col >> 1;
uint64_t rsvdCnt = 0;
LocalTensor<float> reArrLocal = shareTmpUb.ReinterpretCast<float>();
if constexpr (MODE == ROTARY_MODE::HALF) {
DataCopy(reArrLocal, srcLocal[baseAddr + half_col],
{static_cast<uint16_t>(row), static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint16_t>(CeilDivT(actualCol - half_col, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM))});
DataCopy(reArrLocal[half_col], srcLocal[baseAddr],
{static_cast<uint16_t>(row), static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint16_t>(CeilDivT(actualCol - half_col, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM))});
PipeBarrier<PIPE_V>();
Muls(reArrLocal, reArrLocal, float(-1), half_col, row,
{1, 1, static_cast<uint8_t>(CeilDivT(static_cast<uint32_t>(col), FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint8_t>(CeilDivT(static_cast<uint32_t>(col), FP32_BLOCK_ELEMENT_NUM))});
} else if constexpr (MODE == ROTARY_MODE::INTERLEAVE) {
for (uint32_t i = 0; i < row; i++) {
Gather(reArrLocal[i * col], srcLocal[i * actualCol + baseAddr], gatherOffsetcastLocal, 0, col);
}
PipeBarrier<PIPE_V>();
uint32_t repeatTimes = cnt / FP32_REPEAT_ELEMENT_NUM;
uint32_t remainder = cnt % FP32_REPEAT_ELEMENT_NUM;
uint64_t fullMask = 0x5555555555555555;
uint64_t partialMask = 0x55;
SetVectorMask<float, MaskMode::NORMAL>(0, fullMask);
Muls<float, false>(reArrLocal, reArrLocal, float(-1), MASK_PLACEHOLDER, repeatTimes,
{1, 1, FP32_BLOCK_ELEMENT_NUM, FP32_BLOCK_ELEMENT_NUM});
if (unlikely(remainder > 0)) {
SetVectorMask<float, MaskMode::NORMAL>(0, partialMask);
Muls<float, false>(reArrLocal[repeatTimes * FP32_REPEAT_ELEMENT_NUM],
reArrLocal[repeatTimes * FP32_REPEAT_ELEMENT_NUM], float(-1), MASK_PLACEHOLDER,
remainder / FP32_BLOCK_ELEMENT_NUM, {1, 1, 1, 1});
}
ResetMask();
}
PipeBarrier<PIPE_V>();
BinaryRepeatParams computeParams{1,
1,
1,
static_cast<uint8_t>(CeilDivT(actualCol, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint8_t>(CeilDivT(actualCol, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint8_t>(CeilDivT(col, FP32_BLOCK_ELEMENT_NUM))};
Mul(dstLocal[baseAddr], srcLocal[baseAddr], cosLocal, col, row, computeParams);
Mul(reArrLocal, reArrLocal, sinLocal, cnt);
PipeBarrier<PIPE_V>();
Add(dstLocal[baseAddr], dstLocal[baseAddr], reArrLocal, col, row, computeParams);
}
}
#endif