init v0.23.0

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

csrc/attention/load_index_kv_cache/op_host/CMakeLists.txt

@@ -0,0 +1,62 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+# add_ops_compile_options(
+#         OP_NAME LoadIndexKvCache
+#         OPTIONS --cce-auto-sync=off
+#                 -Wno-deprecated-declarations
+#                 -Werror
+#                 -mllvm -cce-aicore-hoist-movemask=false
+#                 --op_relocatable_kernel_binary=true
+# )
+
+# set(load_index_kv_cache_depends transformer/attention/load_index_kv_cache PARENT_SCOPE)
+
+# target_sources(op_host_aclnn PRIVATE
+#         op_host/load_index_kv_cache_def.cpp
+# )
+
+# target_sources(optiling PRIVATE
+#         op_host/load_index_kv_cache_tiling.cpp
+# )
+
+# if (NOT BUILD_OPEN_PROJECT)
+#     target_sources(opmaster_ct PRIVATE
+#         op_host/load_index_kv_cache_tiling.cpp
+#     )
+# endif ()
+
+# target_include_directories(optiling PRIVATE
+#         ${CMAKE_CURRENT_SOURCE_DIR}/op_host
+# )
+
+# target_sources(opsproto PRIVATE
+#         op_host/load_index_kv_cache_proto.cpp
+# )
+
+add_op_to_compiled_list()
+
+if (BUILD_OPEN_PROJECT)
+    target_sources(op_host_aclnn PRIVATE
+        load_index_kv_cache_def.cpp
+     )
+endif()
+
+add_ops_compile_options(
+        OP_NAME LoadIndexKvCache
+        OPTIONS --cce-auto-sync=off
+                -Wno-deprecated-declarations
+                -mllvm -cce-aicore-hoist-movemask=false
+                --op_relocatable_kernel_binary=true
+)
+
+if (NOT BUILD_OPS_RTY_KERNEL)
+    add_modules_sources(OPTYPE load_index_kv_cache ACLNNTYPE aclnn)
+endif()

csrc/attention/load_index_kv_cache/op_host/load_index_kv_cache_def.cpp

@@ -0,0 +1,49 @@
+/**
+ * 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_def.cpp
+ * \brief
+ */
+#include "register/op_def_registry.h"
+
+namespace ops {
+class LoadIndexKvCache : public OpDef {
+public:
+    explicit LoadIndexKvCache(const char* name) : OpDef(name)
+    {
+        this->Input("kv_cache")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_FLOAT8_E4M3FN})
+            .Format({ge::FORMAT_ND})
+            .UnknownShapeFormat({ge::FORMAT_ND})
+            .IgnoreContiguous();
+        this->Input("slot_mapping")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_INT32})
+            .Format({ge::FORMAT_ND})
+            .UnknownShapeFormat({ge::FORMAT_ND});
+        this->Output("kv")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_FLOAT8_E4M3FN})
+            .Format({ge::FORMAT_ND})
+            .UnknownShapeFormat({ge::FORMAT_ND});
+        this->Output("kv_scale")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_FLOAT})
+            .Format({ge::FORMAT_ND})
+            .UnknownShapeFormat({ge::FORMAT_ND});
+        this->Attr("block_stride").AttrType(OPTIONAL).Int(0);
+        this->AICore().AddConfig("ascend950");
+    }
+};
+
+OP_ADD(LoadIndexKvCache);
+}  // namespace ops

csrc/attention/load_index_kv_cache/op_host/load_index_kv_cache_proto.cpp

@@ -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 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_proto.cpp
+ * \brief
+ */
+#include <graph/utils/type_utils.h>
+#include <register/op_impl_registry.h>
+
+#include "error/ops_error.h"
+
+using namespace ge;
+namespace ops {
+
+graphStatus InferShape4LoadIndexKvCache(gert::InferShapeContext* context)
+{
+    return ge::GRAPH_SUCCESS;
+}
+
+graphStatus InferDtype4LoadIndexKvCache(gert::InferDataTypeContext* context)
+{
+    return GRAPH_SUCCESS;
+}
+
+IMPL_OP_INFERSHAPE(LoadIndexKvCache)
+    .InferShape(InferShape4LoadIndexKvCache)
+    .InferDataType(InferDtype4LoadIndexKvCache);
+}  // namespace ops

