EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Revert "[Kernel] add custom moe ops for prefill" (#4806)

Browse Source 
Reverts vllm-project/vllm-ascend#4194 as it broke CI in
https://github.com/vllm-project/vllm-ascend/actions/runs/20030369087/job/57437687382?pr=4791

Co-authored-by: wangxiyuan <wangxiyuan1007@gmail.com>
This commit is contained in:
Mengqing Cao
2025-12-08 23:20:32 +08:00
committed by GitHub
parent 432b861cae
commit 7e70da9fb7
39 changed files with 2 additions and 5562 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
csrc/dispatch_layout/op_host/CMakeLists.txt
View File

@@ -1,49 +0,0 @@
# 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 1.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 DispatchLayout
OPTIONS --cce-auto-sync=off
-Wno-deprecated-declarations
-Werror
)
target_sources(op_host_aclnnInner PRIVATE
dispatch_layout.cpp
)
target_sources(opapi PRIVATE
aclnn_dispatch_layout.cpp
)
if (NOT BUILD_OPEN_PROJECT)
target_sources(aclnn_ops_train PRIVATE
aclnn_dispatch_layout.cpp
)
target_sources(aclnn_ops_infer PRIVATE
aclnn_dispatch_layout.cpp
)
endif ()
target_sources(optiling PRIVATE
dispatch_layout_tiling.cpp
)
target_include_directories(optiling PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}
)
target_sources(opsproto PRIVATE)
file(GLOB _GMM_Aclnn_header "${CMAKE_CURRENT_SOURCE_DIR}/aclnn_dispatch_layout.h")
install(FILES ${_GMM_Aclnn_header}
DESTINATION ${ACLNN_INC_INSTALL_DIR} OPTIONAL
)

64
csrc/dispatch_layout/op_host/aclnn_dispatch_layout.cpp
View File

@@ -1,64 +0,0 @@
#include <string.h>
#include "graph/types.h"
#include "aclnn_dispatch_layout.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 aclnnInnerDispatchLayoutGetWorkspaceSize(
const aclTensor *topkIdx,
int64_t numTokens,
int64_t numRanks,
int64_t numExperts,
int64_t numTopk,
const aclTensor *numTokensPerRank,
const aclTensor *numTokensPerExpert,
const aclTensor *isTokenInRank,
uint64_t *workspaceSize,
aclOpExecutor **executor);
extern aclnnStatus aclnnInnerDispatchLayout(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
aclrtStream stream);
aclnnStatus aclnnDispatchLayoutGetWorkspaceSize(
const aclTensor *topkIdx,
int64_t numTokens,
int64_t numRanks,
int64_t numExperts,
int64_t numTopk,
const aclTensor *numTokensPerRank,
const aclTensor *numTokensPerExpert,
const aclTensor *isTokenInRank,
uint64_t *workspaceSize,
aclOpExecutor **executor)
{
return aclnnInnerDispatchLayoutGetWorkspaceSize(topkIdx, numTokens, numRanks, numExperts, numTopk, numTokensPerRank,
numTokensPerExpert, isTokenInRank, workspaceSize, executor);
}
aclnnStatus aclnnDispatchLayout(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
aclrtStream stream)
{
if (NnopbaseSetHcclServerType) {
NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);
}
return aclnnInnerDispatchLayout(workspace, workspaceSize, executor, stream);
}
#ifdef __cplusplus
}
#endif

50
csrc/dispatch_layout/op_host/aclnn_dispatch_layout.h
View File

@@ -1,50 +0,0 @@
#ifndef ACLNN_DISPATCH_LAYOUT_H_
#define ACLNN_DISPATCH_LAYOUT_H_
#include "aclnn/acl_meta.h"
#ifdef __cplusplus
extern "C" {
#endif
/* funtion: aclnnDispatchLayoutGetWorkspaceSize
* topkIdx : required
* numTokens : required
* numRanks : required
* numExperts : required
* numTopk : required
* numTokensPerRank : required
* numTokensPerExpert : required
* isTokenInRank : required
* workspaceSize : size of workspace(output).
* executor : executor context(output).
*/
__attribute__((visibility("default"))) aclnnStatus aclnnDispatchLayoutGetWorkspaceSize(
const aclTensor *topkIdx,
int64_t numTokens,
int64_t numRanks,
int64_t numExperts,
int64_t numTopk,
const aclTensor *numTokensPerRank,
const aclTensor *numTokensPerExpert,
const aclTensor *isTokenInRank,
uint64_t *workspaceSize,
aclOpExecutor **executor);
/* funtion: aclnnDispatchLayout
* workspace : workspace memory addr(input).
* workspaceSize : size of workspace(input).
* executor : executor context(input).
* stream : acl stream.
*/
__attribute__((visibility("default"))) aclnnStatus aclnnDispatchLayout(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
aclrtStream stream);
#ifdef __cplusplus
}
#endif
#endif

