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

46
csrc/attention/load_index_kv_cache/op_kernel/load_index_kv_cache.cpp Normal file
View File

@@ -0,0 +1,46 @@
/**
* 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 load_index_kv_cache.cpp
* \brief
*/
#include "load_index_kv_cache_perf.h"
using namespace AscendC;
extern "C" __global__ __aicore__ void load_index_kv_cache(
GM_ADDR kv_cache,
GM_ADDR slot_mapping,
GM_ADDR kv,
GM_ADDR kv_scale,
GM_ADDR workspace,
GM_ADDR tiling)
{
if (workspace == nullptr) {
return;
}
GM_ADDR userWs = GetUserWorkspace(workspace);
if (userWs == nullptr) {
return;
}
GET_TILING_DATA(tilingData, tiling);
TPipe pipe;
if (TILING_KEY_IS(0)) {
LoadIndexKvCache::LoadIndexKvCachePerf<DTYPE_KV_CACHE> op;
op.Init(kv_cache, slot_mapping, kv, kv_scale, userWs, &tilingData, &pipe);
op.Process();
return;
}
}

89
csrc/attention/load_index_kv_cache/op_kernel/load_index_kv_cache_base.h Normal file
View File

@@ -0,0 +1,89 @@
/**
* 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 load_index_kv_cache_base.h
* \brief
*/
#ifndef LOAD_INDEX_LV_CACHE_BASE_H
#define LOAD_INDEX_LV_CACHE_BASE_H
#include "kernel_operator.h"
namespace LoadIndexKvCache {
using namespace AscendC;
constexpr int32_t BLOCK_SIZE = 32;
constexpr int32_t VL_FP32 = 64;
constexpr int32_t KV_LAST_DIM = 128;
constexpr int32_t KV_SCALE_LAST_DIM = 4;
__aicore__ inline int32_t CeilDiv(int32_t a, int b)
{
if (b == 0) {
return a;
}
return (a + b - 1) / b;
}
__aicore__ inline int32_t CeilAlign(int32_t a, int b)
{
return CeilDiv(a, b) * b;
}
template <typename T>
__aicore__ inline int32_t RoundUp(int32_t num)
{
int32_t elemNum = BLOCK_SIZE / sizeof(T);
return CeilAlign(num, elemNum);
}
template <typename T>
__aicore__ inline int32_t RoundUp(int32_t num, int32_t elemNum)
{
return CeilAlign(num, elemNum);
}
template <typename T>
__aicore__ inline void CopyIn(
const GlobalTensor<T>& inputGm, const LocalTensor<T>& inputTensor, const uint16_t nBurst, const uint32_t copyLen,
uint32_t srcStride = 0)
{
DataCopyPadExtParams<T> dataCopyPadExtParams;
dataCopyPadExtParams.isPad = false;
dataCopyPadExtParams.leftPadding = 0;
dataCopyPadExtParams.rightPadding = 0;
dataCopyPadExtParams.paddingValue = 0;
DataCopyExtParams dataCoptExtParams;
dataCoptExtParams.blockCount = nBurst;
dataCoptExtParams.blockLen = copyLen * sizeof(T);
dataCoptExtParams.srcStride = srcStride * sizeof(T);
dataCoptExtParams.dstStride = 0;
DataCopyPad(inputTensor, inputGm, dataCoptExtParams, dataCopyPadExtParams);
}
template <typename T, AscendC::PaddingMode mode = AscendC::PaddingMode::Normal>
__aicore__ inline void CopyOut(
const LocalTensor<T>& outputTensor, const GlobalTensor<T>& outputGm, const uint16_t nBurst, const uint32_t copyLen,
uint32_t dstStride = 0)
{
DataCopyExtParams dataCopyParams;
dataCopyParams.blockCount = nBurst;
dataCopyParams.blockLen = copyLen * sizeof(T);
dataCopyParams.srcStride = 0;
dataCopyParams.dstStride = dstStride * sizeof(T);
DataCopyPad<T, mode>(outputGm, outputTensor, dataCopyParams);
}
} // namespace IndexerCompressEpilogV2
#endif

110
csrc/attention/load_index_kv_cache/op_kernel/load_index_kv_cache_perf.h Normal file
View File

