Skip to content

Commit bd8be2e

Browse files
shiro-zzzzshiro-zzzz
authored
[Kernel] Add moe normal ops (#4810)
### What this PR does / why we need it? 1.Add the implementation of normal Aclnn operators: MoeCombineNormal, MoeDispatchNormal, NotifyDispatch，and DispatchLayout. - MoeCombineNormal: Implements the combine logic within MoE operations. - MoeDispatchNormal: Implements the dispatch logic within MoE operations. - NotifyDispatch: Exchanges topk_idx information among different ranks to calculate the device memory required for the dispatch stage. - DispatchLayout: Used to calculate information related to the device memory layout for the dispatch stage. 2.Provide PyTorch interfaces for normal operators—get_dispatch_layout, dispatch_prefill, and combine_prefill—to be used for MoE communication during the prefill stage in vLLM. - get_dispatch_layout: Calculates information related to the device memory layout for the dispatch operator, and is called before dispatch_prefill. - dispatch_prefill: Initiates the dispatch operation. - combine_prefill: Initiates the combine operation. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? The functionality has already been validated using the local Qwen model. Test cases will be added after support for multi-NPU use cases in the CI pipeline is finalized. - vLLM version: v0.12.0 - vLLM main: vllm-project/vllm@ad32e3e Signed-off-by: shiro-zzzz <[email protected]>
1 parent c77dca5 commit bd8be2e

39 files changed

+5366
-5
lines changed

‎csrc/build_aclnn.sh‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then
4545
exit 1
4646
fi
4747
# for dispatch_gmm_combine_decode
48-
yes | cp "${HCCL_STRUCT_FILE_PATH}" "${ROOT_DIR}/csrc/dispatch_gmm_combine_decode/op_kernel"
48+
yes | cp "${HCCL_STRUCT_FILE_PATH}" "${ROOT_DIR}/csrc/utils/inc/kernel"
4949
# for dispatch_ffn_combine
5050
SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
5151
TARGET_DIR="$SCRIPT_DIR/dispatch_ffn_combine/op_kernel/utils/"
@@ -58,7 +58,19 @@ elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then
5858

5959
sed -i 's/struct HcclOpResParam {/struct HcclOpResParamCustom {/g' "$TARGET_FILE"
6060
sed -i 's/struct HcclRankRelationResV2 {/struct HcclRankRelationResV2Custom {/g' "$TARGET_FILE"
61-
CUSTOM_OPS="grouped_matmul_swiglu_quant_weight_nz_tensor_list;lightning_indexer;sparse_flash_attention;dispatch_ffn_combine;dispatch_gmm_combine_decode;"
61+
62+
CUSTOM_OPS_ARRAY=(
63+
"grouped_matmul_swiglu_quant_weight_nz_tensor_list"
64+
"lightning_indexer"
65+
"sparse_flash_attention"
66+
"dispatch_ffn_combine"
67+
"dispatch_gmm_combine_decode"
68+
"moe_combine_normal"
69+
"moe_dispatch_normal"
70+
"dispatch_layout"
71+
"notify_dispatch"
72+
)
73+
CUSTOM_OPS=$(IFS=';'; echo "${CUSTOM_OPS_ARRAY[*]}")
6274
SOC_ARG="ascend910_93"
6375
else
6476
# others
@@ -71,7 +83,7 @@ fi
7183
cd csrc
7284
rm -rf build output
7385
echo "building custom ops $CUSTOM_OPS for $SOC_VERSION"
74-
bash build.sh -n $CUSTOM_OPS -c $SOC_ARG
86+
bash build.sh -n "$CUSTOM_OPS" -c "$SOC_ARG"
7587

7688
# install custom ops to vllm_ascend/_cann_ops_custom
7789
./output/CANN-custom_ops*.run --install-path=$ROOT_DIR/vllm_ascend/_cann_ops_custom

‎csrc/dispatch_gmm_combine_decode/op_kernel/dispatch_gmm_combine_decode_base.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
#ifndef DISPATCH_GMM_COMBINE_DECODE_BASE_H
1111
#define DISPATCH_GMM_COMBINE_DECODE_BASE_H
1212

13-
#include "moe_distribute_base.h"
13+
#include "../common/moe_distribute_base.h"
1414

1515
#define TemplateMC2TypeClass typename ExpandXType, typename ExpandIdxType, bool IsNeedReduceScatter, uint32_t EXEC_FLAG
1616
#define TemplateMC2TypeFunc ExpandXType, ExpandIdxType, IsNeedReduceScatter, EXEC_FLAG

‎csrc/dispatch_layout/op_host/CMakeLists.txt‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
2+
# This file is a part of the CANN Open Software.
3+
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
4+
# Please refer to the License for details. You may not use this file except in compliance with the License.
5+
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
6+
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
7+
# See LICENSE in the root of the software repository for the full text of the License.
8+
# ======================================================================================================================
9+
10+
add_ops_compile_options(
11+
OP_NAME DispatchLayout
12+
OPTIONS --cce-auto-sync=off
13+
-Wno-deprecated-declarations
14+
-Werror
15+
)
16+
17+
target_sources(op_host_aclnnInner PRIVATE
18+
dispatch_layout.cpp
19+
)
20+
21+
target_sources(opapi PRIVATE
22+
aclnn_dispatch_layout.cpp
23+
)
24+
25+
if (NOT BUILD_OPEN_PROJECT)
26+
target_sources(aclnn_ops_train PRIVATE
27+
aclnn_dispatch_layout.cpp
28+
)
29+
30+
target_sources(aclnn_ops_infer PRIVATE
31+
aclnn_dispatch_layout.cpp
32+
)
33+
endif ()
34+
35+
target_sources(optiling PRIVATE
36+
dispatch_layout_tiling.cpp
37+
)
38+
39+
target_include_directories(optiling PRIVATE
40+
${CMAKE_CURRENT_SOURCE_DIR}
41+
)
42+
43+
target_sources(opsproto PRIVATE)
44+
45+
file(GLOB _GMM_Aclnn_header "${CMAKE_CURRENT_SOURCE_DIR}/aclnn_dispatch_layout.h")
46+
47+
install(FILES ${_GMM_Aclnn_header}
48+
DESTINATION ${ACLNN_INC_INSTALL_DIR} OPTIONAL
49+
)

‎csrc/dispatch_layout/op_host/aclnn_dispatch_layout.cpp‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
#include <string.h>
2+
#include "graph/types.h"
3+
#include "aclnn_dispatch_layout.h"
4+
5+
enum NnopbaseHcclServerType {
6+
NNOPBASE_HCCL_SERVER_TYPE_AICPU = 0,
7+
NNOPBASE_HCCL_SERVER_TYPE_MTE,
8+
NNOPBASE_HCCL_SERVER_TYPE_END
9+
};
10+
extern "C" void __attribute__((weak)) NnopbaseSetHcclServerType(void *executor, NnopbaseHcclServerType sType);
11+
12+
#ifdef __cplusplus
13+
extern "C" {
14+
#endif
15+
16+
extern aclnnStatus aclnnInnerDispatchLayoutGetWorkspaceSize(
17+
const aclTensor *topkIdx,
18+
int64_t numTokens,
19+
int64_t numRanks,
20+
int64_t numExperts,
21+
int64_t numTopk,
22+
const aclTensor *numTokensPerRank,
23+
const aclTensor *numTokensPerExpert,
24+
const aclTensor *isTokenInRank,
25+
uint64_t *workspaceSize,
26+
aclOpExecutor **executor);
27+
28+
extern aclnnStatus aclnnInnerDispatchLayout(
29+
void *workspace,
30+
uint64_t workspaceSize,
31+
aclOpExecutor *executor,
32+
aclrtStream stream);
33+
34+
aclnnStatus aclnnDispatchLayoutGetWorkspaceSize(
35+
const aclTensor *topkIdx,
36+
int64_t numTokens,
37+
int64_t numRanks,
38+
int64_t numExperts,
39+
int64_t numTopk,
40+
const aclTensor *numTokensPerRank,
41+
const aclTensor *numTokensPerExpert,
42+
const aclTensor *isTokenInRank,
43+
uint64_t *workspaceSize,
44+
aclOpExecutor **executor)
45+
{
46+
return aclnnInnerDispatchLayoutGetWorkspaceSize(topkIdx, numTokens, numRanks, numExperts, numTopk, numTokensPerRank,
47+
numTokensPerExpert, isTokenInRank, workspaceSize, executor);
48+
}
49+
50+
aclnnStatus aclnnDispatchLayout(
51+
void *workspace,
52+
uint64_t workspaceSize,
53+
aclOpExecutor *executor,
54+
aclrtStream stream)
55+
{
56+
if (NnopbaseSetHcclServerType) {
57+
NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);
58+
}
59+
return aclnnInnerDispatchLayout(workspace, workspaceSize, executor, stream);
60+
}
61+
62+
#ifdef __cplusplus
63+
}
64+
#endif

