init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

csrc/attention/store_kv_block/CMakeLists.txt

@@ -0,0 +1,19 @@
+# -----------------------------------------------------------------------------------------------------------
+# 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(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
+if(NOT ENABLE_TEST AND NOT BENCHMARK)
+    list(REMOVE_ITEM CURRENT_DIRS tests)
+endif()
+foreach(SUB_DIR ${CURRENT_DIRS})
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
+        add_subdirectory(${SUB_DIR})
+    endif()
+endforeach()

csrc/attention/store_kv_block/op_host/CMakeLists.txt

@@ -0,0 +1,30 @@
+# This program is free software, you can redistribute it and/or modify it.
+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
+# This file is a part of the CANN Open Software.
+# Licensed under 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.
+# ======================================================================================================================
+add_op_to_compiled_list()
+
+if (BUILD_OPEN_PROJECT)
+    target_sources(op_host_aclnn PRIVATE
+        store_kv_block_def.cpp
+    )
+endif()
+
+add_ops_compile_options(
+    OP_NAME StoreKVBlock
+    OPTIONS
+        --cce-auto-sync=off
+        -Wno-deprecated-declarations
+)
+
+if (NOT BUILD_OPS_RTY_KERNEL)
+    add_modules_sources(OPTYPE store_kv_block ACLNNTYPE aclnn)
+    target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
+        ${CMAKE_CURRENT_SOURCE_DIR}
+    )
+endif()
+

csrc/attention/store_kv_block/op_host/store_kv_block_def.cpp

@@ -0,0 +1,64 @@
+/**
+ * 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 store_kv_block_def.cpp
+ * \brief Operator definition for StoreKVBlock
+ */
+#include "register/op_def_registry.h"
+
+namespace ops {
+class StoreKVBlock : public OpDef {
+ public:
+  explicit StoreKVBlock(const char* name) : OpDef(name) {
+    this->Input("keyIn")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT8, ge::DT_FLOAT16, ge::DT_BF16})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .AutoContiguous();
+    this->Input("keyCacheIn")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT8, ge::DT_FLOAT16, ge::DT_BF16})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .AutoContiguous();
+    this->Input("groupLen")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32 , ge::DT_INT32 , ge::DT_INT32 })
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .AutoContiguous();
+    this->Input("groupKeyIdx")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32 , ge::DT_INT32 , ge::DT_INT32 })
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .AutoContiguous();
+    this->Input("groupKeyCacheIdx")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32 , ge::DT_INT32 , ge::DT_INT32 })
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .AutoContiguous();
+    this->Attr("blockSize").Int();
+    this->AICore().AddConfig("ascend910b");
+    this->AICore().AddConfig("ascend910_93");
+  }
+};
+
+OP_ADD(StoreKVBlock);
+}  // namespace ops

csrc/attention/store_kv_block/op_host/store_kv_block_infershape.cpp

@@ -0,0 +1,33 @@
+/**
+ * This program is free software, you can redistribute it and/or modify.
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under 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 store_kv_block_infershape.cpp
+ * \brief InferShape implementation for StoreKVBlock
+ */
+#include <graph/utils/type_utils.h>
+#include <register/op_impl_registry.h>
+#include "error/ops_error.h"
+
+using namespace ge;
+namespace ops {
+
+static ge::graphStatus InferShape4StoreKVBlock(gert::InferShapeContext* context)
+{
+    return GRAPH_SUCCESS;
+}
+
+static graphStatus InferDataType4StoreKVBlock(gert::InferDataTypeContext* context)
+{
+    return GRAPH_SUCCESS;
+}
+
+IMPL_OP_INFERSHAPE(StoreKVBlock).InferShape(InferShape4StoreKVBlock).InferDataType(InferDataType4StoreKVBlock);
+} // namespace ops

csrc/attention/store_kv_block/op_host/store_kv_block_tiling.cpp

