Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 0 additions & 37 deletions csrc/attention/store_kv_block/op_host/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -6,43 +6,6 @@
# 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.
# ======================================================================================================================
#add_definitions(-DMAX_BLOCK_NUM_=3)
# add_ops_compile_options(
# OP_NAME StoreKVBlock
# OPTIONS --cce-auto-sync=on
# -Wno-deprecated-declarations
# -Werror
# )

# # -o0
# # -g
# # --cce-ignore-always-inline=true
# target_sources(op_host_aclnn PRIVATE
# store_kv_block_def.cpp
# )

# target_sources(optiling PRIVATE
# store_kv_block_tiling.cpp
# store_kv_block_common.cpp
# )

# if (NOT BUILD_OPEN_PROJECT)
# target_sources(opmaster_ct PRIVATE
# store_kv_block_tiling.cpp
# )
# endif ()

# target_include_directories(optiling PRIVATE
# ${CMAKE_CURRENT_SOURCE_DIR}
# )

# target_sources(opsproto PRIVATE
# store_kv_block_infershape.cpp
# )

# target_link_libraries(optiling
# PRIVATE
# )
add_op_to_compiled_list()

if (BUILD_OPEN_PROJECT)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,9 @@
*/
#include <graph/utils/type_utils.h>
#include <register/op_impl_registry.h>

#include "error/ops_error.h"

static constexpr int IDX_0 = 0;
static constexpr int IDX_1 = 1;
static constexpr int IDX_2 = 2;

using namespace ge;
// using namespace Ops::Base;

namespace ops {

static ge::graphStatus InferShape4StoreKVBlock(gert::InferShapeContext* context)
Expand Down
18 changes: 11 additions & 7 deletions csrc/attention/store_kv_block/op_host/store_kv_block_tiling.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,9 @@ struct StoreKVBlockParams {
uint32_t blockTableSize{0};
uint32_t typeByte{0};
uint32_t tokenSize{1};
uint32_t tilingKey{0};
uint32_t tilingKey{1};
uint64_t workspaceSize{0};
uint64_t groupInfoLen{0};
uint32_t groupInfoLen{0};
uint32_t corepernum{0};
uint32_t coretail{0};
uint64_t sysWorkspaceSize{0};
Expand Down Expand Up @@ -121,11 +121,15 @@ static ge::graphStatus StoreKVBlockTilingFunc(gert::TilingContext* context) {
}

StoreKVBlockTilingData tilingData;
if (params.blockTableSize > 0) tilingData.set_blockTableSize(params.blockTableSize);
if (params.typeByte > 0) tilingData.set_typeByte(params.typeByte);
if (params.tokenSize > 0) tilingData.set_tokenSize(params.tokenSize);
if (params.corepernum > 0 || params.coretail != 0) tilingData.set_corePerNum(params.corepernum);
if (params.coretail < 48) tilingData.set_coreTail(params.coretail);
// if (params.blockTableSize > 0) tilingData.set_blockTableSize(params.blockTableSize);
// if (params.typeByte > 0) tilingData.set_typeByte(params.typeByte);
// if (params.tokenSize > 0) tilingData.set_tokenSize(params.tokenSize);
// if (params.corepernum > 0 || params.coretail != 0) tilingData.set_corePerNum(params.corepernum);
tilingData.set_blockTableSize(params.blockTableSize);
tilingData.set_typeByte(params.typeByte);
tilingData.set_tokenSize(params.tokenSize);
tilingData.set_corePerNum(params.corepernum);
if (params.coretail < params.coreNum) tilingData.set_coreTail(params.coretail);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The variable params.coreNum is used here, but coreNum is not a member of the StoreKVBlockParams struct. This will cause a compilation failure. Please use the correct local variable or platform API to retrieve the core count.

Suggested change
if (params.coretail < params.coreNum) tilingData.set_coreTail(params.coretail);
if (params.coretail < coreNum) tilingData.set_coreTail(params.coretail);

if (params.numTokens > 0) tilingData.set_numTokens(params.numTokens);
if (params.numCache > 0) tilingData.set_numCache(params.numCache);
if (params.groupInfoLen > 0) tilingData.set_groupInfoLen(params.groupInfoLen);
Expand Down
23 changes: 11 additions & 12 deletions csrc/attention/store_kv_block/op_kernel/store_kv_block.h
Original file line number Diff line number Diff line change
Expand Up @@ -62,9 +62,9 @@ class StoreKVBlockBase {
AscendC::LocalTensor<T> tokenLocal;
AscendC::GlobalTensor<T> keyInputGt;
AscendC::GlobalTensor<T> keyCacheInputGt;
AscendC::GlobalTensor<uint32_t> groupLenGt;
AscendC::GlobalTensor<uint32_t> groupKeyIdxGt;
AscendC::GlobalTensor<uint32_t> groupKeyCacheIdxGt;
AscendC::GlobalTensor<int32_t> groupLenGt;
AscendC::GlobalTensor<int32_t> groupKeyIdxGt;
AscendC::GlobalTensor<int32_t> groupKeyCacheIdxGt;
AscendC::TBuf<AscendC::TPosition::VECCALC> tokenBuf;
__aicore__ inline StoreKVBlockBase() {}

Expand All @@ -84,7 +84,6 @@ class StoreKVBlockBase {
numCache = tilingData->numCache;
groupInfoLen = tilingData->groupInfoLen;


coreId = AscendC::GetBlockIdx();
coreTail = tilingData->coreTail;
blockNum = AscendC::GetBlockNum();
Expand All @@ -101,21 +100,21 @@ class StoreKVBlockBase {

keyInputGt.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(keyIn));
keyCacheInputGt.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(keyCacheIn));
groupLenGt.SetGlobalBuffer(reinterpret_cast<__gm__ uint32_t*>(groupLen));
groupKeyIdxGt.SetGlobalBuffer(reinterpret_cast<__gm__ uint32_t*>(groupKeyIdx));
groupKeyCacheIdxGt.SetGlobalBuffer(reinterpret_cast<__gm__ uint32_t*>(groupKeyCacheIdx));
groupLenGt.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(groupLen));
groupKeyIdxGt.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(groupKeyIdx));
groupKeyCacheIdxGt.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(groupKeyCacheIdx));

