init v0.23.0

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

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