@@ -0,0 +1,156 @@
+/**
+ * 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.
+ */
+
+#include "store_kv_block_tiling.h"
+#include "register/op_def_registry.h"
+#include "tiling/platform/platform_ascendc.h"
+#include "tiling_base/error_log.h"
+
+namespace optiling {
+
+constexpr uint32_t DIM_0 = 0;
+constexpr uint32_t DIM_1 = 1;
+constexpr uint32_t DIM_2 = 2;
+constexpr int32_t MAX_UB_USE_SIZE = 180 * 1024;
+
+struct StoreKVBlockParams {
+    uint32_t numTokens{0};
+    uint32_t numCache{0};
+    uint32_t numHeads{1};
+    uint32_t headSize[5]{1, 1, 1, 1, 1};
+    uint32_t blockTableSize{0};
+    uint32_t typeByte{0};
+    uint32_t tokenSize{1};
+    uint32_t tilingKey{1};
+    uint64_t workspaceSize{0};
+    uint32_t groupInfoLen{0};
+    uint32_t corepernum{0};
+    uint32_t coretail{0};
+    uint64_t sysWorkspaceSize{0};
+    uint32_t coreNum{0};
+};
+
+static ge::graphStatus DoCommonTiling(gert::TilingContext* context, StoreKVBlockParams& params) {
+    auto kShape = context->GetInputShape(DIM_0);
+    auto kDimNum = kShape->GetStorageShape().GetDimNum();
+    if (kDimNum < 2 || kDimNum > 7) {
+        OP_LOGE(context->GetNodeName(), "StoreKVBlock Input kDimNum dim < 2 || kDimNum>7");
+        return ge::GRAPH_FAILED;
+    }
+
+    for (int i = 0; i < kDimNum; i++) {
+        if (i == 0) params.numTokens = static_cast<uint32_t>(kShape->GetStorageShape().GetDim(i));
+        else if (i == 1) params.numHeads = static_cast<uint32_t>(kShape->GetStorageShape().GetDim(i));
+        else if (static_cast<uint32_t>(kShape->GetStorageShape().GetDim(i)) != 0)
+            params.headSize[i - 2] = static_cast<uint32_t>(kShape->GetStorageShape().GetDim(i));
+    }
+
+    auto kCacheShape = context->GetInputShape(DIM_1);
+    auto kCacheDimNum = kCacheShape->GetStorageShape().GetDimNum();
+    if (kCacheDimNum < 2 || kCacheDimNum > 7) {
+        OP_LOGE(context->GetNodeName(), "StoreKVBlock Input kCacheDimNum < 2");
+        return ge::GRAPH_FAILED;
+    }
+    params.numCache = kCacheShape->GetStorageShape().GetDim(0) * kCacheShape->GetStorageShape().GetDim(1);
+
+    const int64_t* blockSizePtr = context->GetAttrs()->GetInt(0);
+    uint32_t blockSize = static_cast<uint32_t>(*blockSizePtr);
+    params.tokenSize = params.numHeads * params.headSize[0] * params.headSize[1] * params.headSize[2] * params.headSize[3] * params.headSize[4];
+    params.blockTableSize = blockSize;
+
+    uint32_t typeByte = 0;
+    auto xDataType = context->GetInputDesc(DIM_0)->GetDataType();
+    if (xDataType == ge::DataType::DT_INT8) {
+        typeByte = sizeof(int8_t);
+        params.tilingKey = 1;
+    } else if (xDataType == ge::DataType::DT_FLOAT16 || xDataType == ge::DataType::DT_BF16) {
+        typeByte = sizeof(uint16_t);
+        params.tilingKey = 2;
+    } else if (xDataType == ge::DataType::DT_INT32 || xDataType == ge::DataType::DT_UINT32) {
+        typeByte = sizeof(uint32_t);
+        params.tilingKey = 4;
+    } else {
+        OP_LOGE(context->GetNodeName(), "Unsupported type.");
+        return ge::GRAPH_FAILED;
+    }
+
+    params.typeByte = typeByte;
+
+    auto groupInfoShape = context->GetInputShape(DIM_2);
+    params.groupInfoLen = static_cast<uint32_t>(groupInfoShape->GetStorageShape().GetDim(0));
+    params.corepernum = params.groupInfoLen / params.coreNum;
+    params.coretail = params.groupInfoLen % params.coreNum;
+
+    uint32_t pageBlockEleSize = params.blockTableSize * params.tokenSize;
+    if (pageBlockEleSize > MAX_UB_USE_SIZE) {
+        OP_LOGE(context->GetNodeName(), "pageBlockEleSize > MaxUBSize");
+        return ge::GRAPH_FAILED;
+    }
+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus StoreKVBlockTilingFunc(gert::TilingContext* context) {
+    StoreKVBlockParams params;
+
+    auto platformInfo = context->GetPlatformInfo();
+    OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);
+    auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
+    params.coreNum = ascendcPlatform.GetCoreNum();
+    if (params.coreNum == 0) {
+        OP_LOGE(context->GetNodeName(), "Failed to get core num.");
+        return ge::GRAPH_FAILED;
+    }
+    params.sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
+
+    auto ret = DoCommonTiling(context, params);
+    if (ret != ge::GRAPH_SUCCESS) {
+        return ret;
+    }
+
+    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);
+    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);     
+    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);
+
+    size_t* workspaceSize = context->GetWorkspaceSizes(1);
+    *workspaceSize = params.workspaceSize + params.sysWorkspaceSize;
+    context->SetTilingKey(params.tilingKey);
+    if (params.coreNum > 0) context->SetBlockDim(params.coreNum);
+
+    tilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
+    context->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
+
+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus TilingParseForStoreKVBlock(gert::TilingParseContext* context) {
+    return ge::GRAPH_SUCCESS;
+}
+
+IMPL_OP_OPTILING(StoreKVBlock)
+    .Tiling(StoreKVBlockTilingFunc)
+    .TilingParse<StoreKVBlockCompileInfo>(TilingParseForStoreKVBlock);
+
+} // namespace optiling