pipeThis->InitBuffer(tokenBuf, blockTableSize*tokenByteSize);
tokenLocal = tokenBuf.Get<T>();

AscendC::DataCopyExtParams copyParams{1, 0, 0, 0, 0}; // todo: full block length
AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
for (uint32_t i = 0; i < corePerNum; i++) {
uint32_t idx = (coreId+i*blockNum);
for (int32_t i = 0; i < corePerNum; i++) {
int32_t idx = (coreId+i*blockNum);

// if( groupLenGt.GetValue(idx)<= 0 || groupKeyIdxGt.GetValue(idx)<0 || groupKeyCacheIdxGt.GetValue(idx)<0){
// continue;
// }
if( groupLenGt.GetValue(idx)<= 0 || groupKeyIdxGt.GetValue(idx)<0 || groupKeyCacheIdxGt.GetValue(idx)<0){
continue;
}
Comment on lines +115 to +117

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

In Ascend C, GlobalTensor does not support direct scalar access via GetValue(idx). Attempting to compile this will result in a compilation error. Even if supported by some emulators, performing direct global memory reads inside a loop on the AI Core is extremely inefficient and will severely degrade performance. To resolve this, allocate a small LocalTensor in the Unified Buffer (UB), copy the metadata tensors from Global Memory to UB using DataCopy before the loop, and then access them locally.


copyParams.blockLen = groupLenGt.GetValue(idx)*tokenByteSize; // in bytes
DataCopyPad(tokenLocal, keyInputGt[ groupKeyIdxGt.GetValue(idx)*tokenSize], copyParams, padParams); // note: offset order
Expand Down
92 changes: 0 additions & 92 deletions csrc/attention/store_kv_block/store_kv_block_torch_adpt.h
Original file line number Diff line number Diff line change
Expand Up @@ -20,98 +20,6 @@
#include <climits>
namespace vllm_ascend {

std::tuple<at::Tensor, at::Tensor, at::Tensor> store_kv_block_pre(
const at::Tensor &slot_mapping_npu,
at::IntArrayRef slot_mapping_list,
int64_t block_size)
{

int64_t slot_mapping_len = slot_mapping_list.size();

std::vector<int32_t> length(16, 0);
std::vector<int32_t> key_idx(16, 0);
std::vector<int32_t> key_cache_idx(16, 0);
int32_t idx_slotmap = 0;
int32_t idx_groups = 0;

while (idx_slotmap < slot_mapping_len) {

int32_t current_idx = slot_mapping_list[idx_slotmap];
if(current_idx <0){
idx_slotmap++;
continue;
}

int32_t block_id = current_idx / block_size;
int32_t y= current_idx % block_size;

key_idx[idx_groups] = idx_slotmap;
key_cache_idx[idx_groups] = current_idx;

int32_t j = idx_slotmap;

if(j+1 < slot_mapping_len &&slot_mapping_list[j+1]!=slot_mapping_list[j]+1 ) {
j++;

}else{
int32_t idx_stride = std::min(block_size-y,slot_mapping_len-idx_slotmap)-1;
int32_t expected_last = current_idx + idx_stride;
int32_t expected_last_idx = idx_slotmap + (expected_last-current_idx);

if (expected_last == slot_mapping_list[expected_last_idx]){
j = expected_last_idx+1;
}else{

while(j+1 < slot_mapping_len && slot_mapping_list[j] / block_size == block_id && slot_mapping_list[j+1] ==slot_mapping_list[j]+1) {
j++;
}
}
}

length[idx_groups] = (j - idx_slotmap);
idx_slotmap = j;
idx_groups++;

if(idx_groups>=length.capacity()){
int32_t new_capacity = length.capacity() * 2;
length.reserve(new_capacity);
key_idx.reserve(new_capacity);
key_cache_idx.reserve(new_capacity);

for (int32_t k = idx_groups; k < new_capacity; ++k){
length.emplace_back(0);
key_idx.emplace_back(0);
key_cache_idx.emplace_back(0);
}
}
}

at::Tensor group_len = at::empty({idx_groups},
at::TensorOptions(slot_mapping_npu.options().device()).dtype(torch::kInt32)
);
void* group_len_addr = group_len.data_ptr();

at::Tensor group_key_idx = at::empty({idx_groups},
at::TensorOptions(slot_mapping_npu.options().device()).dtype(torch::kInt32)
);
void* group_key_idx_addr = group_key_idx.data_ptr();

at::Tensor group_key_cache_idx = at::empty({idx_groups},
at::TensorOptions(slot_mapping_npu.options().device()).dtype(torch::kInt32)
);
void* group_key_cache_idx_addr = group_key_cache_idx.data_ptr();

uint32_t device_size=idx_groups*sizeof(length[0]);
aclrtStream stream = c10_npu::getCurrentNPUStream().stream();
aclrtMemcpyKind memcpy_type=ACL_MEMCPY_HOST_TO_DEVICE;
aclrtMemcpyAsync(group_len_addr, device_size, &length[0], device_size, ACL_MEMCPY_HOST_TO_DEVICE, stream);
aclrtMemcpyAsync(group_key_idx_addr, device_size, &key_idx[0], device_size, ACL_MEMCPY_HOST_TO_DEVICE, stream);
aclrtMemcpyAsync(group_key_cache_idx_addr, device_size, &key_cache_idx[0], device_size, ACL_MEMCPY_HOST_TO_DEVICE, stream);

return std::tuple<at::Tensor, at::Tensor, at::Tensor>(group_len, group_key_idx, group_key_cache_idx);

}

void store_kv_block(
const at::Tensor &key_in,
const at::Tensor &key_cache_in,
Expand Down
10 changes: 10 additions & 0 deletions csrc/attention/store_kv_block_metadata/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
# ---------------------------------------------------------------------------------------------------------
# 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 should 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.
# ---------------------------------------------------------------------------------------------------------
add_modules_sources_aicpu(OPTYPE store_kv_block ACLNNTYPE aclnn)
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
/**
* 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 should 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 aclnn_store_kv_block_metadata.cpp
* \brief AClnn interface for StoreKvBlockMetadata operator
*/

#include "aclnn_store_kv_block_metadata.h"
#include "l0_store_kv_block_metadata.h"
#include "aclnn_kernels/contiguous.h"
#include "aclnn_kernels/reshape.h"
#include "aclnn/aclnn_base.h"
#include "aclnn_kernels/common/op_error_check.h"
#include "opdev/common_types.h"
#include "opdev/data_type_utils.h"
#include "opdev/format_utils.h"
#include "opdev/op_dfx.h"
#include "opdev/op_executor.h"
#include "opdev/op_log.h"
#include "opdev/tensor_view_utils.h"
#include "opdev/make_op_executor.h"

#ifdef __cplusplus
extern "C" {
#endif

aclnnStatus aclnnStoreKvBlockMetadataGetWorkspaceSize(
const aclTensor *slotMapping,
const aclTensor *groupLen,
const aclTensor *groupKeyIdx,
const aclTensor *groupKeyCacheIdx,
int64_t blockSize,
uint64_t *workspaceSize,
aclOpExecutor **executor)
{
L2_DFX_PHASE_1(aclnnStoreKvBlockMetadata,
DFX_IN(slotMapping, blockSize),
DFX_OUT(groupLen, groupKeyIdx, groupKeyCacheIdx));

auto uniqueExecutor = CREATE_EXECUTOR();
CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);

// basic parameter checks
CHECK_RET(slotMapping != nullptr, ACLNN_ERR_INNER_NULLPTR);
CHECK_RET(groupLen != nullptr, ACLNN_ERR_INNER_NULLPTR);
CHECK_RET(groupKeyIdx != nullptr, ACLNN_ERR_INNER_NULLPTR);
CHECK_RET(groupKeyCacheIdx != nullptr, ACLNN_ERR_INNER_NULLPTR);
CHECK_RET(blockSize > 0, ACLNN_ERR_PARAM_INVALID);

auto slotMappingContiguous = l0op::Contiguous(slotMapping, uniqueExecutor.get());
CHECK_RET(slotMappingContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);

auto ret = l0op::StoreKvBlockMetadata(slotMappingContiguous, groupLen, groupKeyIdx, groupKeyCacheIdx,blockSize,
uniqueExecutor.get());
CHECK_RET(ret != nullptr, ACLNN_ERR_INNER_NULLPTR);

*workspaceSize = 0;
uniqueExecutor.ReleaseTo(executor);
return ACLNN_SUCCESS;
}

__attribute__((visibility("default"))) aclnnStatus aclnnStoreKvBlockMetadata(void *workspace, uint64_t workspaceSize,
aclOpExecutor *executor, aclrtStream stream)
{
L2_DFX_PHASE_2(aclnnStoreKvBlockMetadata);
return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
}

#ifdef __cplusplus
}
#endif
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
/**
* 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 should 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.
*/

#ifndef ACLNN_STORE_KV_BLOCK_METADATA_H
#define ACLNN_STORE_KV_BLOCK_METADATA_H

#include "aclnn/aclnn_base.h"

#ifdef __cplusplus
extern "C" {
#endif

__attribute__((visibility("default"))) aclnnStatus aclnnStoreKvBlockMetadataGetWorkspaceSize(
const aclTensor *slotMapping,
const aclTensor *groupLen,
const aclTensor *groupKeyIdx,
const aclTensor *groupKeyCacheIdx,
int64_t blockSize,
uint64_t *workspaceSize,
aclOpExecutor **executor);

__attribute__((visibility("default"))) aclnnStatus aclnnStoreKvBlockMetadata(void *workspace, uint64_t workspaceSize,
aclOpExecutor *executor, aclrtStream stream);

#ifdef __cplusplus
}
#endif

#endif // ACLNN_STORE_KV_BLOCK_METADATA_H
Loading
Loading