@@ -0,0 +1,110 @@
/**
* 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 load_index_kv_cache_single_row.h
* \brief
*/
#ifndef LOAD_INDEX_KV_CACH_PERF_H
#define LOAD_INDEX_KV_CACH_PERF_H
#include "kernel_operator.h"
#include "load_index_kv_cache_base.h"
namespace LoadIndexKvCache {
using namespace AscendC;
template <typename T>
class LoadIndexKvCachePerf {
public:
__aicore__ inline LoadIndexKvCachePerf()
{}
__aicore__ inline void Init(
GM_ADDR kvCache, GM_ADDR slotMapping, GM_ADDR kv, GM_ADDR kvScale, GM_ADDR workspace,
const LoadIndexKvCacheTilingData* tilingDataPtr, TPipe* pipePtr)
{
pipe = pipePtr;
tilingData = tilingDataPtr;
kvCacheGm.SetGlobalBuffer((__gm__ T*)kvCache);
slotMappingGm.SetGlobalBuffer((__gm__ int32_t*)slotMapping);
kvGm.SetGlobalBuffer((__gm__ T*)kv);
kvScaleGm.SetGlobalBuffer((__gm__ float*)kvScale);
pipe->InitBuffer(kvCacheQue, 2, RoundUp<T>(tilingData->d) * sizeof(T));
pipe->InitBuffer(kvAndKvScaleQue, 2, RoundUp<T>(tilingData->d) * sizeof(T));
}
__aicore__ inline void Process()
{
int64_t curBlockIdx = GetBlockIdx();
int64_t rowOuterLoop =
(curBlockIdx == GetBlockNum() - 1) ? tilingData->rowOfTailBlock : tilingData->rowOfFormerBlock;
int64_t baseSlotMappingOffset = curBlockIdx * tilingData->rowOfFormerBlock;
int64_t kvBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock * KV_LAST_DIM;
int64_t kvScaleBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock;
for (int64_t rowOuterIdx = 0; rowOuterIdx < rowOuterLoop; rowOuterIdx++) {
int64_t curSlotIdx = baseSlotMappingOffset + rowOuterIdx;
int64_t slot = slotMappingGm.GetValue(curSlotIdx);
if (slot == -1) {
// slot为-1时，需要将对应位置的kv和kvscale全部刷成0
kvAndKvScaleLocal = kvAndKvScaleQue.template AllocTensor<T>();
Duplicate(kvAndKvScaleLocal.template ReinterpretCast<uint8_t>(), uint8_t(0), KV_LAST_DIM);
Duplicate(kvAndKvScaleLocal[KV_LAST_DIM].template ReinterpretCast<float>(), float(0), 1);
kvAndKvScaleQue.template EnQue(kvAndKvScaleLocal);
kvAndKvScaleLocal = kvAndKvScaleQue.template DeQue<T>();
CopyOut(kvAndKvScaleLocal, kvGm[kvBaseOffset + rowOuterIdx * KV_LAST_DIM], 1, KV_LAST_DIM);
CopyOut(kvAndKvScaleLocal[KV_LAST_DIM].template ReinterpretCast<float>(), kvScaleGm[kvScaleBaseOffset + rowOuterIdx], 1, 1);
kvAndKvScaleQue.template FreeTensor(kvAndKvScaleLocal);
continue;
}
// 处理kvCacheLocal
kvCacheLocal = kvCacheQue.template AllocTensor<T>();
int64_t kvCacheGmOffset = (slot / tilingData->bs) * tilingData->blockStride + (slot % tilingData->bs) * KV_LAST_DIM;
int64_t kvCacheScaleGmOffset = (slot / tilingData->bs) * tilingData->blockStride + tilingData->bs * KV_LAST_DIM + (slot % tilingData->bs) * KV_SCALE_LAST_DIM;
CopyIn(kvCacheGm[kvCacheGmOffset].template ReinterpretCast<uint8_t>(), kvCacheLocal.template ReinterpretCast<uint8_t>(), 1, KV_LAST_DIM);
CopyIn(kvCacheGm[kvCacheScaleGmOffset].template ReinterpretCast<uint8_t>(), kvCacheLocal[KV_LAST_DIM].template ReinterpretCast<uint8_t>(), 1, KV_SCALE_LAST_DIM);
kvCacheQue.template EnQue(kvCacheLocal);
kvCacheLocal = kvCacheQue.template DeQue<T>();
kvAndKvScaleLocal = kvAndKvScaleQue.AllocTensor<T>();
DataCopy(kvAndKvScaleLocal, kvCacheLocal, RoundUp<T>(tilingData->d));
kvCacheQue.template FreeTensor(kvCacheLocal);
kvAndKvScaleQue.template EnQue(kvAndKvScaleLocal);
kvAndKvScaleLocal = kvAndKvScaleQue.template DeQue<T>();
CopyOut(kvAndKvScaleLocal, kvGm[kvBaseOffset + rowOuterIdx * KV_LAST_DIM], 1, KV_LAST_DIM);
CopyOut(kvAndKvScaleLocal[KV_LAST_DIM].template ReinterpretCast<float>(), kvScaleGm[kvScaleBaseOffset + rowOuterIdx * 1], 1, 1);
kvAndKvScaleQue.template FreeTensor(kvAndKvScaleLocal);
}
}
private:
TPipe* pipe;
const LoadIndexKvCacheTilingData* tilingData;
GlobalTensor<T> kvCacheGm;
GlobalTensor<int32_t> slotMappingGm;
GlobalTensor<T> kvGm;
GlobalTensor<float> kvScaleGm;
TQue<QuePosition::VECIN, 1> kvCacheQue;
TQue<QuePosition::VECOUT, 1> kvAndKvScaleQue;
TQue<QuePosition::VECOUT, 1> kvScaleQue;
LocalTensor<T> kvCacheLocal;
LocalTensor<T> kvAndKvScaleLocal;
};
} // namespace LoadIndexKvCache
#endif