csrc/attention/store_kv_block/op_host/store_kv_block_tiling.h

@@ -0,0 +1,23 @@
+#include "register/tilingdata_base.h"
+
+namespace optiling {
+BEGIN_TILING_DATA_DEF(StoreKVBlockTilingData)
+    TILING_DATA_FIELD_DEF(uint32_t, blockTableSize);
+    TILING_DATA_FIELD_DEF(uint32_t, typeByte);
+    TILING_DATA_FIELD_DEF(uint32_t, tokenSize);
+    TILING_DATA_FIELD_DEF(uint32_t, corePerNum);
+    TILING_DATA_FIELD_DEF(uint32_t, coreTail);
+    TILING_DATA_FIELD_DEF(uint32_t, numTokens);
+    TILING_DATA_FIELD_DEF(uint32_t, numCache);
+    TILING_DATA_FIELD_DEF(uint32_t, groupInfoLen);
+END_TILING_DATA_DEF;
+
+REGISTER_TILING_DATA_CLASS(StoreKVBlock, StoreKVBlockTilingData)
+
+struct StoreKVBlockCompileInfo {
+    uint32_t coreNum;
+    uint64_t ubSizePlatForm;
+    uint32_t sysWorkspaceSize;
+};
+
+} // namespace optiling

csrc/attention/store_kv_block/op_kernel/store_kv_block.cpp

@@ -0,0 +1,43 @@
+/**
+ * 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 store_kv_block.cpp
+ * \brief Kernel entry for StoreKVBlock operator
+ */
+#include "store_kv_block.h"
+
+extern "C" __global__ __aicore__ void store_kv_block(
+    GM_ADDR keyIn, GM_ADDR keyCacheIn, GM_ADDR groupLen, GM_ADDR groupKeyIdx, GM_ADDR groupKeyCacheIdx, GM_ADDR workspace, GM_ADDR tiling)
+{
+    AscendC::TPipe pipe;
+    REGISTER_TILING_DEFAULT(StoreKVBlock::StoreKVBlockTilingData);
+    GET_TILING_DATA(tilingData, tiling);
+
+    if (TILING_KEY_IS(1)) {
+        StoreKVBlock::StoreKVBlockBase<uint8_t> op;
+        op.Init( &pipe, &tilingData);
+        op.Process(keyIn,keyCacheIn, groupLen, groupKeyIdx, groupKeyCacheIdx);
+    } else if (TILING_KEY_IS(2)) {
+        StoreKVBlock::StoreKVBlockBase<half> op;
+        op.Init( &pipe, &tilingData);
+        op.Process(keyIn,keyCacheIn, groupLen, groupKeyIdx, groupKeyCacheIdx);
+    } else if (TILING_KEY_IS(4)) {
+        StoreKVBlock::StoreKVBlockBase<int32_t> op;
+        op.Init( &pipe, &tilingData);
+        op.Process(keyIn,keyCacheIn, groupLen, groupKeyIdx, groupKeyCacheIdx);
+    }
+}

csrc/attention/store_kv_block/op_kernel/store_kv_block.h

