EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

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

22
csrc/moe/copy_and_expand_eagle_inputs/op_host/CMakeLists.txt Normal file
View File

@@ -0,0 +1,22 @@
add_op_to_compiled_list()
if (BUILD_OPEN_PROJECT)
target_sources(op_host_aclnn PRIVATE
copy_and_expand_eagle_inputs_def.cpp
)
endif()
add_ops_compile_options(
OP_NAME CopyAndExpandEagleInputs
OPTIONS
--cce-auto-sync=on
-Wno-deprecated-declarations
)
if (NOT BUILD_OPS_RTY_KERNEL)
add_modules_sources(OPTYPE copy_and_expand_eagle_inputs ACLNNTYPE aclnn)
target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}
)
endif()

88
csrc/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_def.cpp Normal file
View File

@@ -0,0 +1,88 @@
/**
* @file copy_and_expand_eagle_inputs_def.cpp
* @brief CopyAndExpandEagleInputs OpDef registration
*/
#include "register/op_def_registry.h"
namespace ops {
class CopyAndExpandEagleInputs : public OpDef {
public:
explicit CopyAndExpandEagleInputs(const char* name) : OpDef(name)
{
// -------------------- Inputs --------------------
this->Input("target_token_ids")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Input("target_positions")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Input("next_token_ids")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Input("query_start_loc")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Input("query_end_loc")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
// -------------------- Outputs --------------------
this->Output("out_input_ids")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Output("out_positions")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Output("out_is_rejected_token_mask")
.ParamType(REQUIRED)
.DataType({ge::DT_INT8})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Output("out_is_masked_token_mask")
.ParamType(REQUIRED)
.DataType({ge::DT_INT8})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Output("out_new_token_indices")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Output("out_hidden_state_mapping")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
// -------------------- Attributes --------------------
this->Attr("padding_token_id").Int();
this->Attr("parallel_drafting_token_id").Int();
this->Attr("num_padding_slots_per_request").Int();
this->Attr("shift_input_ids").Bool();
this->Attr("total_input_tokens").Int();
// -------------------- Platform --------------------
this->AICore().AddConfig("ascend910b");
this->AICore().AddConfig("ascend910_93");
}
};
OP_ADD(CopyAndExpandEagleInputs);
} // namespace ops

111
csrc/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_infershape.cpp Normal file
View File