51
csrc/dispatch_layout/op_host/dispatch_layout.cpp
View File

@@ -1,51 +0,0 @@
#include "register/op_def_registry.h"
namespace ops {
class DispatchLayout : public OpDef {
public:
explicit DispatchLayout(const char *name) : OpDef(name)
{
this->Input("topkIdx")
.ParamType(REQUIRED)
.DataType({ge::DT_INT64})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Attr("num_tokens").Int();
this->Attr("num_ranks").Int();
this->Attr("num_experts").Int();
this->Attr("num_topk").Int();
this->Output("numTokensPerRank")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Output("numTokensPerExpert")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32})
.Format({ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND});
this->Output("isTokenInRank")
.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("ascend910_93", aicore_config);
}
};
OP_ADD(DispatchLayout);
} // namespace ops

211
csrc/dispatch_layout/op_host/dispatch_layout_tiling.cpp
View File

@@ -1,211 +0,0 @@
#include <queue>
#include <vector>
#include <dlfcn.h>
#include <fcntl.h>
#include <cstdio>
#include <cstdlib>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <cmath>
#include <cstdint>
#include <string>
#include "log/ops_log.h"
#include "graph/utils/type_utils.h"
#include "register/op_def_registry.h"
#include "../op_kernel/dispatch_layout_tiling.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling/hccl/hccl_tiling.h"
#include "experiment/platform/platform/platform_infos_def.h"
using namespace ge;
namespace {
constexpr uint32_t INPUT_TOPK_IDX_INDEX = 0;
constexpr uint32_t OUTPUT_NUM_TOKEN_PER_RANK_INDEX = 0;
constexpr uint32_t OUTPUT_NUM_TOKEN_PER_EXPERT_INDEX = 1;
constexpr uint32_t OUTPUT_IS_TOKEN_IN_RANK_INDEX = 2;
constexpr uint32_t ATTR_NUM_TOKENS_INDEX = 0;
constexpr uint32_t ATTR_NUM_RANKS_INDEX = 1;
constexpr uint32_t ATTR_NUM_EXPERTS_INDEX = 2;
constexpr uint32_t ATTR_NUM_TOPK_INDEX = 3;
const int64_t MAX_COMM_WORLD_SIZE = 384;
const int64_t MAX_MOE_EXPERTS_NUM = 384;
constexpr uint32_t SYSTEM_NEED_WORKSPACE = 16 * 1024 * 1024;
constexpr uint32_t KERNEL_USE_WORKSPACE = 1 * 1024 * 1024;
constexpr uint32_t KERNEL_A2_ARG_SIZE = 1 * 1024 * 1024;
constexpr uint32_t TWO_DIMS = 2;
constexpr uint32_t K_MAX = 16;
} // namespace
namespace optiling {
static void PrintTilingDataInfo(const char *nodeName, DispatchLayoutTilingData &tilingData)
{
OPS_LOG_D(nodeName, "numToken is %u.", tilingData.dispatchLayoutInfo.numTokens);
OPS_LOG_D(nodeName, "numRanks is %u.", tilingData.dispatchLayoutInfo.numRanks);
OPS_LOG_D(nodeName, "numExperts is %u.", tilingData.dispatchLayoutInfo.numExperts);
OPS_LOG_D(nodeName, "numTopk is %u.", tilingData.dispatchLayoutInfo.numTopk);
OPS_LOG_D(nodeName, "totalUbSize is %lu.", tilingData.dispatchLayoutInfo.totalUbSize);
}
static ge::graphStatus GetAttrAndSetTilingData(gert::TilingContext *context, const char *nodeName,
DispatchLayoutTilingData &tilingData)
{
auto attrs = context->GetAttrs();
OPS_CHECK(attrs == nullptr, OPS_LOG_E(nodeName, "attrs is nullptr."), return ge::GRAPH_FAILED);
auto numTokensPtr = attrs->GetAttrPointer<int64_t>(static_cast<int>(ATTR_NUM_TOKENS_INDEX));
auto numRanksPtr = attrs->GetAttrPointer<int64_t>(static_cast<int>(ATTR_NUM_RANKS_INDEX));
auto numExpertsPtr = attrs->GetAttrPointer<int64_t>(ATTR_NUM_EXPERTS_INDEX);
auto numTopkPtr = attrs->GetAttrPointer<int64_t>(static_cast<int>(ATTR_NUM_TOPK_INDEX));
OPS_CHECK(numTokensPtr == nullptr, OPS_LOG_E(nodeName, "numTokensPtr is null."), return ge::GRAPH_FAILED);
OPS_CHECK(numRanksPtr == nullptr, OPS_LOG_E(nodeName, "numRanksPtr is null."), return ge::GRAPH_FAILED);
OPS_CHECK(numExpertsPtr == nullptr, OPS_LOG_E(nodeName, "numExpertsPtr is null."), return ge::GRAPH_FAILED);
OPS_CHECK(numTopkPtr == nullptr, OPS_LOG_E(nodeName, "numTopkPtr is null."), return ge::GRAPH_FAILED);
OPS_CHECK((*numRanksPtr <= 0) || (*numRanksPtr > MAX_COMM_WORLD_SIZE),
OPS_LOG_E(nodeName, "rankSize is invalid, only support (0, %ld], but got rankSize=%ld.", MAX_COMM_WORLD_SIZE, *numRanksPtr),
return ge::GRAPH_FAILED);
OPS_CHECK((*numExpertsPtr <= 0) || (*numExpertsPtr > MAX_MOE_EXPERTS_NUM),
OPS_LOG_E(nodeName, "numExperts is invalid, only support (0, %ld], but got numExperts=%ld.", MAX_MOE_EXPERTS_NUM, *numExpertsPtr),
return ge::GRAPH_FAILED);
OPS_CHECK((*numTopkPtr <= 0) || (*numTopkPtr > K_MAX),
OPS_LOG_E(nodeName, "numTopkPtr is invalid, only support (0, %u], but got numTopk=%ld.", K_MAX, *numTopkPtr),
return ge::GRAPH_FAILED);
tilingData.dispatchLayoutInfo.numTokens = static_cast<uint32_t>(*numTokensPtr);
tilingData.dispatchLayoutInfo.numRanks = static_cast<uint32_t>(*numRanksPtr);
tilingData.dispatchLayoutInfo.numExperts = static_cast<uint32_t>(*numExpertsPtr);
tilingData.dispatchLayoutInfo.numTopk = static_cast<uint32_t>(*numTopkPtr);
return ge::GRAPH_SUCCESS;
}
static ge::graphStatus SetWorkSpace(gert::TilingContext *context, const char *nodeName)
{
size_t *workSpaces = context->GetWorkspaceSizes(1);
OPS_CHECK(workSpaces == nullptr, OPS_LOG_E(nodeName, "workSpaces is nullptr."), return ge::GRAPH_FAILED);
workSpaces[0] = SYSTEM_NEED_WORKSPACE + KERNEL_USE_WORKSPACE + KERNEL_A2_ARG_SIZE;
return ge::GRAPH_SUCCESS;
}
static bool CheckTensorDataType(gert::TilingContext *context, const char *nodeName)
{
auto topkIdx = context->GetInputDesc(INPUT_TOPK_IDX_INDEX);
auto numTokensPerRank = context->GetOutputDesc(OUTPUT_NUM_TOKEN_PER_RANK_INDEX);
auto numTokensPerExpert = context->GetOutputDesc(OUTPUT_NUM_TOKEN_PER_EXPERT_INDEX);
auto isTokenInRank = context->GetOutputDesc(OUTPUT_IS_TOKEN_IN_RANK_INDEX);
OPS_CHECK(topkIdx == nullptr, OPS_LOG_E(nodeName, "topkIdx is null."), return false);
OPS_CHECK(numTokensPerRank == nullptr, OPS_LOG_E(nodeName, "numTokensPerRank is null."), return false);
OPS_CHECK(numTokensPerExpert == nullptr, OPS_LOG_E(nodeName, "numTokensPerExpert is null."), return false);
OPS_CHECK(isTokenInRank == nullptr, OPS_LOG_E(nodeName, "isTokenInRank is null."), return false);
OPS_CHECK((topkIdx->GetDataType() != ge::DT_INT64),
OPS_LOG_E(nodeName, "topkIdx datatype is invalid, datatype should be int, but is %d.",
static_cast<ge::DataType>(topkIdx->GetDataType())), return false);
OPS_CHECK((numTokensPerRank->GetDataType() != ge::DT_INT32),
OPS_LOG_E(nodeName, "numTokensPerRank datatype is invalid, datatype should be int, but is %d.",
static_cast<ge::DataType>(numTokensPerRank->GetDataType())), return false);
OPS_CHECK((numTokensPerExpert->GetDataType() != ge::DT_INT32),
OPS_LOG_E(nodeName, "numTokensPerExpert datatype is invalid, datatype should be int, but is %d.",
static_cast<ge::DataType>(numTokensPerExpert->GetDataType())), return false);
OPS_CHECK((isTokenInRank->GetDataType() != ge::DT_INT32),
OPS_LOG_E(nodeName, "isTokenInRank datatype is invalid, datatype should be int, but is %d.",
static_cast<ge::DataType>(isTokenInRank->GetDataType())), return false);
return true;
}
static bool CheckTensorShape(gert::TilingContext *context, const char *nodeName)
{
const gert::StorageShape *topkIdxStorageShape = context->GetInputShape(INPUT_TOPK_IDX_INDEX);
int64_t topkIdxDim0 = topkIdxStorageShape->GetStorageShape().GetDim(0);
int64_t topkIdxDim1 = topkIdxStorageShape->GetStorageShape().GetDim(1);
OPS_CHECK((topkIdxStorageShape->GetStorageShape().GetDimNum() != TWO_DIMS),
OPS_LOG_E(nodeName, "topkIdx must be 2-dimension, but get %lu dim.",
topkIdxStorageShape->GetStorageShape().GetDimNum()), return false);
return true;
}
static ge::graphStatus TilingCheckTensor(
gert::TilingContext *context, const char *nodeName)
{
OPS_CHECK(!CheckTensorDataType(context, nodeName),
OPS_LOG_E(nodeName, "params dataType is invalid."),
return ge::GRAPH_FAILED);
OPS_CHECK(!CheckTensorShape(context, nodeName),
OPS_LOG_E(nodeName, "params dataType is invalid."),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
static ge::graphStatus DispatchLayoutTilingFuncImpl(gert::TilingContext *context)
{
const char *nodeName = context->GetNodeName();
DispatchLayoutTilingData *tilingData = context->GetTilingData<DispatchLayoutTilingData>();
OPS_CHECK(tilingData == nullptr, OPS_LOG_E(nodeName, "tilingData is nullptr."), return ge::GRAPH_FAILED);
OPS_LOG_I(nodeName, "Enter NotifyDispatch tiling check func.");
OPS_CHECK(GetAttrAndSetTilingData(context, nodeName, *tilingData) != ge::GRAPH_SUCCESS,
OPS_LOG_E(nodeName, "Get attr and set tiling data failed."),
return ge::GRAPH_FAILED);
OPS_CHECK(TilingCheckTensor(context, nodeName) != ge::GRAPH_SUCCESS,
OPS_LOG_E(nodeName, "Tiling check param failed."),
return ge::GRAPH_FAILED);
OPS_CHECK(SetWorkSpace(context, nodeName) != ge::GRAPH_SUCCESS,
OPS_LOG_E(nodeName, "Tiling set workspace failed."),
return ge::GRAPH_FAILED);
fe::PlatFormInfos *platformInfoPtr = context->GetPlatformInfo();
fe::PlatFormInfos &platformInfo = *platformInfoPtr;
std::string socVersion;
(void)platformInfo.GetPlatformResWithLock("version", "Short_SoC_version", socVersion);
auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
uint32_t blockDim;
uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
uint64_t ubSize = 0UL;
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
blockDim = aivNum;
context->SetBlockDim(blockDim);
tilingData->dispatchLayoutInfo.totalUbSize = ubSize;
OPS_LOG_D(nodeName, "blockDim=%u, aivNum=%u, ubSize=%lu", blockDim, aivNum, ubSize);
PrintTilingDataInfo(nodeName, *tilingData);
return ge::GRAPH_SUCCESS;
}
static ge::graphStatus DispatchLayoutTilingFunc(gert::TilingContext *context)
{
fe::PlatFormInfos *platformInfoPtr = context->GetPlatformInfo();
fe::PlatFormInfos &platformInfo = *platformInfoPtr;
std::string socVersion;
ge::graphStatus ret;
ret = DispatchLayoutTilingFuncImpl(context);
return ret;
}
struct DispatchLayoutCompileInfo {};
ge::graphStatus TilingParseForDispatchLayout(gert::TilingParseContext *context)
{
(void)context;
return ge::GRAPH_SUCCESS;
}
IMPL_OP_OPTILING(DispatchLayout)
.Tiling(DispatchLayoutTilingFunc)
.TilingParse<DispatchLayoutCompileInfo>(TilingParseForDispatchLayout);
} // namespace optiling