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,22 @@
add_op_to_compiled_list()
if (BUILD_OPEN_PROJECT)
target_sources(op_host_aclnnInner PRIVATE
ngram_spec_decode_def.cpp
)
endif()
add_ops_compile_options(
OP_NAME NgramSpecDecode
OPTIONS
--cce-auto-sync=on
-Wno-deprecated-declarations
)
if (NOT BUILD_OPS_RTY_KERNEL)
add_modules_sources(OPTYPE ngram_spec_decode ACLNNTYPE aclnn_inner)
target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}
)
endif()

View File

@@ -0,0 +1,72 @@
#include "register/op_def_registry.h"
namespace ops {
class NgramSpecDecode : public OpDef {
public:
explicit NgramSpecDecode(const char *name) : OpDef(name)
{
this->Input("tokenIds")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Input("numTokensNoSpec")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Input("sampledTokenIds")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Input("discardRequestMask")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Attr("vocab_size").Int();
this->Attr("min_n").Int();
this->Attr("max_n").Int();
this->Attr("k").Int();
this->Output("nextTokenIds")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Output("draftTokenIds")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Output("numValidDraftTokens")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
OpAICoreConfig aicore_config;
aicore_config.DynamicCompileStaticFlag(true)
.DynamicFormatFlag(true)
.DynamicRankSupportFlag(true)
.DynamicShapeSupportFlag(true)
.NeedCheckSupportFlag(false)
.PrecisionReduceFlag(true)
.ExtendCfgInfo("aclnnSupport.value", "support_aclnn")
.ExtendCfgInfo("jitCompile.flag", "static_true")
.ExtendCfgInfo("multiKernelSupportDynamicGraph.value", "multi_kernel");
this->AICore().AddConfig("ascend910b", aicore_config);
this->AICore().AddConfig("ascend910_93", aicore_config);
}
};
OP_ADD(NgramSpecDecode);
} // namespace ops

View File