@@ -0,0 +1,133 @@
+/**
+ * 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 store_kv_block.h
+ * \brief StoreKVBlock kernel operator
+ */
+
+#ifndef ASCEND_STORE_KV_BLOCK_H
+#define ASCEND_STORE_KV_BLOCK_H
+
+#include "kernel_operator.h"
+
+namespace StoreKVBlock {
+using namespace AscendC;
+
+
+#ifndef STORE_KV_BLOCK_TILING_DATA_H_
+#define STORE_KV_BLOCK_TILING_DATA_H_
+struct StoreKVBlockTilingData{
+    uint32_t blockTableSize;    
+    uint32_t typeByte;
+    uint32_t tokenSize;
+    uint32_t corePerNum;
+    uint32_t coreTail;
+    uint32_t numTokens;
+    uint32_t numCache;
+    uint32_t groupInfoLen;
+
+};
+#endif
+template <typename T>
+class StoreKVBlockBase {
+public:
+
+    uint32_t tokenSize = 0;
+    uint32_t tokenByteSize = 0;
+    uint32_t blockTableSize = 0;
+    uint32_t typeByte = 0;
+    uint32_t numTokens = 0;
+    uint32_t numCache = 0;
+    uint32_t groupInfoLen = 0;
+
+    uint32_t coreId = 0;
+    uint32_t coreTail = 0;
+    uint32_t corePerNum = 0;
+    uint32_t blockNum = 0;
+    AscendC::TPipe* pipeThis;
+    AscendC::LocalTensor<T> tokenLocal;
+    AscendC::GlobalTensor<T> keyInputGt;
+    AscendC::GlobalTensor<T> keyCacheInputGt;
+    AscendC::GlobalTensor<int32_t> groupLenGt;
+    AscendC::GlobalTensor<int32_t> groupKeyIdxGt;
+    AscendC::GlobalTensor<int32_t> groupKeyCacheIdxGt;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> tokenBuf;
+    __aicore__ inline StoreKVBlockBase() {}
+
+    __aicore__ inline uint32_t RoundUp(uint32_t x, uint32_t y = 16)
+    {
+        return y == 0 ? 0 : (x + y - 1) / y * y;
+    }
+
+    __aicore__ inline void Init( AscendC::TPipe *pipe, StoreKVBlockTilingData *tilingData)
+    {
+        pipeThis = pipe;
+        typeByte = tilingData->typeByte;
+        tokenSize = tilingData->tokenSize;
+        tokenByteSize = tokenSize*typeByte;
+        blockTableSize = tilingData->blockTableSize;
+        numTokens = tilingData->numTokens;
+        numCache = tilingData->numCache;
+        groupInfoLen = tilingData->groupInfoLen;
+
+        coreId = AscendC::GetBlockIdx();
+        coreTail = tilingData->coreTail;
+        blockNum = AscendC::GetBlockNum();
+        if (coreId < coreTail){
+            // Not all cores have corePerNum+1 items; only coreTail cores get one extra.
+            // If corePerNum is 0, cores beyond coreTail have no work and will not access any address.
+            corePerNum = tilingData->corePerNum+1;
+        }else {
+            corePerNum = tilingData->corePerNum;
+        }
+    }
+    __aicore__ inline void Process(GM_ADDR keyIn, GM_ADDR keyCacheIn, GM_ADDR groupLen, GM_ADDR groupKeyIdx, GM_ADDR groupKeyCacheIdx)
+    {
+        
+        keyInputGt.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(keyIn));
+        keyCacheInputGt.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(keyCacheIn));
+        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 (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;
+            }
+           
+            copyParams.blockLen = groupLenGt.GetValue(idx)*tokenByteSize; // in bytes
+            DataCopyPad(tokenLocal, keyInputGt[ groupKeyIdxGt.GetValue(idx)*tokenSize], copyParams, padParams); // note: offset order
+            AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE3>(EVENT_ID1);
+            AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE3>(EVENT_ID1);
+            DataCopyPad(keyCacheInputGt[groupKeyCacheIdxGt.GetValue(idx)*tokenSize], tokenLocal, copyParams);
+            AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);
+            AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);
+        }
+
+    }
+
+};
+}
+
+#endif

csrc/attention/store_kv_block/store_kv_block_torch_adpt.h

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. 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.
+ */
+//  #include "../aclnn_torch_adapter/op_api_common.h"
+
+#ifndef STORE_KV_BLOCK_TORCH_ADPT_H
+#define STORE_KV_BLOCK_TORCH_ADPT_H
+#include <climits>  
+namespace vllm_ascend {
+
+void store_kv_block(
+    const at::Tensor &key_in,
+    const at::Tensor &key_cache_in,
+    const at::Tensor &group_len,
+    const at::Tensor &group_key_idx,
+    const at::Tensor &group_key_cache_idx,
+    int64_t block_size)
+{
+
+    EXEC_NPU_CMD(aclnnStoreKVBlock, key_in, key_cache_in,group_len, group_key_idx, group_key_cache_idx, block_size);
+    
+} 
+
+}
+#endif