From c4b0c4e52f86b0f45bfc24294f597453e1cd069d Mon Sep 17 00:00:00 2001 From: wangyibo1005 <2633333316@qq.com> Date: Wed, 26 Aug 2026 17:41:36 +0800 Subject: [PATCH] [NPU] add A5 support for A2E and E2A ops A5 HCCL rejects the 910C MC2 tiling/window ABI, so register CANN 9.1 ascend950 and the CANN 8.5 ascend910_95 alias, use the MTE comm engine, and read A5 window addresses without changing the 910C path. Signed-off-by: wangyibo1005 <2633333316@qq.com> --- csrc/npu/README.md | 4 +- csrc/npu/a2e/op_host/a2e_def.cpp | 2 + csrc/npu/a2e/op_host/a2e_tiling.cpp | 15 ++++++ csrc/npu/a2e/op_kernel/a2e.h | 20 +++----- csrc/npu/a2e/op_kernel/moe_distribute_base.h | 47 +++++++++++++++++++ csrc/npu/build_aclnn.sh | 17 ++++++- csrc/npu/e2a/op_host/e2a_def.cpp | 2 + csrc/npu/e2a/op_host/e2a_tiling.cpp | 15 ++++++ csrc/npu/e2a/op_kernel/e2a.h | 22 ++++----- csrc/npu/e2a/op_kernel/moe_distribute_base.h | 47 +++++++++++++++++++ tests/unit/package/test_ascend_build_files.py | 21 +++++++++ 11 files changed, 183 insertions(+), 29 deletions(-) mode change 100644 => 100755 csrc/npu/build_aclnn.sh diff --git a/csrc/npu/README.md b/csrc/npu/README.md index 3d472c46..e3598b04 100644 --- a/csrc/npu/README.md +++ b/csrc/npu/README.md @@ -40,7 +40,9 @@ Common environment variables: `/usr/local/Ascend/ascend-toolkit/latest`. - `TORCH_NPU_PATH`: optional path to the `torch_npu` package. - `SOC_VERSION`: `910c`, `ascend910_93*`, and `ascend910_9392` build - `a2e;e2a`. + `a2e;e2a` for Ascend 910C. `950`, `a5`, `ascend950*`, and `ascend910_95*` + build the same operators for Atlas A5 (DAV_3510). CANN 9.1 compiles + `ascend950`; CANN 8.5 uses the older `ascend910_95` alias. - `MAX_JOBS`: number of parallel CMake build jobs for the PyTorch extension. - `AFD_SKIP_ACLNN_BUILD=1`: skip rebuilding the ACLNN operator package and build the PyTorch extension against an existing custom-op installation. diff --git a/csrc/npu/a2e/op_host/a2e_def.cpp b/csrc/npu/a2e/op_host/a2e_def.cpp index 803afced..44465f99 100644 --- a/csrc/npu/a2e/op_host/a2e_def.cpp +++ b/csrc/npu/a2e/op_host/a2e_def.cpp @@ -75,6 +75,8 @@ class A2e : public OpDef { this->MC2().HcclGroup({"group_ep"}); this->AICore().AddConfig("ascend910_93"); + this->AICore().AddConfig("ascend910_95"); + this->AICore().AddConfig("ascend950"); } }; diff --git a/csrc/npu/a2e/op_host/a2e_tiling.cpp b/csrc/npu/a2e/op_host/a2e_tiling.cpp index 5479ce3c..4615b4c4 100644 --- a/csrc/npu/a2e/op_host/a2e_tiling.cpp +++ b/csrc/npu/a2e/op_host/a2e_tiling.cpp @@ -37,6 +37,18 @@ constexpr int ATTR_AIV_NUM = 7; constexpr int ATTR_COMPUTE_GATE = 8; constexpr uint32_t OP_TYPE_ALL_TO_ALL = 8; +constexpr uint32_t DAV_3510_ARCH = 3510; +constexpr uint32_t HCCL_COMM_ENGINE_MTE = 3; + +bool IsDav3510(gert::TilingContext *context) +{ + auto *platformInfo = context->GetPlatformInfo(); + if (platformInfo == nullptr) { + return false; + } + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + return static_cast(ascendcPlatform.GetCurNpuArch()) == DAV_3510_ARCH; +} } namespace optiling { @@ -87,6 +99,9 @@ namespace optiling { std::string algConfigAllToAllStr = "AlltoAll=level0:fullmesh;level1:pairwise"; AscendC::Mc2CcTilingConfig mc2CcTilingConfig(groupEp, opType1, algConfigAllToAllStr); + if (IsDav3510(context)) { + mc2CcTilingConfig.SetCommEngine(HCCL_COMM_ENGINE_MTE); + } mc2CcTilingConfig.GetTiling(tiling->mc2InitTiling); mc2CcTilingConfig.GetTiling(tiling->mc2CcTiling1); diff --git a/csrc/npu/a2e/op_kernel/a2e.h b/csrc/npu/a2e/op_kernel/a2e.h index a2081f4e..62dd0613 100644 --- a/csrc/npu/a2e/op_kernel/a2e.h +++ b/csrc/npu/a2e/op_kernel/a2e.h @@ -73,10 +73,10 @@ class A2e { pipe.InitBuffer(tBuf, UB_SINGLE_TOTAL_SIZE_MAX); - epWinContext_ = (__gm__ HcclOpResParam *)AscendC::GetHcclContext(); + GM_ADDR localWindowsIn = GetHcclLocalWindowsIn(); - magicTensor_.SetGlobalBuffer((__gm__ int32_t*)((epWinContext_->localWindowsIn) + - IPC_DATA_OFFSET - blockNum * sizeof(int32_t) * INT32_COUNT_PER_BLOCK)); + magicTensor_.SetGlobalBuffer((__gm__ int32_t*)(localWindowsIn + + IPC_DATA_OFFSET - blockNum * sizeof(int32_t) * INT32_COUNT_PER_BLOCK)); LocalTensor tempLocal = tBuf.GetWithOffset(INT32_COUNT_PER_BLOCK, 0); tempLocal(0) = 1; @@ -91,20 +91,15 @@ class A2e { PipeBarrier(); if (rank >= expertRankSize) { - shareAddrs[rank] = (GM_ADDR)(epWinContext_->localWindowsIn) + rank * OPT_RANK_OFFSET; - shareAddrs[rank % expertRankSize] = (GM_ADDR)(((HcclRankRelationResV2 *)(epWinContext_-> - remoteRes[rank % expertRankSize].nextDevicePtr))->windowsIn) + (rank % expertRankSize) * OPT_RANK_OFFSET; + shareAddrs[rank] = localWindowsIn + rank * OPT_RANK_OFFSET; + int32_t remoteRank = rank % expertRankSize; + shareAddrs[remoteRank] = GetHcclRankWindowsIn(remoteRank, rank) + remoteRank * OPT_RANK_OFFSET; pipe_barrier(PIPE_ALL); } else { pipe_barrier(PIPE_ALL); for (int i = 0; i < rankSize; i++) { - if (i == rank) { - shareAddrs[i] = (GM_ADDR)(epWinContext_->localWindowsIn) + rank * OPT_RANK_OFFSET; - continue; - } - shareAddrs[i] = (GM_ADDR)(((HcclRankRelationResV2 *)(epWinContext_->remoteRes[i].nextDevicePtr))-> - windowsIn) + i * OPT_RANK_OFFSET; + shareAddrs[i] = GetHcclRankWindowsIn(i, rank) + i * OPT_RANK_OFFSET; } } @@ -357,7 +352,6 @@ class A2e { __gm__ T *x; __gm__ TQ *expandX; __gm__ float *dynamicScales; - __gm__ HcclOpResParam *epWinContext_{nullptr}; TPipe pipe; TBuf tBuf; GM_ADDR shareAddrs[CAM_MAX_RANK_SIZE]; diff --git a/csrc/npu/a2e/op_kernel/moe_distribute_base.h b/csrc/npu/a2e/op_kernel/moe_distribute_base.h index b74df9a5..ba3eeef0 100644 --- a/csrc/npu/a2e/op_kernel/moe_distribute_base.h +++ b/csrc/npu/a2e/op_kernel/moe_distribute_base.h @@ -174,6 +174,53 @@ struct HcclOpResParam { bool utraceStatusFlag; }; +#if (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)) || defined(__DAV_C310__) +constexpr uint32_t HCCL_MTE_MAX_RANK_NUM = 64; +constexpr uint64_t A5_MTE_STATE_WIN_SIZE = 1024UL * 1024UL; + +struct HcclA5OpResParam { + uint64_t workSpace; + uint64_t workSpaceSize; + uint32_t rankId; + uint32_t rankDim; + uint64_t winSize; + uint64_t windowsIn[HCCL_MTE_MAX_RANK_NUM]; + uint64_t windowsOut[HCCL_MTE_MAX_RANK_NUM]; + uint64_t xnAddr; + uint64_t ckeAddr; + uint64_t msAddr; + uint64_t msSize; +}; + +__aicore__ inline GM_ADDR GetHcclLocalWindowsIn() +{ + __gm__ HcclA5OpResParam *ctx = (__gm__ HcclA5OpResParam *)AscendC::GetHcclContext(); + return (GM_ADDR)(ctx->windowsIn[ctx->rankId] + A5_MTE_STATE_WIN_SIZE); +} + +__aicore__ inline GM_ADDR GetHcclRankWindowsIn(int32_t rankId, int32_t curRank) +{ + (void)curRank; + __gm__ HcclA5OpResParam *ctx = (__gm__ HcclA5OpResParam *)AscendC::GetHcclContext(); + return (GM_ADDR)(ctx->windowsIn[rankId] + A5_MTE_STATE_WIN_SIZE); +} +#else +__aicore__ inline GM_ADDR GetHcclLocalWindowsIn() +{ + __gm__ HcclOpResParam *ctx = (__gm__ HcclOpResParam *)AscendC::GetHcclContext(); + return (GM_ADDR)(ctx->localWindowsIn); +} + +__aicore__ inline GM_ADDR GetHcclRankWindowsIn(int32_t rankId, int32_t curRank) +{ + __gm__ HcclOpResParam *ctx = (__gm__ HcclOpResParam *)AscendC::GetHcclContext(); + if (rankId == curRank) { + return (GM_ADDR)(ctx->localWindowsIn); + } + return (GM_ADDR)(((HcclRankRelationResV2 *)(ctx->remoteRes[rankId].nextDevicePtr))->windowsIn); +} +#endif + // Transport enum class HcclAiRMAMemType : uint32_t { LOCAL_INPUT = 0, diff --git a/csrc/npu/build_aclnn.sh b/csrc/npu/build_aclnn.sh old mode 100644 new mode 100755 index 698b6584..f5eacd2a --- a/csrc/npu/build_aclnn.sh +++ b/csrc/npu/build_aclnn.sh @@ -6,13 +6,28 @@ set -euo pipefail ROOT_DIR=$1 SOC_VERSION=$2 +CANN_HOME="${ASCEND_HOME_PATH:-/usr/local/Ascend/ascend-toolkit/latest}" +OPP_CONFIG_DIR="${CANN_HOME}/opp/built-in/op_impl/ai_core/tbe/config" + +resolve_a5_compute_unit() { + # CANN 9.1+ names Atlas A5 as ascend950. CANN 8.5 used the + # transitional ascend910_95 alias for the same DAV_3510 arch. + if [ -d "${OPP_CONFIG_DIR}/ascend950" ]; then + echo "ascend950" + else + echo "ascend910_95" + fi +} case "$SOC_VERSION" in 910c|ascend910_93*|ascend910_9392) SOC_ARG="ascend910_93" ;; + 950|a5|ascend950*|ascend910_95*) + SOC_ARG="$(resolve_a5_compute_unit)" + ;; *) - echo "AFD A2E/E2A custom ACLNN ops are currently built only for Ascend 910C; got ${SOC_VERSION}." + echo "AFD A2E/E2A custom ACLNN ops are currently built for Ascend 910C or Ascend 950/A5; got ${SOC_VERSION}." exit 0 ;; esac diff --git a/csrc/npu/e2a/op_host/e2a_def.cpp b/csrc/npu/e2a/op_host/e2a_def.cpp index a348140c..0170ca04 100644 --- a/csrc/npu/e2a/op_host/e2a_def.cpp +++ b/csrc/npu/e2a/op_host/e2a_def.cpp @@ -44,6 +44,8 @@ class E2a : public OpDef { this->Attr("aiv_num").Int(); this->AICore().AddConfig("ascend910_93"); + this->AICore().AddConfig("ascend910_95"); + this->AICore().AddConfig("ascend950"); this->MC2().HcclGroup({"group_ep"}); } }; diff --git a/csrc/npu/e2a/op_host/e2a_tiling.cpp b/csrc/npu/e2a/op_host/e2a_tiling.cpp index 2e597a3f..2f3be625 100644 --- a/csrc/npu/e2a/op_host/e2a_tiling.cpp +++ b/csrc/npu/e2a/op_host/e2a_tiling.cpp @@ -30,6 +30,18 @@ constexpr int ATTR_ENUM_GROUP_EP = 6; constexpr int ATTR_AIV_NUM = 7; constexpr uint32_t OP_TYPE_ALL_TO_ALL = 8; +constexpr uint32_t DAV_3510_ARCH = 3510; +constexpr uint32_t HCCL_COMM_ENGINE_MTE = 3; + +bool IsDav3510(gert::TilingContext *context) +{ + auto *platformInfo = context->GetPlatformInfo(); + if (platformInfo == nullptr) { + return false; + } + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + return static_cast(ascendcPlatform.GetCurNpuArch()) == DAV_3510_ARCH; +} } namespace optiling { @@ -78,6 +90,9 @@ namespace optiling { std::string algConfigAllToAllStr = "AlltoAll=level0:fullmesh;level1:pairwise"; AscendC::Mc2CcTilingConfig mc2CcTilingConfig(groupEp, opType1, algConfigAllToAllStr); + if (IsDav3510(context)) { + mc2CcTilingConfig.SetCommEngine(HCCL_COMM_ENGINE_MTE); + } mc2CcTilingConfig.GetTiling(tiling->mc2InitTiling); mc2CcTilingConfig.GetTiling(tiling->mc2CcTiling1); diff --git a/csrc/npu/e2a/op_kernel/e2a.h b/csrc/npu/e2a/op_kernel/e2a.h index 2bc0c890..44723128 100644 --- a/csrc/npu/e2a/op_kernel/e2a.h +++ b/csrc/npu/e2a/op_kernel/e2a.h @@ -52,11 +52,11 @@ class E2a { this->blockNum = GetBlockNum(); pipe.InitBuffer(tBuf, UB_SINGLE_TOTAL_SIZE_MAX); - - epWinContext_ = (__gm__ HcclOpResParam *)AscendC::GetHcclContext(); - magicTensor_.SetGlobalBuffer((__gm__ int32_t*)((epWinContext_->localWindowsIn) + - IPC_DATA_OFFSET - blockNum * sizeof(int32_t) * INT32_COUNT_PER_BLOCK)); + GM_ADDR localWindowsIn = GetHcclLocalWindowsIn(); + + magicTensor_.SetGlobalBuffer((__gm__ int32_t*)(localWindowsIn + + IPC_DATA_OFFSET - blockNum * sizeof(int32_t) * INT32_COUNT_PER_BLOCK)); LocalTensor tempLocal = tBuf.GetWithOffset(INT32_COUNT_PER_BLOCK, 0); tempLocal(0) = 1; @@ -71,20 +71,15 @@ class E2a { PipeBarrier(); if (rank >= expertRankSize) { - shareAddrs[rank] = (GM_ADDR)(epWinContext_->localWindowsIn) + rank * OPT_RANK_OFFSET; - shareAddrs[rank % expertRankSize] = (GM_ADDR)(((HcclRankRelationResV2 *)(epWinContext_-> - remoteRes[rank % expertRankSize].nextDevicePtr))->windowsIn) + (rank % expertRankSize) * OPT_RANK_OFFSET; + shareAddrs[rank] = localWindowsIn + rank * OPT_RANK_OFFSET; + int32_t remoteRank = rank % expertRankSize; + shareAddrs[remoteRank] = GetHcclRankWindowsIn(remoteRank, rank) + remoteRank * OPT_RANK_OFFSET; pipe_barrier(PIPE_ALL); } else { pipe_barrier(PIPE_ALL); for (int i = 0; i < rankSize; i++) { - if (i == rank) { - shareAddrs[i] = (GM_ADDR)(epWinContext_->localWindowsIn) + rank * OPT_RANK_OFFSET; - continue; - } - shareAddrs[i] = (GM_ADDR)(((HcclRankRelationResV2 *)(epWinContext_->remoteRes[i].nextDevicePtr))-> - windowsIn) + i * OPT_RANK_OFFSET; + shareAddrs[i] = GetHcclRankWindowsIn(i, rank) + i * OPT_RANK_OFFSET; } } @@ -174,7 +169,6 @@ class E2a { __gm__ T *x; __gm__ T *expandX; - __gm__ HcclOpResParam *epWinContext_{nullptr}; TPipe pipe; TBuf tBuf; GM_ADDR shareAddrs[CAM_MAX_RANK_SIZE]; diff --git a/csrc/npu/e2a/op_kernel/moe_distribute_base.h b/csrc/npu/e2a/op_kernel/moe_distribute_base.h index b74df9a5..ba3eeef0 100644 --- a/csrc/npu/e2a/op_kernel/moe_distribute_base.h +++ b/csrc/npu/e2a/op_kernel/moe_distribute_base.h @@ -174,6 +174,53 @@ struct HcclOpResParam { bool utraceStatusFlag; }; +#if (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)) || defined(__DAV_C310__) +constexpr uint32_t HCCL_MTE_MAX_RANK_NUM = 64; +constexpr uint64_t A5_MTE_STATE_WIN_SIZE = 1024UL * 1024UL; + +struct HcclA5OpResParam { + uint64_t workSpace; + uint64_t workSpaceSize; + uint32_t rankId; + uint32_t rankDim; + uint64_t winSize; + uint64_t windowsIn[HCCL_MTE_MAX_RANK_NUM]; + uint64_t windowsOut[HCCL_MTE_MAX_RANK_NUM]; + uint64_t xnAddr; + uint64_t ckeAddr; + uint64_t msAddr; + uint64_t msSize; +}; + +__aicore__ inline GM_ADDR GetHcclLocalWindowsIn() +{ + __gm__ HcclA5OpResParam *ctx = (__gm__ HcclA5OpResParam *)AscendC::GetHcclContext(); + return (GM_ADDR)(ctx->windowsIn[ctx->rankId] + A5_MTE_STATE_WIN_SIZE); +} + +__aicore__ inline GM_ADDR GetHcclRankWindowsIn(int32_t rankId, int32_t curRank) +{ + (void)curRank; + __gm__ HcclA5OpResParam *ctx = (__gm__ HcclA5OpResParam *)AscendC::GetHcclContext(); + return (GM_ADDR)(ctx->windowsIn[rankId] + A5_MTE_STATE_WIN_SIZE); +} +#else +__aicore__ inline GM_ADDR GetHcclLocalWindowsIn() +{ + __gm__ HcclOpResParam *ctx = (__gm__ HcclOpResParam *)AscendC::GetHcclContext(); + return (GM_ADDR)(ctx->localWindowsIn); +} + +__aicore__ inline GM_ADDR GetHcclRankWindowsIn(int32_t rankId, int32_t curRank) +{ + __gm__ HcclOpResParam *ctx = (__gm__ HcclOpResParam *)AscendC::GetHcclContext(); + if (rankId == curRank) { + return (GM_ADDR)(ctx->localWindowsIn); + } + return (GM_ADDR)(((HcclRankRelationResV2 *)(ctx->remoteRes[rankId].nextDevicePtr))->windowsIn); +} +#endif + // Transport enum class HcclAiRMAMemType : uint32_t { LOCAL_INPUT = 0, diff --git a/tests/unit/package/test_ascend_build_files.py b/tests/unit/package/test_ascend_build_files.py index cf8254e2..afaf9527 100644 --- a/tests/unit/package/test_ascend_build_files.py +++ b/tests/unit/package/test_ascend_build_files.py @@ -171,3 +171,24 @@ def test_ascend_ops_use_isolated_namespace_and_vendor_path(): assert '"vllm-ascend"' not in cann_cmake assert "AFD_CUST_OPAPI_LIB_PATH" in op_api_common assert 'return "libcust_opapi.so"' not in op_api_common + + +def test_a2e_e2a_ops_are_registered_for_910c_and_a5(): + root = Path(__file__).resolve().parents[3] + a2e_def = (root / "csrc/npu/a2e/op_host/a2e_def.cpp").read_text() + e2a_def = (root / "csrc/npu/e2a/op_host/e2a_def.cpp").read_text() + build_script = (root / "csrc/npu/build_aclnn.sh").read_text() + a2e_tiling = (root / "csrc/npu/a2e/op_host/a2e_tiling.cpp").read_text() + e2a_tiling = (root / "csrc/npu/e2a/op_host/e2a_tiling.cpp").read_text() + + assert 'AddConfig("ascend910_93")' in a2e_def + assert 'AddConfig("ascend910_95")' in a2e_def + assert 'AddConfig("ascend950")' in a2e_def + assert 'AddConfig("ascend910_93")' in e2a_def + assert 'AddConfig("ascend910_95")' in e2a_def + assert 'AddConfig("ascend950")' in e2a_def + assert "resolve_a5_compute_unit" in build_script + assert "ascend950" in build_script + assert "ascend910_95" in build_script + assert "SetCommEngine(HCCL_COMM_ENGINE_MTE)" in a2e_tiling + assert "SetCommEngine(HCCL_COMM_ENGINE_MTE)" in e2a_tiling