@@ -0,0 +1,118 @@
#include <cstdint>
#include <algorithm>
#include "log/ops_log.h"
#include "graph/utils/type_utils.h"
#include "register/op_def_registry.h"
#include "../op_kernel/ngram_spec_decode.h"
#include "tiling/platform/platform_ascendc.h"
#include "platform/platform_infos_def.h"
using namespace ge;
namespace {
constexpr uint32_t INPUT_TOKEN_IDS_INDEX = 0;
constexpr uint32_t INPUT_NUM_TOKENS_INDEX = 1;
constexpr uint32_t INPUT_SAMPLED_INDEX = 2;
constexpr uint32_t INPUT_DISCARD_INDEX = 3;
constexpr uint32_t ATTR_VOCAB_SIZE_INDEX = 0;
constexpr uint32_t ATTR_MIN_N_INDEX = 1;
constexpr uint32_t ATTR_MAX_N_INDEX = 2;
constexpr uint32_t ATTR_K_INDEX = 3;
constexpr int64_t ELEM_SIZE = 4; // int32
} // namespace
namespace optiling {
static ge::graphStatus NgramSpecDecodeTilingFunc(gert::TilingContext *context)
{
const char *nodeName = context->GetNodeName();
NgramSpecDecodeTilingData *tilingData = context->GetTilingData<NgramSpecDecodeTilingData>();
OPS_CHECK(tilingData == nullptr,
OPS_LOG_E(nodeName, "tilingData is nullptr."), return ge::GRAPH_FAILED);
auto attrs = context->GetAttrs();
OPS_CHECK(attrs == nullptr,
OPS_LOG_E(nodeName, "attrs is nullptr."), return ge::GRAPH_FAILED);
auto vocabSizePtr = attrs->GetAttrPointer<int64_t>(static_cast<int>(ATTR_VOCAB_SIZE_INDEX));
auto minNPtr = attrs->GetAttrPointer<int64_t>(static_cast<int>(ATTR_MIN_N_INDEX));
auto maxNPtr = attrs->GetAttrPointer<int64_t>(static_cast<int>(ATTR_MAX_N_INDEX));
auto kPtr = attrs->GetAttrPointer<int64_t>(static_cast<int>(ATTR_K_INDEX));
OPS_CHECK(vocabSizePtr == nullptr, OPS_LOG_E(nodeName, "vocabSizePtr is null."), return ge::GRAPH_FAILED);
OPS_CHECK(minNPtr == nullptr, OPS_LOG_E(nodeName, "minNPtr is null."), return ge::GRAPH_FAILED);
OPS_CHECK(maxNPtr == nullptr, OPS_LOG_E(nodeName, "maxNPtr is null."), return ge::GRAPH_FAILED);
OPS_CHECK(kPtr == nullptr, OPS_LOG_E(nodeName, "kPtr is null."), return ge::GRAPH_FAILED);
int64_t vocab_size = *vocabSizePtr;
int64_t min_n = *minNPtr;
int64_t max_n = *maxNPtr;
int64_t k = *kPtr;
const gert::StorageShape *tokenIdsShape = context->GetInputShape(INPUT_TOKEN_IDS_INDEX);
const gert::StorageShape *sampledShape = context->GetInputShape(INPUT_SAMPLED_INDEX);
OPS_CHECK(tokenIdsShape == nullptr, OPS_LOG_E(nodeName, "tokenIdsShape is null."), return ge::GRAPH_FAILED);
OPS_CHECK(sampledShape == nullptr, OPS_LOG_E(nodeName, "sampledShape is null."), return ge::GRAPH_FAILED);
int64_t batch_size = tokenIdsShape->GetStorageShape().GetDim(0);
int64_t max_seq_len = tokenIdsShape->GetStorageShape().GetDim(1);
int64_t max_new_tokens = sampledShape->GetStorageShape().GetDim(1);
auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
uint64_t ubSize = 0UL;
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
int64_t ub_size_limit = static_cast<int64_t>(ubSize);
int64_t align_elems = 32 / ELEM_SIZE;
int64_t max_seq_len_align = ((max_seq_len + align_elems - 1) / align_elems) * align_elems;
int64_t max_new_tokens_align = ((max_new_tokens + align_elems - 1) / align_elems) * align_elems;
int64_t k_align = ((k + align_elems - 1) / align_elems) * align_elems;
int64_t ub_per_row = (max_seq_len_align + max_new_tokens_align + k_align) * ELEM_SIZE;
int64_t ub_overhead = 4 * 32 + static_cast<int64_t>(max_n) * ELEM_SIZE
+ ((max_seq_len_align + 7) / 8); // maskBuf
int64_t ub_available = ub_size_limit - ub_overhead;
int64_t max_block_rows = (ub_available > 0) ? (ub_available / ub_per_row) : 1;
max_block_rows = std::max(max_block_rows, static_cast<int64_t>(1));
int64_t block_dim = std::min(batch_size, static_cast<int64_t>(aivNum));
int64_t rows_per_core = (block_dim > 0) ? (batch_size / block_dim) : 0;
int64_t former_num = (block_dim > 0) ? (block_dim - 1) : 0;
int64_t tail_rows = batch_size - former_num * rows_per_core;
int64_t block_rows = std::min(rows_per_core, max_block_rows);
tilingData->ngramInfo.batchSize = static_cast<uint32_t>(batch_size);
tilingData->ngramInfo.maxSeqLen = static_cast<uint32_t>(max_seq_len);
tilingData->ngramInfo.maxNewTokens = static_cast<uint32_t>(max_new_tokens);
tilingData->ngramInfo.vocabSize = static_cast<uint32_t>(vocab_size);
tilingData->ngramInfo.minN = static_cast<uint32_t>(min_n);
tilingData->ngramInfo.maxN = static_cast<uint32_t>(max_n);
tilingData->ngramInfo.k = static_cast<uint32_t>(k);
tilingData->ngramInfo.formerNum = static_cast<uint32_t>(former_num);
tilingData->ngramInfo.rowsPerCore = static_cast<uint32_t>(rows_per_core);
tilingData->ngramInfo.tailRows = static_cast<uint32_t>(tail_rows);
tilingData->ngramInfo.blockRows = static_cast<uint32_t>(block_rows);
context->SetBlockDim(static_cast<uint32_t>(block_dim));
OPS_LOG_D(nodeName, "batchSize=%lu, maxSeqLen=%lu, maxNewTokens=%lu, k=%lu, blockDim=%lu, blockRows=%lu",
batch_size, max_seq_len, max_new_tokens, k, block_dim, block_rows);
return ge::GRAPH_SUCCESS;
}
struct NgramSpecDecodeCompileInfo {};
ge::graphStatus TilingParseForNgramSpecDecode(gert::TilingParseContext *context)
{
(void)context;
return ge::GRAPH_SUCCESS;
}
IMPL_OP_OPTILING(NgramSpecDecode)
.Tiling(NgramSpecDecodeTilingFunc)
.TilingParse<NgramSpecDecodeCompileInfo>(TilingParseForNgramSpecDecode);
} // namespace optiling

View File