@@ -0,0 +1,111 @@
/**
* @file copy_and_expand_eagle_inputs_infershape.cpp
* @brief InferShape and InferDataType for CopyAndExpandEagleInputs
*/
#include "register/op_def_registry.h"
#include "log/ops_log.h"
#define unlikely(x) __builtin_expect((x), 0)
#define OP_CHECK_NULL_WITH_CONTEXT(context, ptr) \
do { \
if (unlikely((ptr) == nullptr)) { \
const char* name = (unlikely(((context) == nullptr) || (context)->GetNodeName() == nullptr)) ? \
"nil" : \
(context)->GetNodeName(); \
OPS_LOG_E(name, "%s is nullptr!", #ptr); \
return ge::GRAPH_FAILED; \
} \
} while (0)
static constexpr int IDX_TARGET_TOKEN_IDS = 0;
static constexpr int IDX_TARGET_POSITIONS = 1;
static constexpr int IDX_NEXT_TOKEN_IDS = 2;
static constexpr int IDX_QUERY_START_LOC = 3;
static constexpr int IDX_QUERY_END_LOC = 4;
static constexpr int OUT_INPUT_IDS = 0;
static constexpr int OUT_POSITIONS = 1;
static constexpr int OUT_REJECTED_MASK = 2;
static constexpr int OUT_MASKED_MASK = 3;
static constexpr int OUT_NEW_TOKEN_INDICES = 4;
static constexpr int OUT_HIDDEN_STATE_MAPPING = 5;
static constexpr int OUTPUT_NUM = 6;
static constexpr int ATTR_NUM_PADDING_SLOTS = 2;
static constexpr int ATTR_SHIFT_INPUT_IDS = 3;
using namespace ge;
namespace ops {
static ge::graphStatus InferShape4CopyAndExpandEagleInputs(gert::InferShapeContext* context)
{
// Get input shapes
const gert::Shape* targetTokenIdsShape = context->GetInputShape(IDX_TARGET_TOKEN_IDS);
OP_CHECK_NULL_WITH_CONTEXT(context, targetTokenIdsShape);
const gert::Shape* queryStartLocShape = context->GetInputShape(IDX_QUERY_START_LOC);
OP_CHECK_NULL_WITH_CONTEXT(context, queryStartLocShape);
// Derive dimensions from input shapes
int64_t totalInputTokens = targetTokenIdsShape->GetDim(0);
int64_t numReqs = queryStartLocShape->GetDim(0) - 1;
// Get num_padding_slots_per_request and shift_input_ids from attributes.
auto attrs = context->GetAttrs();
OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
int64_t numPaddingSlotsPerReq = *(attrs->GetAttrPointer<int64_t>(ATTR_NUM_PADDING_SLOTS));
bool shiftInputIds = *(attrs->GetAttrPointer<bool>(ATTR_SHIFT_INPUT_IDS));
// total_input_tokens already includes rejected tokens. The only delta between
// shift=false and shift=true is whether each request keeps its first accepted
// token in the draft sequence.
int64_t totalDraftTokens = totalInputTokens +
(numPaddingSlotsPerReq - (shiftInputIds ? 1 : 0)) * numReqs;
// Get and validate all output shapes
gert::Shape* outShapes[OUTPUT_NUM];
for (int i = 0; i < OUTPUT_NUM; ++i) {
outShapes[i] = context->GetOutputShape(i);
OP_CHECK_NULL_WITH_CONTEXT(context, outShapes[i]);
outShapes[i]->SetDimNum(1);
}
// out_input_ids, out_positions, out_rejected_mask, out_masked_mask: [total_draft_tokens]
outShapes[OUT_INPUT_IDS]->SetDim(0, totalDraftTokens);
outShapes[OUT_POSITIONS]->SetDim(0, totalDraftTokens);
outShapes[OUT_REJECTED_MASK]->SetDim(0, totalDraftTokens);
outShapes[OUT_MASKED_MASK]->SetDim(0, totalDraftTokens);
// out_new_token_indices: [num_reqs * num_padding_slots_per_request]
outShapes[OUT_NEW_TOKEN_INDICES]->SetDim(0, numReqs * numPaddingSlotsPerReq);
// out_hidden_state_mapping: [total_input_tokens]
outShapes[OUT_HIDDEN_STATE_MAPPING]->SetDim(0, totalInputTokens);
return GRAPH_SUCCESS;
}
static ge::graphStatus InferDataType4CopyAndExpandEagleInputs(gert::InferDataTypeContext* context)
{
// out_input_ids: INT32
context->SetOutputDataType(OUT_INPUT_IDS, DT_INT32);
// out_positions: INT32
context->SetOutputDataType(OUT_POSITIONS, DT_INT32);
// out_is_rejected_token_mask: INT8
context->SetOutputDataType(OUT_REJECTED_MASK, DT_INT8);
// out_is_masked_token_mask: INT8
context->SetOutputDataType(OUT_MASKED_MASK, DT_INT8);
// out_new_token_indices: INT32
context->SetOutputDataType(OUT_NEW_TOKEN_INDICES, DT_INT32);
// out_hidden_state_mapping: INT32
context->SetOutputDataType(OUT_HIDDEN_STATE_MAPPING, DT_INT32);
return GRAPH_SUCCESS;
}
IMPL_OP_INFERSHAPE(CopyAndExpandEagleInputs)
.InferShape(InferShape4CopyAndExpandEagleInputs)
.InferDataType(InferDataType4CopyAndExpandEagleInputs);
} // namespace ops

123
csrc/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_tiling.cpp Normal file
View File

@@ -0,0 +1,123 @@
/**
* @file copy_and_expand_eagle_inputs_tiling.cpp
* @brief CopyAndExpandEagleInputs TilingFunc implementation
*/
#include "copy_and_expand_eagle_inputs_tiling.h"
#include "tiling_base/error_log.h"
#include "register/op_def_registry.h"
#include "log/ops_log.h"
#include <algorithm>
namespace optiling {
static void GetCompileParameters(
gert::TilingContext* context, uint32_t& coreNum)
{
auto ptrCompileInfo = reinterpret_cast<const CopyAndExpandEagleInputsCompileInfo*>(context->GetCompileInfo());
if (ptrCompileInfo == nullptr) {
auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
coreNum = ascendcPlatform.GetCoreNum();
} else {
coreNum = ptrCompileInfo->totalCoreNum;
}
}
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
OPS_LOG_I(context, "Enter TilingFunc for CopyAndExpandEagleInputs");
OPS_LOG_D(context, "TilingFunc running.");
// ========== 1. Get hardware core count ==========
uint32_t coreNum;
GetCompileParameters(context, coreNum);
// ========== 2. Derive num_reqs from query_start_loc shape ==========
// query_start_loc is the 4th input (index 3), shape [num_reqs + 1]
auto queryStartLocShape = context->GetInputShape(3);
uint32_t numReqs = 0;
if (queryStartLocShape != nullptr &&
queryStartLocShape->GetStorageShape().GetDimNum() > 0) {
int64_t dim0 = queryStartLocShape->GetStorageShape().GetDim(0);
numReqs = (dim0 > 1) ? static_cast<uint32_t>(dim0 - 1) : 0;
}
// ========== 3. Get operator attributes ==========
auto attrs = context->GetAttrs();
int32_t paddingTokenId = *(attrs->GetAttrPointer<int32_t>(0));
int32_t parallelDraftingTokenId = *(attrs->GetAttrPointer<int32_t>(1));
int32_t numPaddingSlotsPerReq = *(attrs->GetAttrPointer<int32_t>(2));
bool shiftInputIds = *(attrs->GetAttrPointer<bool>(3));
int32_t totalInputTokens = *(attrs->GetAttrPointer<int32_t>(4));
// ========== 4. Compute core distribution ==========
// Adjacent requests do not guarantee block-aligned output boundaries.
// The runtime write path emits block-granular stores, so request-parallel
// execution can cause neighboring segments to overlap at block tails.
// Run this op on a single core to preserve correctness.
uint32_t usedCoreNum = 1;
uint32_t reqsPerCore = numReqs / usedCoreNum;
uint32_t remainderReqs = numReqs % usedCoreNum;
// ========== 5. Set tiling_key ==========
context->SetTilingKey(1);
// ========== 6. Get output shape ==========
uint32_t totalDraftTokens = 0;
auto outShape = context->GetOutputShape(0);
if (outShape != nullptr &&
outShape->GetStorageShape().GetDimNum() > 0) {
totalDraftTokens = static_cast<uint32_t>(outShape->GetStorageShape().GetDim(0));
}
// ========== 7. Fill TilingData ==========
CopyAndExpandEagleInputsTilingData tiling;
tiling.set_usedCoreNum(usedCoreNum);
tiling.set_numReqs(numReqs);
tiling.set_reqsPerCore(reqsPerCore);
tiling.set_remainderReqs(remainderReqs);
tiling.set_paddingTokenId(paddingTokenId);
tiling.set_parallelDraftingTokenId(parallelDraftingTokenId);
tiling.set_numPaddingSlotsPerReq(static_cast<uint32_t>(numPaddingSlotsPerReq));
tiling.set_totalInputTokens(static_cast<uint32_t>(totalInputTokens));
tiling.set_shiftInputIds(shiftInputIds ? 1u : 0u);
tiling.set_totalDraftTokens(totalDraftTokens);
tiling.SaveToBuffer(
context->GetRawTilingData()->GetData(),
context->GetRawTilingData()->GetCapacity());
context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
// ========== 8. Set block_dim ==========
context->SetBlockDim(usedCoreNum);
OPS_LOG_I(context, "Block Dim: %u", usedCoreNum);
OPS_LOG_I(context,
"numReqs: %u, reqsPerCore: %u, remainderReqs: %u, totalInputTokens: %d, totalDraftTokens: %u",
numReqs, reqsPerCore, remainderReqs, totalInputTokens, totalDraftTokens);
return ge::GRAPH_SUCCESS;
}
static ge::graphStatus TilingPrepare4CopyAndExpandEagleInputs(gert::TilingParseContext* context)
{
OPS_LOG_D(context, "TilingPrepare4CopyAndExpandEagleInputs running.");
OPS_LOG_I(context, "TilingPrepare4CopyAndExpandEagleInputs running.");
auto compileInfo = context->GetCompiledInfo<CopyAndExpandEagleInputsCompileInfo>();
OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
auto platformInfo = context->GetPlatformInfo();
OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
compileInfo->totalCoreNum = ascendcPlatform.GetCoreNum();
return ge::GRAPH_SUCCESS;
}
IMPL_OP_OPTILING(CopyAndExpandEagleInputs)
.Tiling(TilingFunc)
.TilingParse<CopyAndExpandEagleInputsCompileInfo>(TilingPrepare4CopyAndExpandEagleInputs);
} // namespace optiling