csrc/attention/load_index_kv_cache/op_host/load_index_kv_cache_tiling.cpp

@@ -0,0 +1,204 @@
+/**
+ * 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_tiling.cpp
+ * \brief
+ */
+
+#include <sstream>
+#include "load_index_kv_cache_tiling.h"
+
+using namespace ge;
+namespace optiling {
+namespace {
+constexpr uint64_t WORKSPACE_SIZE = 32;
+int64_t CeilDiv(int64_t x, int64_t y)
+{
+    if (y != 0) {
+        return (x + y - 1) / y;
+    }
+    return x;
+}
+int64_t DownAlign(int64_t x, int64_t y) {
+    if (y == 0) {
+        return x;
+    }
+    return (x / y) * y;
+}
+int64_t RoundUp(int64_t x, int64_t y) {
+    return CeilDiv(x, y) * y;
+}
+
+constexpr int64_t INPUT_KV_CACHE_IDX = 0;
+constexpr int64_t INPUT_SLOT_MAPPING_IDX = 1;
+constexpr int64_t ATTR_BLOCK_STRIDE_INDEX = 0;
+constexpr int64_t DIM_0 = 0;
+constexpr int64_t DIM_1 = 1;
+constexpr int64_t DIM_2 = 2;
+constexpr int64_t DIM_3 = 3;
+constexpr int64_t INPUTS_DIM_LIMIT = 4;
+}
+
+ge::graphStatus LoadIndexKvCacheTiling::GetPlatformInfo()
+{
+    auto platformInfo = context_->GetPlatformInfo();
+    if (platformInfo == nullptr) {
+        auto compileInfoPtr = context_->GetCompileInfo<LoadIndexKvCacheCompileInfo>();
+        OPS_ERR_IF(compileInfoPtr == nullptr, OPS_LOG_E(context_, "compile info is null"),
+                      return ge::GRAPH_FAILED);
+        coreNum_ = compileInfoPtr->coreNum;
+        ubSize_ = compileInfoPtr->ubSize;
+    } else {
+        auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
+        coreNum_ = ascendcPlatform.GetCoreNumAiv();
+        uint64_t ubSizePlatForm;
+        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);
+        ubSize_ = ubSizePlatForm;
+        socVersion_ = ascendcPlatform.GetSocVersion();
+    }
+    return ge::GRAPH_SUCCESS;
+}
+
+ge::graphStatus LoadIndexKvCacheTiling::GetAttr()
+{
+    auto* attrs = context_->GetAttrs();
+    OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED);
+
+    auto blockStrideAttr = attrs->GetAttrPointer<int64_t>(ATTR_BLOCK_STRIDE_INDEX);
+    blockStride_ = blockStrideAttr != nullptr ? *blockStrideAttr : 0;
+    return ge::GRAPH_SUCCESS;
+}
+
+ge::graphStatus LoadIndexKvCacheTiling::GetShapeAttrsInfoInner()
+{
+    auto shapeKvCache = context_->GetInputShape(INPUT_KV_CACHE_IDX);
+    OPS_LOG_E_IF_NULL(context_, shapeKvCache, return ge::GRAPH_FAILED);
+    auto kvCacheStorageShape = shapeKvCache->GetStorageShape();
+    OPS_ERR_IF(
+        (kvCacheStorageShape.GetDimNum() != INPUTS_DIM_LIMIT),
+        OPS_LOG_E(context_, "the dim of kv_cache only support %d, please check.", INPUTS_DIM_LIMIT),
+        return ge::GRAPH_FAILED);
+
+    bn_ = kvCacheStorageShape.GetDim(DIM_0);
+    bs_ = kvCacheStorageShape.GetDim(DIM_1);
+    d_ = kvCacheStorageShape.GetDim(DIM_3);
+
+    auto shapeSlotMapping = context_->GetInputShape(INPUT_SLOT_MAPPING_IDX);
+    OPS_LOG_E_IF_NULL(context_, shapeSlotMapping, return ge::GRAPH_FAILED);
+    auto slotMappingStorageShape = shapeSlotMapping->GetStorageShape();
+    OPS_ERR_IF(
+        (slotMappingStorageShape.GetDimNum() != 1),
+        OPS_LOG_E(context_, "the dim of slot_mapping must be 1, please check."),
+        return ge::GRAPH_FAILED);
+
+    n_ = slotMappingStorageShape.GetDim(0);
+
+    OPS_ERR_IF(GetAttr() != ge::GRAPH_SUCCESS,
+                  OPS_LOG_E(context_->GetNodeName(), "get attr failed."),
+                  return ge::GRAPH_FAILED);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+ge::graphStatus LoadIndexKvCacheTiling::CheckShapesAndAttrs()
+{
+    blockStride_ = blockStride_ == 0 ? bs_ * d_ : blockStride_;
+    OPS_ERR_IF(blockStride_ < bs_ * d_,
+            OPS_LOG_E(context_, "stride_kvcache must be greater than last dim of kv_cache %ld.", d_),
+            return ge::GRAPH_FAILED);
+    // 如果blockStride_为0,则认为连续
+    return ge::GRAPH_SUCCESS;
+}
+
+
+ge::graphStatus LoadIndexKvCacheTiling::CalcOpTiling()
+{
+    // 默认尾轴在ub内全载
+    rowOfFormerBlock_ = CeilDiv(n_, static_cast<int64_t>(coreNum_));
+    usedCoreNums_ = std::min(CeilDiv(n_, rowOfFormerBlock_), static_cast<int64_t>(coreNum_));
+    rowOfTailBlock_ = n_ - (usedCoreNums_ - 1) * rowOfFormerBlock_;
+
+    tilingData_.set_bn(bn_);
+    tilingData_.set_bs(bs_);
+    tilingData_.set_d(d_);
+    tilingData_.set_n(n_);
+    tilingData_.set_rowOfFormerBlock(rowOfFormerBlock_);
+    tilingData_.set_rowOfTailBlock(rowOfTailBlock_);
+    tilingData_.set_blockStride(blockStride_);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+ge::graphStatus LoadIndexKvCacheTiling::DoOpTiling()
+{
+    if (GetPlatformInfo() == ge::GRAPH_FAILED) {
+        return ge::GRAPH_FAILED;
+    }
+
+    if (GetShapeAttrsInfoInner() == ge::GRAPH_FAILED) {
+        return ge::GRAPH_FAILED;
+    }
+
+    if (CheckShapesAndAttrs() == ge::GRAPH_FAILED) {
+        return ge::GRAPH_FAILED;
+    }
+
+    if (CalcOpTiling() == ge::GRAPH_FAILED) {
+        return ge::GRAPH_FAILED;
+    }
+
+    if (GetWorkspaceSize() == ge::GRAPH_FAILED) {
+        return ge::GRAPH_FAILED;
+    }
+
+    if (PostTiling() == ge::GRAPH_FAILED) {
+        return ge::GRAPH_FAILED;
+    }
+    context_->SetTilingKey(0);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+ge::graphStatus LoadIndexKvCacheTiling::GetWorkspaceSize()
+{
+    workspaceSize_ = WORKSPACE_SIZE;
+    return ge::GRAPH_SUCCESS;
+}
+
+ge::graphStatus LoadIndexKvCacheTiling::PostTiling()
+{
+    context_->SetBlockDim(usedCoreNums_);
+    size_t* workspaces = context_->GetWorkspaceSizes(1);
+    workspaces[0] = workspaceSize_;
+    tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
+    context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize());
+    return ge::GRAPH_SUCCESS;
+}
+
+ge::graphStatus TilingPrepareForLoadIndexKvCache(gert::TilingParseContext *context)
+{
+    (void)context;
+    return ge::GRAPH_SUCCESS;
+}
+
+ge::graphStatus TilingForLoadIndexKvCache(gert::TilingContext *context)
+{
+    OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("LoadIndexKvCache", "Tiling context is null"),
+               return ge::GRAPH_FAILED);
+    LoadIndexKvCacheTiling LoadIndexKvCacheTiling(context);
+    return LoadIndexKvCacheTiling.DoOpTiling();
+}
+
+IMPL_OP_OPTILING(LoadIndexKvCache)
+    .Tiling(TilingForLoadIndexKvCache)
+    .TilingParse<LoadIndexKvCacheCompileInfo>(TilingPrepareForLoadIndexKvCache);
+
+}  // namespace optiling