@@ -0,0 +1,26 @@
#ifndef NGRAM_SPEC_DECODE_TILING_H
#define NGRAM_SPEC_DECODE_TILING_H
#include "kernel_tiling/kernel_tiling.h"
struct NgramSpecDecodeInfo {
uint32_t batchSize;
uint32_t maxSeqLen;
uint32_t maxNewTokens;
uint32_t vocabSize;
uint32_t minN;
uint32_t maxN;
uint32_t k;
uint32_t formerNum;
uint32_t rowsPerCore;
uint32_t tailRows;
uint32_t blockRows;
};
struct NgramSpecDecodeTilingData {
Mc2InitTiling mc2InitTiling;
Mc2CcTiling mc2CcTiling1;
NgramSpecDecodeInfo ngramInfo;
};
#endif // NGRAM_SPEC_DECODE_TILING_H

View File

@@ -0,0 +1,73 @@
#include <string.h>
#include "graph/types.h"
#include "aclnn_ngram_spec_decode.h"
enum NnopbaseHcclServerType {
NNOPBASE_HCCL_SERVER_TYPE_AICPU = 0,
NNOPBASE_HCCL_SERVER_TYPE_MTE,
NNOPBASE_HCCL_SERVER_TYPE_END
};
extern "C" void __attribute__((weak)) NnopbaseSetHcclServerType(void *executor, NnopbaseHcclServerType sType);
#ifdef __cplusplus
extern "C" {
#endif
extern aclnnStatus aclnnInnerNgramSpecDecodeGetWorkspaceSize(
const aclTensor *tokenIds,
const aclTensor *numTokensNoSpec,
const aclTensor *sampledTokenIds,
const aclTensor *discardRequestMask,
int64_t vocabSize,
int64_t minN,
int64_t maxN,
int64_t k,
const aclTensor *nextTokenIds,
const aclTensor *draftTokenIds,
const aclTensor *numValidDraftTokens,
uint64_t *workspaceSize,
aclOpExecutor **executor);
extern aclnnStatus aclnnInnerNgramSpecDecode(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
aclrtStream stream);
aclnnStatus aclnnNgramSpecDecodeGetWorkspaceSize(
const aclTensor *tokenIds,
const aclTensor *numTokensNoSpec,
const aclTensor *sampledTokenIds,
const aclTensor *discardRequestMask,
int64_t vocabSize,
int64_t minN,
int64_t maxN,
int64_t k,
const aclTensor *nextTokenIds,
const aclTensor *draftTokenIds,
const aclTensor *numValidDraftTokens,
uint64_t *workspaceSize,
aclOpExecutor **executor)
{
return aclnnInnerNgramSpecDecodeGetWorkspaceSize(
tokenIds, numTokensNoSpec, sampledTokenIds, discardRequestMask,
vocabSize, minN, maxN, k,
nextTokenIds, draftTokenIds, numValidDraftTokens,
workspaceSize, executor);
}
aclnnStatus aclnnNgramSpecDecode(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
aclrtStream stream)
{
if (NnopbaseSetHcclServerType) {
NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);
}
return aclnnInnerNgramSpecDecode(workspace, workspaceSize, executor, stream);
}
#ifdef __cplusplus
}
#endif

View File

@@ -0,0 +1,56 @@
#ifndef ACLNN_NGRAM_SPEC_DECODE_H_
#define ACLNN_NGRAM_SPEC_DECODE_H_
#include "aclnn/acl_meta.h"
#ifdef __cplusplus
extern "C" {
#endif
/* function: aclnnNgramSpecDecodeGetWorkspaceSize
* tokenIds : required, [batch_size, max_seq_len], int32
* numTokensNoSpec : required, [batch_size], int32
* sampledTokenIds : required, [batch_size, max_new_tokens], int32
* discardRequestMask : required, [batch_size], int32
* vocabSize : required, int
* minN : required, int
* maxN : required, int
* k : required, int
* nextTokenIds : required, [batch_size], int32
* draftTokenIds : required, [batch_size, k], int32
* numValidDraftTokens : required, [batch_size], int32
* workspaceSize : size of workspace(output).
* executor : executor context(output).
*/
__attribute__((visibility("default"))) aclnnStatus aclnnNgramSpecDecodeGetWorkspaceSize(
const aclTensor *tokenIds,
const aclTensor *numTokensNoSpec,
const aclTensor *sampledTokenIds,
const aclTensor *discardRequestMask,
int64_t vocabSize,
int64_t minN,
int64_t maxN,
int64_t k,
const aclTensor *nextTokenIds,
const aclTensor *draftTokenIds,
const aclTensor *numValidDraftTokens,
uint64_t *workspaceSize,
aclOpExecutor **executor);
/* function: aclnnNgramSpecDecode
* workspace : workspace memory addr(input).
* workspaceSize : size of workspace(input).
* executor : executor context(input).
* stream : acl stream.
*/
__attribute__((visibility("default"))) aclnnStatus aclnnNgramSpecDecode(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
aclrtStream stream);
#ifdef __cplusplus
}
#endif
#endif // ACLNN_NGRAM_SPEC_DECODE_H_