37
csrc/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_tiling.h Normal file
View File

@@ -0,0 +1,37 @@
#ifndef COPY_AND_EXPAND_EAGLE_INPUTS_TILING_H
#define COPY_AND_EXPAND_EAGLE_INPUTS_TILING_H
#include "register/tilingdata_base.h"
#include "tiling_base/error_log.h"
#include "register/op_impl_registry.h"
#include "tiling/platform/platform_ascendc.h"
namespace optiling {
BEGIN_TILING_DATA_DEF(CopyAndExpandEagleInputsTilingData)
// ---- 分核参数 ----
TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum); // 实际使用的核数
TILING_DATA_FIELD_DEF(uint32_t, numReqs); // 总请求数
TILING_DATA_FIELD_DEF(uint32_t, reqsPerCore); // 每核基础请求数
TILING_DATA_FIELD_DEF(uint32_t, remainderReqs); // 余数（前 remainder 个核多处理 1 个请求）
// ---- 算子属性 ----
TILING_DATA_FIELD_DEF(int32_t, paddingTokenId); // 填充 token ID
TILING_DATA_FIELD_DEF(int32_t, parallelDraftingTokenId); // 并行推测解码 token ID
TILING_DATA_FIELD_DEF(uint32_t, numPaddingSlotsPerReq); // 每个请求的 padding 槽位数
TILING_DATA_FIELD_DEF(uint32_t, totalInputTokens); // 输入 token 总数（用于 clamp）
TILING_DATA_FIELD_DEF(uint32_t, shiftInputIds); // 0 = false, 1 = true
// ---- 输出尺寸 ----
TILING_DATA_FIELD_DEF(uint32_t, totalDraftTokens); // 输出 token 总数
END_TILING_DATA_DEF;
struct CopyAndExpandEagleInputsCompileInfo {
uint32_t totalCoreNum = 0;
};
REGISTER_TILING_DATA_CLASS(CopyAndExpandEagleInputs, CopyAndExpandEagleInputsTilingData)
} // namespace optiling
#endif // COPY_AND_EXPAND_EAGLE_INPUTS_TILING_H