csrc/attention/load_index_kv_cache/op_host/load_index_kv_cache_tiling.h

@@ -0,0 +1,98 @@
+/**
+ * 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_tiling.h
+ * \brief
+ */
+
+#ifndef LOAD_INDEX_KV_CACHE_TILING_H
+#define LOAD_INDEX_KV_CACHE_TILING_H
+
+
+#include <vector>
+#include <iostream>
+#include "register/op_impl_registry.h"
+#include "platform/platform_infos_def.h"
+#include "exe_graph/runtime/tiling_context.h"
+#include "tiling/platform/platform_ascendc.h"
+#include "register/op_def_registry.h"
+#include "register/tilingdata_base.h"
+#include "tiling/tiling_api.h"
+#include "error/ops_error.h"
+#include "platform/platform_info.h"
+
+namespace optiling {
+// ----------公共定义----------
+struct TilingRequiredParaInfo {
+    const gert::CompileTimeTensorDesc *desc;
+    const gert::StorageShape *shape;
+};
+
+struct TilingOptionalParaInfo {
+    const gert::CompileTimeTensorDesc *desc;
+    const gert::Tensor *tensor;
+};
+
+// ----------算子TilingData定义----------
+BEGIN_TILING_DATA_DEF(LoadIndexKvCacheTilingData)
+TILING_DATA_FIELD_DEF(int64_t, bn); // block_num
+TILING_DATA_FIELD_DEF(int64_t, bs); // block_size
+TILING_DATA_FIELD_DEF(int64_t, d); // head_dim
+TILING_DATA_FIELD_DEF(int64_t, n);
+TILING_DATA_FIELD_DEF(int64_t, rowOfFormerBlock);  // 头核共需要处理多少行
+TILING_DATA_FIELD_DEF(int64_t, rowOfTailBlock);  // 尾核共需要处理多少行
+TILING_DATA_FIELD_DEF(int64_t, blockStride);
+END_TILING_DATA_DEF;
+
+REGISTER_TILING_DATA_CLASS(LoadIndexKvCache, LoadIndexKvCacheTilingData)
+
+// ----------算子CompileInfo定义----------
+struct LoadIndexKvCacheCompileInfo {
+    uint64_t coreNum = 0;
+    uint64_t ubSize = 0;
+};
+
+// ----------算子Tiling入参信息解析及check类----------
+class LoadIndexKvCacheTiling {
+public:
+    explicit LoadIndexKvCacheTiling(gert::TilingContext* tilingContext) : context_(tilingContext)
+    {
+    }
+    ~LoadIndexKvCacheTiling() = default;
+
+    ge::graphStatus GetPlatformInfo();
+    ge::graphStatus DoOpTiling();
+    ge::graphStatus GetWorkspaceSize();
+    ge::graphStatus PostTiling();
+    ge::graphStatus GetAttr();
+    ge::graphStatus GetShapeAttrsInfoInner();
+    ge::graphStatus CheckShapesAndAttrs();
+    ge::graphStatus CalcOpTiling();
+private:
+    gert::TilingContext *context_ = nullptr;
+    LoadIndexKvCacheTilingData tilingData_;
+    uint64_t coreNum_ = 0;
+    uint64_t workspaceSize_ = 0;
+    uint64_t usedCoreNums_ = 0;
+    uint64_t ubSize_ = 0;
+    int64_t bn_ = 0;
+    int64_t bs_ = 0;
+    int64_t d_ = 0;
+    int64_t n_ = 0;
+    int64_t rowOfFormerBlock_ = 0;
+    int64_t rowOfTailBlock_ = 0;
+    int64_t blockStride_ = 0;
+    platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B;
+    int64_t tilingKey_ = 0;
+};
+
+}  // namespace optiling
+#endif  // LOAD_INDEX_KV_CACHE_TILING_H