init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions

View File

@@ -0,0 +1,33 @@
# -----------------------------------------------------------------------------------------------------------
# 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_op_to_compiled_list()
if (BUILD_OPEN_PROJECT)
target_sources(op_host_aclnn PRIVATE
compressor_metadata_def.cpp
)
endif()
add_ops_compile_options(
OP_NAME CompressorMetadata
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 compressor_metadata ACLNNTYPE aclnn)
target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}
)
endif()

View File

@@ -0,0 +1,70 @@
/*
* Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
*/
#include "register/op_def_registry.h"
namespace ops {
class CompressorMetadata : public OpDef {
public:
explicit CompressorMetadata(const char* name) : OpDef(name)
{
this->Input("ropeCos")
.ParamType(REQUIRED)
.DataType({ge::DT_FLOAT, 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("ropeSin")
.ParamType(REQUIRED)
.DataType({ge::DT_FLOAT, 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("cuSeqlens")
.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("startPos")
.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("kvBlockTable")
.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->Output("compressCos")
.ParamType(REQUIRED)
.DataType({ge::DT_FLOAT, 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});
this->Output("compressSin")
.ParamType(REQUIRED)
.DataType({ge::DT_FLOAT, 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});
this->Output("slotMapping")
.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});
this->Attr("kvBlockSize").Int();
this->Attr("slotMappingFormat").Int();
this->Attr("cmpRatio").Int();
this->Attr("actualNumReqs").Int();
this->AICore().AddConfig("ascend910b");
this->AICore().AddConfig("ascend910_93");
this->AICore().AddConfig("ascend950");
}
};
OP_ADD(CompressorMetadata);
} // namespace ops

View File

@@ -0,0 +1,289 @@
/*
* Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
*/
#include "compressor_metadata_tiling.h"
#include <algorithm>
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling_base/error_log.h"
namespace optiling {
namespace {
constexpr uint32_t ROPE_COS_INDEX = 0;
constexpr uint32_t ROPE_SIN_INDEX = 1;
constexpr uint32_t CU_SEQLENS_INDEX = 2;
constexpr uint32_t START_POS_INDEX = 3;
constexpr uint32_t KV_BLOCK_TABLE_INDEX = 4;
constexpr uint32_t COMPRESS_COS_INDEX = 0;
constexpr uint32_t COMPRESS_SIN_INDEX = 1;
constexpr uint32_t SLOT_MAPPING_INDEX = 2;
constexpr uint32_t SLOT_MAPPING_FLAT = 1;
constexpr uint32_t SLOT_MAPPING_BLOCK_OFFSET = 2;
constexpr int64_t MAX_UINT32_VALUE = 0xFFFFFFFFLL;
constexpr int64_t MAX_INT32_VALUE = 0x7FFFFFFFLL;
constexpr uint32_t TILING_KEY_FLOAT = 1;
constexpr uint32_t TILING_KEY_FLOAT16 = 2;
constexpr uint32_t TILING_KEY_BF16 = 3;
constexpr uint32_t ALIGN_BYTES = 32;
constexpr uint32_t BUFFER_NUM = 2;
constexpr uint32_t MAX_TILE_ROWS = 512;
constexpr uint32_t MAX_DATACOPY_BLOCK_COUNT = 4095;
constexpr uint32_t ROWS_PER_CORE_TARGET = 64;
constexpr uint32_t UB_RESERVED_BYTES = 16 * 1024;
uint32_t AlignUp(uint64_t value, uint32_t align)
{
return static_cast<uint32_t>((value + align - 1) / align * align);
}
uint32_t CeilDiv(uint64_t lhs, uint64_t rhs)
{
return static_cast<uint32_t>((lhs + rhs - 1) / rhs);
}
} // namespace
static ge::graphStatus CompressorMetadataTilingFunc(gert::TilingContext* context)
{
auto platformInfo = context->GetPlatformInfo();
OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
uint32_t aivCoreNum = ascendcPlatform.GetCoreNumAiv();
if (aivCoreNum == 0) {
aivCoreNum = ascendcPlatform.GetCoreNum();
}
if (aivCoreNum == 0) {
OP_LOGE(context->GetNodeName(), "Failed to get AIV core num.");
return ge::GRAPH_FAILED;
}
uint64_t ubSize = 0;
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
if (ubSize == 0) {
OP_LOGE(context->GetNodeName(), "Failed to get UB size.");
return ge::GRAPH_FAILED;
}
auto outputShape = context->GetOutputShape(COMPRESS_COS_INDEX);
auto compressSinShape = context->GetOutputShape(COMPRESS_SIN_INDEX);
auto slotMappingShape = context->GetOutputShape(SLOT_MAPPING_INDEX);
OP_CHECK_NULL_WITH_CONTEXT(context, outputShape);
OP_CHECK_NULL_WITH_CONTEXT(context, compressSinShape);
OP_CHECK_NULL_WITH_CONTEXT(context, slotMappingShape);
auto outputDimNum = outputShape->GetStorageShape().GetDimNum();
if (outputDimNum < 2) {
OP_LOGE(context->GetNodeName(), "compressCos dim num should be at least 2.");
return ge::GRAPH_FAILED;
}
if (compressSinShape->GetStorageShape().GetDimNum() != outputDimNum) {
OP_LOGE(context->GetNodeName(), "compressCos and compressSin dim num mismatch.");
return ge::GRAPH_FAILED;
}
for (size_t dimIdx = 0; dimIdx < outputDimNum; ++dimIdx) {
if (compressSinShape->GetStorageShape().GetDim(dimIdx) != outputShape->GetStorageShape().GetDim(dimIdx)) {
OP_LOGE(context->GetNodeName(), "compressCos and compressSin shape mismatch.");
return ge::GRAPH_FAILED;
}
}
int64_t numRows = outputShape->GetStorageShape().GetDim(0);
int64_t ropeDim = outputShape->GetStorageShape().GetDim(outputDimNum - 1);
if (numRows <= 0 || ropeDim <= 0 || numRows > MAX_UINT32_VALUE || ropeDim > MAX_UINT32_VALUE) {
OP_LOGE(context->GetNodeName(), "compressCos shape is invalid.");
return ge::GRAPH_FAILED;
}
auto ropeCosShape = context->GetInputShape(ROPE_COS_INDEX);
auto ropeSinShape = context->GetInputShape(ROPE_SIN_INDEX);
auto cuSeqlensShape = context->GetInputShape(CU_SEQLENS_INDEX);
auto startPosShape = context->GetInputShape(START_POS_INDEX);
auto kvBlockTableShape = context->GetInputShape(KV_BLOCK_TABLE_INDEX);
OP_CHECK_NULL_WITH_CONTEXT(context, ropeCosShape);
OP_CHECK_NULL_WITH_CONTEXT(context, ropeSinShape);
OP_CHECK_NULL_WITH_CONTEXT(context, cuSeqlensShape);
OP_CHECK_NULL_WITH_CONTEXT(context, startPosShape);
OP_CHECK_NULL_WITH_CONTEXT(context, kvBlockTableShape);
if (ropeCosShape->GetStorageShape().GetDimNum() != 2 || ropeSinShape->GetStorageShape().GetDimNum() != 2) {
OP_LOGE(context->GetNodeName(), "ropeCos and ropeSin should be 2D tensors.");
return ge::GRAPH_FAILED;
}
int64_t ropeRows = ropeCosShape->GetStorageShape().GetDim(0);
int64_t ropeCosDim = ropeCosShape->GetStorageShape().GetDim(1);
if (ropeRows <= 0 || ropeCosDim <= 0 || ropeRows > MAX_UINT32_VALUE || ropeCosDim != ropeDim ||
ropeSinShape->GetStorageShape().GetDim(0) != ropeRows ||
ropeSinShape->GetStorageShape().GetDim(1) != ropeCosDim) {
OP_LOGE(context->GetNodeName(), "ropeCos and ropeSin shape mismatch.");
return ge::GRAPH_FAILED;
}
int64_t cuSeqlensDim0 = cuSeqlensShape->GetStorageShape().GetDim(0);
if (cuSeqlensDim0 < 2 || cuSeqlensDim0 > MAX_UINT32_VALUE) {
OP_LOGE(context->GetNodeName(), "cuSeqlens dim0 should be at least 2.");
return ge::GRAPH_FAILED;
}
if (startPosShape->GetStorageShape().GetDimNum() != 1 ||
startPosShape->GetStorageShape().GetDim(0) <= 0 ||
startPosShape->GetStorageShape().GetDim(0) > MAX_UINT32_VALUE) {
OP_LOGE(context->GetNodeName(), "startPos should be a non-empty 1D tensor.");
return ge::GRAPH_FAILED;
}
if (kvBlockTableShape->GetStorageShape().GetDimNum() != 2) {
OP_LOGE(context->GetNodeName(), "kvBlockTable should be a 2D tensor.");
return ge::GRAPH_FAILED;
}
int64_t kvBlockTableRows = kvBlockTableShape->GetStorageShape().GetDim(0);
int64_t kvBlockTableStride = kvBlockTableShape->GetStorageShape().GetDim(1);
if (kvBlockTableRows <= 0 || kvBlockTableStride <= 0 || kvBlockTableRows > MAX_UINT32_VALUE ||
kvBlockTableStride > MAX_UINT32_VALUE) {
OP_LOGE(context->GetNodeName(), "kvBlockTable shape is invalid.");
return ge::GRAPH_FAILED;
}
auto attrs = context->GetAttrs();
OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
const int64_t* kvBlockSizePtr = attrs->GetInt(0);
const int64_t* slotMappingFormatPtr = attrs->GetInt(1);
const int64_t* cmpRatioPtr = attrs->GetInt(2);
const int64_t* actualNumReqsPtr = attrs->GetInt(3);
OP_CHECK_NULL_WITH_CONTEXT(context, kvBlockSizePtr);
OP_CHECK_NULL_WITH_CONTEXT(context, slotMappingFormatPtr);
OP_CHECK_NULL_WITH_CONTEXT(context, cmpRatioPtr);
OP_CHECK_NULL_WITH_CONTEXT(context, actualNumReqsPtr);
if (*kvBlockSizePtr <= 0 || *kvBlockSizePtr > MAX_INT32_VALUE) {
OP_LOGE(context->GetNodeName(), "kvBlockSize should be in (0, INT32_MAX].");
return ge::GRAPH_FAILED;
}
if (*cmpRatioPtr <= 0 || *cmpRatioPtr > MAX_UINT32_VALUE) {
OP_LOGE(context->GetNodeName(), "cmpRatio should be in (0, UINT32_MAX].");
return ge::GRAPH_FAILED;
}
if (*slotMappingFormatPtr != SLOT_MAPPING_FLAT && *slotMappingFormatPtr != SLOT_MAPPING_BLOCK_OFFSET) {
OP_LOGE(context->GetNodeName(), "slotMappingFormat should be 1(flat) or 2(block_offset).");
return ge::GRAPH_FAILED;
}
auto slotMappingDimNum = slotMappingShape->GetStorageShape().GetDimNum();
if ((*slotMappingFormatPtr == SLOT_MAPPING_FLAT &&
(slotMappingDimNum != 1 || slotMappingShape->GetStorageShape().GetDim(0) != numRows)) ||
(*slotMappingFormatPtr == SLOT_MAPPING_BLOCK_OFFSET &&
(slotMappingDimNum != 2 || slotMappingShape->GetStorageShape().GetDim(0) != numRows ||
slotMappingShape->GetStorageShape().GetDim(1) != 2))) {
OP_LOGE(context->GetNodeName(), "slotMapping shape does not match slotMappingFormat.");
return ge::GRAPH_FAILED;
}
if (*actualNumReqsPtr <= 0 || *actualNumReqsPtr >= cuSeqlensDim0 ||
*actualNumReqsPtr > startPosShape->GetStorageShape().GetDim(0) ||
*actualNumReqsPtr > kvBlockTableRows ||
*actualNumReqsPtr > MAX_UINT32_VALUE) {
OP_LOGE(context->GetNodeName(), "actualNumReqs is invalid.");
return ge::GRAPH_FAILED;
}
CompressorMetadataTilingData tilingData;
tilingData.set_numRows(static_cast<uint32_t>(numRows));
tilingData.set_numReqs(static_cast<uint32_t>(cuSeqlensDim0 - 1));
tilingData.set_actualNumReqs(static_cast<uint32_t>(*actualNumReqsPtr));
tilingData.set_ropeRows(static_cast<uint32_t>(ropeRows));
tilingData.set_ropeDim(static_cast<uint32_t>(ropeDim));
tilingData.set_kvBlockTableStride(static_cast<uint32_t>(kvBlockTableStride));
tilingData.set_kvBlockSize(static_cast<uint32_t>(*kvBlockSizePtr));
tilingData.set_slotMappingFormat(static_cast<uint32_t>(*slotMappingFormatPtr));
tilingData.set_cmpRatio(static_cast<uint32_t>(*cmpRatioPtr));
auto ropeDesc = context->GetInputDesc(ROPE_COS_INDEX);
auto ropeSinDesc = context->GetInputDesc(ROPE_SIN_INDEX);
auto compressCosDesc = context->GetOutputDesc(COMPRESS_COS_INDEX);
auto compressSinDesc = context->GetOutputDesc(COMPRESS_SIN_INDEX);
OP_CHECK_NULL_WITH_CONTEXT(context, ropeDesc);
OP_CHECK_NULL_WITH_CONTEXT(context, ropeSinDesc);
OP_CHECK_NULL_WITH_CONTEXT(context, compressCosDesc);
OP_CHECK_NULL_WITH_CONTEXT(context, compressSinDesc);
auto ropeDtype = ropeDesc->GetDataType();
if (ropeSinDesc->GetDataType() != ropeDtype ||
compressCosDesc->GetDataType() != ropeDtype ||
compressSinDesc->GetDataType() != ropeDtype) {
OP_LOGE(context->GetNodeName(), "rope and compress output dtypes should match.");
return ge::GRAPH_FAILED;
}
uint64_t tilingKey = 0;
uint32_t dtypeSize = 0;
if (ropeDtype == ge::DataType::DT_FLOAT) {
tilingKey = TILING_KEY_FLOAT;
dtypeSize = sizeof(float);
} else if (ropeDtype == ge::DataType::DT_FLOAT16) {
tilingKey = TILING_KEY_FLOAT16;
dtypeSize = sizeof(uint16_t);
} else if (ropeDtype == ge::DataType::DT_BF16) {
tilingKey = TILING_KEY_BF16;
dtypeSize = sizeof(uint16_t);
} else {
OP_LOGE(context->GetNodeName(), "Unsupported rope dtype.");
return ge::GRAPH_FAILED;
}
uint32_t actualNumReqs = static_cast<uint32_t>(*actualNumReqsPtr);
uint32_t cmpRatio = static_cast<uint32_t>(*cmpRatioPtr);
if (static_cast<uint64_t>(ropeDim) * dtypeSize > MAX_UINT32_VALUE ||
(static_cast<uint64_t>(actualNumReqs) + 1) * sizeof(int32_t) > MAX_UINT32_VALUE) {
OP_LOGE(context->GetNodeName(), "tiling byte size exceeds UINT32_MAX.");
return ge::GRAPH_FAILED;
}
uint32_t ropeRowBytes = static_cast<uint32_t>(ropeDim) * dtypeSize;
if (static_cast<uint64_t>(cmpRatio - 1) * ropeRowBytes > MAX_UINT32_VALUE) {
OP_LOGE(context->GetNodeName(), "rope stride exceeds UINT32_MAX.");
return ge::GRAPH_FAILED;
}
uint32_t ropeRowBytesAligned = AlignUp(ropeRowBytes, ALIGN_BYTES);
uint32_t slotCols = (*slotMappingFormatPtr == SLOT_MAPPING_FLAT) ? 1U : 2U;
uint32_t reqTableBytes = AlignUp((static_cast<uint64_t>(actualNumReqs) + 1) * sizeof(int32_t), ALIGN_BYTES);
uint64_t fixedUbBytes = static_cast<uint64_t>(reqTableBytes) * 3 + ALIGN_BYTES + UB_RESERVED_BYTES;
uint64_t rowUbBytes =
static_cast<uint64_t>(BUFFER_NUM) * ropeRowBytesAligned * 2 + slotCols * sizeof(int32_t) + sizeof(int32_t);
if (rowUbBytes > MAX_UINT32_VALUE) {
OP_LOGE(context->GetNodeName(), "row UB footprint exceeds UINT32_MAX.");
return ge::GRAPH_FAILED;
}
uint64_t minUbBytes = static_cast<uint64_t>(reqTableBytes) * 3 + ALIGN_BYTES + rowUbBytes;
if (ubSize <= minUbBytes) {
OP_LOGE(context->GetNodeName(), "UB size is insufficient for compressor metadata.");
return ge::GRAPH_FAILED;
}
uint32_t tileRows = 1;
if (ubSize > fixedUbBytes && rowUbBytes > 0) {
tileRows = static_cast<uint32_t>((ubSize - fixedUbBytes) / rowUbBytes);
tileRows = std::max(tileRows, 1U);
}
tileRows = std::min(tileRows, MAX_TILE_ROWS);
tileRows = std::min(tileRows, MAX_DATACOPY_BLOCK_COUNT);
uint32_t usedCoreNum =
std::min(aivCoreNum, std::max(1U, CeilDiv(static_cast<uint64_t>(numRows), ROWS_PER_CORE_TARGET)));
tilingData.set_usedCoreNum(usedCoreNum);
tilingData.set_tileRows(tileRows);
tilingData.set_ropeRowBytes(ropeRowBytes);
tilingData.set_ropeRowBytesAligned(ropeRowBytesAligned);
tilingData.set_slotCols(slotCols);
size_t* workspaceSize = context->GetWorkspaceSizes(1);
OP_CHECK_NULL_WITH_CONTEXT(context, workspaceSize);
*workspaceSize = 0;
context->SetBlockDim(usedCoreNum);
context->SetTilingKey(tilingKey);
auto rawTilingData = context->GetRawTilingData();
OP_CHECK_NULL_WITH_CONTEXT(context, rawTilingData);
tilingData.SaveToBuffer(rawTilingData->GetData(), rawTilingData->GetCapacity());
rawTilingData->SetDataSize(tilingData.GetDataSize());
return ge::GRAPH_SUCCESS;
}
static ge::graphStatus TilingParseForCompressorMetadata(gert::TilingParseContext* context)
{
return ge::GRAPH_SUCCESS;
}
IMPL_OP_OPTILING(CompressorMetadata)
.Tiling(CompressorMetadataTilingFunc)
.TilingParse<CompressorMetadataCompileInfo>(TilingParseForCompressorMetadata);
} // namespace optiling

View File

@@ -0,0 +1,32 @@
#ifndef COMPRESSOR_METADATA_TILING_H
#define COMPRESSOR_METADATA_TILING_H
#include "register/tilingdata_base.h"
namespace optiling {
BEGIN_TILING_DATA_DEF(CompressorMetadataTilingData)
TILING_DATA_FIELD_DEF(uint32_t, numRows);
TILING_DATA_FIELD_DEF(uint32_t, numReqs);
TILING_DATA_FIELD_DEF(uint32_t, actualNumReqs);
TILING_DATA_FIELD_DEF(uint32_t, ropeRows);
TILING_DATA_FIELD_DEF(uint32_t, ropeDim);
TILING_DATA_FIELD_DEF(uint32_t, kvBlockTableStride);
TILING_DATA_FIELD_DEF(uint32_t, kvBlockSize);
TILING_DATA_FIELD_DEF(uint32_t, slotMappingFormat);
TILING_DATA_FIELD_DEF(uint32_t, cmpRatio);
TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum);
TILING_DATA_FIELD_DEF(uint32_t, tileRows);
TILING_DATA_FIELD_DEF(uint32_t, ropeRowBytes);
TILING_DATA_FIELD_DEF(uint32_t, ropeRowBytesAligned);
TILING_DATA_FIELD_DEF(uint32_t, slotCols);
END_TILING_DATA_DEF;
REGISTER_TILING_DATA_CLASS(CompressorMetadata, CompressorMetadataTilingData)
struct CompressorMetadataCompileInfo {
uint32_t coreNum;
uint64_t ubSizePlatForm;
};
} // namespace optiling
#endif