‎csrc/dispatch_layout/op_host/aclnn_dispatch_layout.h‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
#ifndef ACLNN_DISPATCH_LAYOUT_H_
2+
#define ACLNN_DISPATCH_LAYOUT_H_
3+
4+
#include "aclnn/acl_meta.h"
5+
6+
#ifdef __cplusplus
7+
extern "C" {
8+
#endif
9+
10+
/* funtion: aclnnDispatchLayoutGetWorkspaceSize
11+
* topkIdx : required
12+
* numTokens : required
13+
* numRanks : required
14+
* numExperts : required
15+
* numTopk : required
16+
* numTokensPerRank : required
17+
* numTokensPerExpert : required
18+
* isTokenInRank : required
19+
* workspaceSize : size of workspace(output).
20+
* executor : executor context(output).
21+
*/
22+
__attribute__((visibility("default"))) aclnnStatus aclnnDispatchLayoutGetWorkspaceSize(
23+
const aclTensor *topkIdx,
24+
int64_t numTokens,
25+
int64_t numRanks,
26+
int64_t numExperts,
27+
int64_t numTopk,
28+
const aclTensor *numTokensPerRank,
29+
const aclTensor *numTokensPerExpert,
30+
const aclTensor *isTokenInRank,
31+
uint64_t *workspaceSize,
32+
aclOpExecutor **executor);
33+
34+
/* funtion: aclnnDispatchLayout
35+
* workspace : workspace memory addr(input).
36+
* workspaceSize : size of workspace(input).
37+
* executor : executor context(input).
38+
* stream : acl stream.
39+
*/
40+
__attribute__((visibility("default"))) aclnnStatus aclnnDispatchLayout(
41+
void *workspace,
42+
uint64_t workspaceSize,
43+
aclOpExecutor *executor,
44+
aclrtStream stream);
45+
46+
#ifdef __cplusplus
47+
}
48+
#endif
49+
50+
#endif

‎csrc/dispatch_layout/op_host/dispatch_layout.cpp‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
#include "register/op_def_registry.h"
2+
3+
namespace ops {
4+
class DispatchLayout : public OpDef {
5+
public:
6+
explicit DispatchLayout(const char *name) : OpDef(name)
7+
{
8+
this->Input("topkIdx")
9+
.ParamType(REQUIRED)
10+
.DataType({ge::DT_INT64})
11+
.Format({ge::FORMAT_ND})
12+
.UnknownShapeFormat({ge::FORMAT_ND});
13+
14+
this->Attr("num_tokens").Int();
15+
this->Attr("num_ranks").Int();
16+
this->Attr("num_experts").Int();
17+
this->Attr("num_topk").Int();
18+
19+
this->Output("numTokensPerRank")
20+
.ParamType(REQUIRED)
21+
.DataType({ge::DT_INT32})
22+
.Format({ge::FORMAT_ND})
23+
.UnknownShapeFormat({ge::FORMAT_ND});
24+
this->Output("numTokensPerExpert")
25+
.ParamType(REQUIRED)
26+
.DataType({ge::DT_INT32})
27+
.Format({ge::FORMAT_ND})
28+
.UnknownShapeFormat({ge::FORMAT_ND});
29+
this->Output("isTokenInRank")
30+
.ParamType(REQUIRED)
31+
.DataType({ge::DT_INT32})
32+
.Format({ge::FORMAT_ND})
33+
.UnknownShapeFormat({ge::FORMAT_ND});
34+
35+
OpAICoreConfig aicore_config;
36+
aicore_config.DynamicCompileStaticFlag(true)
37+
.DynamicFormatFlag(true)
38+
.DynamicRankSupportFlag(true)
39+
.DynamicShapeSupportFlag(true)
40+
.NeedCheckSupportFlag(false)
41+
.PrecisionReduceFlag(true)
42+
.ExtendCfgInfo("aclnnSupport.value", "support_aclnn")
43+
.ExtendCfgInfo("jitCompile.flag", "static_true")
44+
.ExtendCfgInfo("multiKernelSupportDynamicGraph.value", "multi_kernel");
45+
46+
this->AICore().AddConfig("ascend910_93", aicore_config);
47+
}
48+
};
49+
50+
OP_ADD(DispatchLayout);
51+
} // namespace ops

0 commit comments

Comments
 (0)