diff --git a/csrc/npu/CMakeLists.txt b/csrc/npu/CMakeLists.txt index 98f4bfb2..f9990138 100644 --- a/csrc/npu/CMakeLists.txt +++ b/csrc/npu/CMakeLists.txt @@ -28,13 +28,13 @@ if (BUILD_OPEN_PROJECT) register c_sec ) - set(CMAKE_MODULE_PATH - ${CMAKE_MODULE_PATH} + set(CMAKE_MODULE_PATH + ${CMAKE_MODULE_PATH} ${CMAKE_CURRENT_LIST_DIR}/cmake/modules ) - set(CMAKE_PREFIX_PATH - ${CMAKE_PREFIX_PATH} + set(CMAKE_PREFIX_PATH + ${CMAKE_PREFIX_PATH} ${ASCEND_CANN_PACKAGE_PATH} ) @@ -155,6 +155,11 @@ if (BUILD_OPEN_PROJECT) LOG_CPP PROCESS_LOG ) + # Mc2CcTilingConfig::SetCommEngine only exists in CANN releases that support + # Ascend 950; gate its use so 910C builds against older CANN still compile. + if ("${ASCEND_COMPUTE_UNIT}" MATCHES "ascend950") + target_compile_definitions(optiling PRIVATE AFD_TILING_HAS_COMM_ENGINE=1) + endif () target_link_libraries(optiling PRIVATE $ -Wl,--whole-archive diff --git a/csrc/npu/README.md b/csrc/npu/README.md index 3d472c46..2bebc85f 100644 --- a/csrc/npu/README.md +++ b/csrc/npu/README.md @@ -40,7 +40,8 @@ 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`, `ascend950*`, and `Ascend950*` build + the same operators for Atlas A5 (`ascend950`). - `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..66bb0d7a 100644 --- a/csrc/npu/a2e/op_host/a2e_def.cpp +++ b/csrc/npu/a2e/op_host/a2e_def.cpp @@ -75,6 +75,7 @@ class A2e : public OpDef { this->MC2().HcclGroup({"group_ep"}); this->AICore().AddConfig("ascend910_93"); + 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..df6c9234 100644 --- a/csrc/npu/a2e/op_host/a2e_tiling.cpp +++ b/csrc/npu/a2e/op_host/a2e_tiling.cpp @@ -87,6 +87,17 @@ namespace optiling { std::string algConfigAllToAllStr = "AlltoAll=level0:fullmesh;level1:pairwise"; AscendC::Mc2CcTilingConfig mc2CcTilingConfig(groupEp, opType1, algConfigAllToAllStr); +#ifdef AFD_TILING_HAS_COMM_ENGINE + // On A5 (Ascend 950) the MC2 tiling's commEngine field is a + // HcclAccelerator, not a CommEngine. The AIV value (=3, see CANN + // hccl_params/MAKE_ENUM HcclAccelerator: DEFAULT,HOSTCPU_TS,AICPU_TS, + // AIV) is required for the MTE path. SetCommEngine(2)=AICPU_TS is + // rejected by HCCL GetTilingAccelerator with HCCL_E_NOT_SUPPORT. + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + if (ascendcPlatform.GetSocVersion() == platform_ascendc::SocVersion::ASCEND950) { + mc2CcTilingConfig.SetCommEngine(3); + } +#endif 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..d00f9400 100644 --- a/csrc/npu/a2e/op_kernel/a2e.h +++ b/csrc/npu/a2e/op_kernel/a2e.h @@ -17,7 +17,7 @@ using namespace AscendC; -#define DATA_FULSH(_gm_tensor, _type) \ +#define DATA_FLUSH(_gm_tensor, _type) \ Barrier(); \ DataCacheCleanAndInvalid<_type, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(_gm_tensor); \ __asm__("NOP"); \ @@ -34,7 +34,12 @@ class A2e { constexpr static uint32_t DOUBLE_BUFFER_COUNT = 2; constexpr static uint32_t BLOCK_IDX_USED_2 = 2; constexpr static uint32_t UB_OFFSET = 32; +#if defined(AFD_ARCH_A5) + // A5 per-rank window step is EP_RANK_OFFSET_STEP (1KB), not the A3 512B step. + constexpr static uint32_t OPT_RANK_OFFSET = 1024; +#else constexpr static uint32_t OPT_RANK_OFFSET = 512; +#endif constexpr static uint32_t INT64_COUNT_PER_BLOCK = 4; constexpr static uint32_t INT32_COUNT_PER_BLOCK = 8; public: @@ -73,10 +78,14 @@ class A2e { pipe.InitBuffer(tBuf, UB_SINGLE_TOTAL_SIZE_MAX); +#ifdef AFD_ARCH_A5 + epWinContextA5_ = (__gm__ Moe::HcclCombinOpParam *)AscendC::GetHcclContext(); +#else epWinContext_ = (__gm__ HcclOpResParam *)AscendC::GetHcclContext(); +#endif - magicTensor_.SetGlobalBuffer((__gm__ int32_t*)((epWinContext_->localWindowsIn) + - IPC_DATA_OFFSET - blockNum * sizeof(int32_t) * INT32_COUNT_PER_BLOCK)); + magicTensor_.SetGlobalBuffer((__gm__ int32_t*)(winBaseOf(rank) + + 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 +100,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] = winBaseOf(rank) + rank * OPT_RANK_OFFSET; + shareAddrs[rank % expertRankSize] = winBaseOf(rank % expertRankSize) + + (rank % expertRankSize) * 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] = winBaseOf(i) + i * OPT_RANK_OFFSET; } } @@ -150,11 +154,26 @@ class A2e { } private: + // Window base address of `rankId` as mapped on this device. On A5 HCCL hands + // out a flat cross-card window array; on A3 the local window and the + // remoteRes tree are separate. + __aicore__ inline GM_ADDR winBaseOf(int rankId) + { +#ifdef AFD_ARCH_A5 + return (GM_ADDR)(epWinContextA5_->windowsIn[rankId]); +#else + if (rankId == this->rank) { + return (GM_ADDR)(epWinContext_->localWindowsIn); + } + return (GM_ADDR)(((HcclRankRelationResV2 *)(epWinContext_->remoteRes[rankId].nextDevicePtr))->windowsIn); +#endif + } + __aicore__ inline void waitFlagWithScalar(int addr, uint32_t magic) { GlobalTensor flagGt; flagGt.SetGlobalBuffer((__gm__ uint32_t *)(shareAddrs[rank] + addr)); while(1) { - DATA_FULSH(flagGt, uint32_t); + DATA_FLUSH(flagGt, uint32_t); if (flagGt.GetValue(0) == magic) { return; } @@ -170,7 +189,7 @@ class A2e { flagLt.SetValue(0, mergeMagicWithValue(magic, 0)); GlobalTensor shareFlagGt; - shareFlagGt.SetGlobalBuffer((__gm__ uint64_t *)(shareAddrs[rank % expertRankSize]) + + shareFlagGt.SetGlobalBuffer((__gm__ uint64_t *)(shareAddrs[rank % expertRankSize]) + (1 * attnToMoeRatio + sendOffset) * flagUnitInt64Num); AscendC::SetFlag(EVENT_ID0); @@ -193,7 +212,7 @@ class A2e { flagLt.SetValue(0, mergeMagicWithValue(magic, 0)); GlobalTensor shareFlagGt; - shareFlagGt.SetGlobalBuffer((__gm__ uint64_t *)(shareAddrs[rank % expertRankSize]) + + shareFlagGt.SetGlobalBuffer((__gm__ uint64_t *)(shareAddrs[rank % expertRankSize]) + (WAIT_FLAG_OFFSET_2 * attnToMoeRatio + sendOffset) * flagUnitInt64Num); AscendC::SetFlag(EVENT_ID0); @@ -220,7 +239,7 @@ class A2e { flagLt.SetValue(0, mergeMagicWithValue(magic, 0)); GlobalTensor shareFlagGt; - shareFlagGt.SetGlobalBuffer((__gm__ uint64_t *)(shareAddrs[rank % expertRankSize]) + + shareFlagGt.SetGlobalBuffer((__gm__ uint64_t *)(shareAddrs[rank % expertRankSize]) + (WAIT_FLAG_OFFSET_4 * attnToMoeRatio + sendOffset) * flagUnitInt64Num); AscendC::SetFlag(EVENT_ID0); @@ -236,9 +255,9 @@ class A2e { __aicore__ inline void sendX(int sendOffset, int xOffset) { GlobalTensor shareXGt; shareXGt.SetGlobalBuffer((__gm__ T *)(shareAddrs[rank] + IPC_DATA_OFFSET + xOffset)); - + int actualBlockNum = (computeGate == 0) ? blockNum : (blockNum - BLOCK_IDX_USED_2); - + copyGmToGmWithBlocks(shareXGt, xGt, batchSize * hiddenSize, actualBlockNum, blockIdx); AscendC::SetFlag(EVENT_ID0); @@ -250,7 +269,7 @@ class A2e { flagLt.SetValue(0, mergeMagicWithValue(magic, 0)); GlobalTensor shareFlagGt; - shareFlagGt.SetGlobalBuffer((__gm__ uint64_t *)(shareAddrs[rank % expertRankSize]) + + shareFlagGt.SetGlobalBuffer((__gm__ uint64_t *)(shareAddrs[rank % expertRankSize]) + (WAIT_FLAG_OFFSET_3 * attnToMoeRatio + sendOffset) * flagUnitInt64Num); AscendC::SetFlag(EVENT_ID0); @@ -292,7 +311,7 @@ class A2e { int expertScalesOffset = sizeof(int32_t) + expertIdsReserveSize; int sendRank = rank + (index + 1) * expertRankSize; - + if (computeGate == 0) { waitFlagWithScalar((WAIT_FLAG_OFFSET_3 * attnToMoeRatio + index) * flagUnitInt64Num * sizeof(uint64_t) + sizeof(uint32_t), magic); @@ -308,7 +327,7 @@ class A2e { GlobalTensor shareExpertScalesGt; shareExpertScalesGt.SetGlobalBuffer((__gm__ float *)(shareAddrs[sendRank] + IPC_DATA_OFFSET + expertScalesOffset)); - + copyGmToGmWithBlocks(simulateExpertScalesGt[expertScalesOutputOffset], shareExpertScalesGt, recvBatchSize * topk, blockNum - 1, blockIdx); expertScalesOutputOffset += recvBatchSize * topk; @@ -358,6 +377,9 @@ class A2e { __gm__ TQ *expandX; __gm__ float *dynamicScales; __gm__ HcclOpResParam *epWinContext_{nullptr}; +#ifdef AFD_ARCH_A5 + __gm__ Moe::HcclCombinOpParam *epWinContextA5_{nullptr}; +#endif TPipe pipe; TBuf tBuf; GM_ADDR shareAddrs[CAM_MAX_RANK_SIZE]; diff --git a/csrc/npu/a2e/op_kernel/comm_args.h b/csrc/npu/a2e/op_kernel/comm_args.h index 3aadb840..af9ce2c9 100644 --- a/csrc/npu/a2e/op_kernel/comm_args.h +++ b/csrc/npu/a2e/op_kernel/comm_args.h @@ -5,11 +5,49 @@ #define FORCE_INLINE_AICORE __attribute__((always_inline)) inline __aicore__ #include "kernel_operator.h" +// A5 (Ascend 950, __NPU_ARCH__ == 3510 / __DAV_C310__) exposes a flat per-rank +// window array (HcclCombinOpParam) instead of the A3 HcclOpResParam + remoteRes tree. +#if (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)) || defined(__DAV_C310__) +#define AFD_ARCH_A5 1 +#endif + namespace Moe { constexpr int CAM_MAX_RANK_SIZE = 384; // Maximum number of NPU cards supported by the communication library +#ifdef AFD_ARCH_A5 +constexpr uint32_t HCCL_MTE_MAX_RANK_NUM = 64; +// A5 (Ascend 950) MTE per-rank window layout, taken from CANN ops_transformer +// moe_distribute_base.h: the state region occupies A5_MTE_STATE_WIN_SIZE at the +// head of each rank's base, the data region follows, and each rank's per-rank +// step is EP_RANK_OFFSET_STEP. This replaces the A3 512B-step/2MB-data offsets. +constexpr uint64_t A5_MTE_STATE_WIN_SIZE = 1024UL * 1024UL; // state region per rank (bytes) +constexpr uint32_t EP_RANK_OFFSET_STEP = 1024; // per-rank window step (bytes) + +struct HcclCombinOpParam { + uint64_t workSpace; // client和server之间通信的地址 + uint64_t workSpaceSize; // client和server之间通信的空间大小 + uint32_t rankId; // 当前卡rankId + uint32_t rankDim; // 总卡数 + uint64_t winSize; // ccu不使用 + uint64_t windowsIn[HCCL_MTE_MAX_RANK_NUM]; // ccu不使用, MTE 数据区 + uint64_t windowsOut[HCCL_MTE_MAX_RANK_NUM]; // ccu不使用, MTE 状态区 + // for ccu + uint64_t xnAddr; + uint64_t ckeAddr; + uint64_t msAddr; + uint64_t msSize; +}; +#endif // AFD_ARCH_A5 + constexpr int64_t IPC_BUFF_MAX_SIZE = 100 * 1024 * 1024; +#ifdef AFD_ARCH_A5 +// On A5 each rank's window keeps a 1MB state region at its head and the data +// region right after it, so the "data offset" from the state base is the state +// region size (not the A3 2MB flag+data split). +constexpr int64_t IPC_DATA_OFFSET = A5_MTE_STATE_WIN_SIZE; +#else constexpr int64_t IPC_DATA_OFFSET = 2 * 1024 * 1024; // First 2MB as flag, then 100MB as data storage +#endif constexpr int64_t PING_PONG_SIZE = 2; constexpr int64_t UB_SINGLE_DMA_SIZE_MAX = 190 * 1024; constexpr int64_t SMALL_DATA_SIZE = 1 * 1024 * 1024; diff --git a/csrc/npu/build_aclnn.sh b/csrc/npu/build_aclnn.sh index 698b6584..971bcd35 100644 --- a/csrc/npu/build_aclnn.sh +++ b/csrc/npu/build_aclnn.sh @@ -8,11 +8,14 @@ ROOT_DIR=$1 SOC_VERSION=$2 case "$SOC_VERSION" in - 910c|ascend910_93*|ascend910_9392) + 910c|ascend910_9392|ascend910_93*) SOC_ARG="ascend910_93" ;; + 950|ascend950*|Ascend950*) + SOC_ARG="ascend950" + ;; *) - 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 only for Ascend 910C/950; got ${SOC_VERSION}." exit 0 ;; esac @@ -25,4 +28,10 @@ bash build.sh -n "a2e;e2a" -c "${SOC_ARG}" INSTALL_PATH="${ROOT_DIR}/afd_plugin/_cann_ops_custom" rm -rf "${INSTALL_PATH}" mkdir -p "${INSTALL_PATH}" -./output/CANN-custom_ops*.run --install-path="${INSTALL_PATH}" +shopt -s nullglob +run_files=(./output/CANN-custom_ops*.run) +if [ "${#run_files[@]}" -ne 1 ]; then + echo "expected one CANN-custom_ops installer under output/, found ${#run_files[@]}" + exit 1 +fi +"${run_files[0]}" --install-path="${INSTALL_PATH}" diff --git a/csrc/npu/e2a/op_host/e2a_def.cpp b/csrc/npu/e2a/op_host/e2a_def.cpp index a348140c..ff98b699 100644 --- a/csrc/npu/e2a/op_host/e2a_def.cpp +++ b/csrc/npu/e2a/op_host/e2a_def.cpp @@ -44,6 +44,7 @@ class E2a : public OpDef { this->Attr("aiv_num").Int(); this->AICore().AddConfig("ascend910_93"); + 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..e5a9a835 100644 --- a/csrc/npu/e2a/op_host/e2a_tiling.cpp +++ b/csrc/npu/e2a/op_host/e2a_tiling.cpp @@ -78,6 +78,17 @@ namespace optiling { std::string algConfigAllToAllStr = "AlltoAll=level0:fullmesh;level1:pairwise"; AscendC::Mc2CcTilingConfig mc2CcTilingConfig(groupEp, opType1, algConfigAllToAllStr); +#ifdef AFD_TILING_HAS_COMM_ENGINE + // On A5 (Ascend 950) the MC2 tiling's commEngine field is a + // HcclAccelerator, not a CommEngine. The AIV value (=3, see CANN + // hccl_params/MAKE_ENUM HcclAccelerator: DEFAULT,HOSTCPU_TS,AICPU_TS, + // AIV) is required for the MTE path. SetCommEngine(2)=AICPU_TS is + // rejected by HCCL GetTilingAccelerator with HCCL_E_NOT_SUPPORT. + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + if (ascendcPlatform.GetSocVersion() == platform_ascendc::SocVersion::ASCEND950) { + mc2CcTilingConfig.SetCommEngine(3); + } +#endif mc2CcTilingConfig.GetTiling(tiling->mc2InitTiling); mc2CcTilingConfig.GetTiling(tiling->mc2CcTiling1); diff --git a/csrc/npu/e2a/op_kernel/comm_args.h b/csrc/npu/e2a/op_kernel/comm_args.h index 3aadb840..af9ce2c9 100644 --- a/csrc/npu/e2a/op_kernel/comm_args.h +++ b/csrc/npu/e2a/op_kernel/comm_args.h @@ -5,11 +5,49 @@ #define FORCE_INLINE_AICORE __attribute__((always_inline)) inline __aicore__ #include "kernel_operator.h" +// A5 (Ascend 950, __NPU_ARCH__ == 3510 / __DAV_C310__) exposes a flat per-rank +// window array (HcclCombinOpParam) instead of the A3 HcclOpResParam + remoteRes tree. +#if (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)) || defined(__DAV_C310__) +#define AFD_ARCH_A5 1 +#endif + namespace Moe { constexpr int CAM_MAX_RANK_SIZE = 384; // Maximum number of NPU cards supported by the communication library +#ifdef AFD_ARCH_A5 +constexpr uint32_t HCCL_MTE_MAX_RANK_NUM = 64; +// A5 (Ascend 950) MTE per-rank window layout, taken from CANN ops_transformer +// moe_distribute_base.h: the state region occupies A5_MTE_STATE_WIN_SIZE at the +// head of each rank's base, the data region follows, and each rank's per-rank +// step is EP_RANK_OFFSET_STEP. This replaces the A3 512B-step/2MB-data offsets. +constexpr uint64_t A5_MTE_STATE_WIN_SIZE = 1024UL * 1024UL; // state region per rank (bytes) +constexpr uint32_t EP_RANK_OFFSET_STEP = 1024; // per-rank window step (bytes) + +struct HcclCombinOpParam { + uint64_t workSpace; // client和server之间通信的地址 + uint64_t workSpaceSize; // client和server之间通信的空间大小 + uint32_t rankId; // 当前卡rankId + uint32_t rankDim; // 总卡数 + uint64_t winSize; // ccu不使用 + uint64_t windowsIn[HCCL_MTE_MAX_RANK_NUM]; // ccu不使用, MTE 数据区 + uint64_t windowsOut[HCCL_MTE_MAX_RANK_NUM]; // ccu不使用, MTE 状态区 + // for ccu + uint64_t xnAddr; + uint64_t ckeAddr; + uint64_t msAddr; + uint64_t msSize; +}; +#endif // AFD_ARCH_A5 + constexpr int64_t IPC_BUFF_MAX_SIZE = 100 * 1024 * 1024; +#ifdef AFD_ARCH_A5 +// On A5 each rank's window keeps a 1MB state region at its head and the data +// region right after it, so the "data offset" from the state base is the state +// region size (not the A3 2MB flag+data split). +constexpr int64_t IPC_DATA_OFFSET = A5_MTE_STATE_WIN_SIZE; +#else constexpr int64_t IPC_DATA_OFFSET = 2 * 1024 * 1024; // First 2MB as flag, then 100MB as data storage +#endif constexpr int64_t PING_PONG_SIZE = 2; constexpr int64_t UB_SINGLE_DMA_SIZE_MAX = 190 * 1024; constexpr int64_t SMALL_DATA_SIZE = 1 * 1024 * 1024; diff --git a/csrc/npu/e2a/op_kernel/e2a.h b/csrc/npu/e2a/op_kernel/e2a.h index 2bc0c890..c7af3d8d 100644 --- a/csrc/npu/e2a/op_kernel/e2a.h +++ b/csrc/npu/e2a/op_kernel/e2a.h @@ -17,7 +17,7 @@ using namespace AscendC; -#define DATA_FULSH_E2A(_gm_tensor, _type) \ +#define DATA_FLUSH_E2A(_gm_tensor, _type) \ Barrier(); \ DataCacheCleanAndInvalid<_type, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(_gm_tensor); \ __asm__("NOP"); \ @@ -30,7 +30,12 @@ class E2a { constexpr static int64_t MAGIC_OFFSET = 32; constexpr static uint32_t INT32_COUNT_PER_BLOCK = 8; constexpr static uint32_t DOUBLE_BUFFER_COUNT = 2; +#if defined(AFD_ARCH_A5) + // A5 per-rank window step is EP_RANK_OFFSET_STEP (1KB), not the A3 512B step. + constexpr static uint32_t OPT_RANK_OFFSET = 1024; +#else constexpr static uint32_t OPT_RANK_OFFSET = 512; +#endif public: __aicore__ inline E2a(int rank, int rankSize) { @@ -52,11 +57,15 @@ class E2a { this->blockNum = GetBlockNum(); pipe.InitBuffer(tBuf, UB_SINGLE_TOTAL_SIZE_MAX); - + +#ifdef AFD_ARCH_A5 + epWinContextA5_ = (__gm__ Moe::HcclCombinOpParam *)AscendC::GetHcclContext(); +#else epWinContext_ = (__gm__ HcclOpResParam *)AscendC::GetHcclContext(); +#endif - magicTensor_.SetGlobalBuffer((__gm__ int32_t*)((epWinContext_->localWindowsIn) + - IPC_DATA_OFFSET - blockNum * sizeof(int32_t) * INT32_COUNT_PER_BLOCK)); + magicTensor_.SetGlobalBuffer((__gm__ int32_t*)(winBaseOf(rank) + + 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 +80,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] = winBaseOf(rank) + rank * OPT_RANK_OFFSET; + shareAddrs[rank % expertRankSize] = winBaseOf(rank % expertRankSize) + + (rank % expertRankSize) * 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] = winBaseOf(i) + i * OPT_RANK_OFFSET; } } @@ -110,6 +114,20 @@ class E2a { } private: + // Window base address of `rankId` as mapped on this device. On A5 HCCL hands + // out a flat cross-card window array; on A3 the local window and the + // remoteRes tree are separate. + __aicore__ inline GM_ADDR winBaseOf(int rankId) + { +#ifdef AFD_ARCH_A5 + return (GM_ADDR)(epWinContextA5_->windowsIn[rankId]); +#else + if (rankId == this->rank) { + return (GM_ADDR)(epWinContext_->localWindowsIn); + } + return (GM_ADDR)(((HcclRankRelationResV2 *)(epWinContext_->remoteRes[rankId].nextDevicePtr))->windowsIn); +#endif + } __aicore__ inline uint64_t mergeMagicWithValue(uint32_t magic, uint16_t offset, uint16_t count) { @@ -175,6 +193,9 @@ class E2a { __gm__ T *x; __gm__ T *expandX; __gm__ HcclOpResParam *epWinContext_{nullptr}; +#ifdef AFD_ARCH_A5 + __gm__ Moe::HcclCombinOpParam *epWinContextA5_{nullptr}; +#endif TPipe pipe; TBuf tBuf; GM_ADDR shareAddrs[CAM_MAX_RANK_SIZE]; diff --git a/pyproject.toml b/pyproject.toml index 39fb96b6..a91ceeae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -146,6 +146,8 @@ extend-ignore-identifiers-re = [ ".*nin.*", ".*[Oo]no.*[Aa]nna.*", ".*cann.*", + "alog", + ".*alog.*", ".*[Nn][Dd]3.*", ".*updator.*", "byt5.*", diff --git a/tests/unit/compat/npu/a2e_e2a_a5.py b/tests/unit/compat/npu/a2e_e2a_a5.py new file mode 100644 index 00000000..4482f32e --- /dev/null +++ b/tests/unit/compat/npu/a2e_e2a_a5.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""A5 / Ascend 950 1A1F hardware sample for A2E and E2A. + +This is a standalone mp.spawn script, not a default pytest case. The filename +avoids the ``test_*.py`` pattern so collection does not import torch_npu. + +Run: + + AFD_TEST_DEVICES=2,3 python tests/unit/compat/npu/a2e_e2a_a5.py + +Environment: + AFD_TEST_DEVICES physical NPU ids, default ``2,3`` + AFD_TEST_E expert/FFN rank count, default 1 + AFD_TEST_A attention rank count, default same as E + AFD_TEST_WORLD HCCL world size, default E + A + AFD_TEST_PORT MASTER_PORT, default 29601 +""" + +import os + +import numpy as np +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch_npu # noqa: F401 + +# afd registers a2e/e2a in torch.ops.afd_ascend via this extension. Each +# mp.spawn child is a fresh interpreter, so import it in every process. +# ensure_cam_p2p_ops_available() also sets AFD_CUST_OPAPI_LIB_PATH and the +# vendor paths so the aclnnA2e/aclnnE2a symbols are resolvable at runtime. +from afd_plugin.compat.npu.ops import ensure_cam_p2p_ops_available + +ensure_cam_p2p_ops_available() +import afd_plugin._C_ascend # noqa: E402, F401 + + +class A2eE2aModule(torch.nn.Module): + def forward( + self, + x, + expert_ids, + scales, + batch_size, + hidden_size, + topk, + expert_rank_size, + attention_rank_size, + rank, + group_ep, + aiv_num, + ): + is_attention_side = rank >= expert_rank_size + + if is_attention_side: + # Attention side: send data to MOE side via A2E. + a2e_output = torch.ops.afd_ascend.a2e( + x=x, + expert_ids=expert_ids, + scales=scales, + batch_size=batch_size, + hidden_size=hidden_size, + topk=topk, + expert_rank_size=expert_rank_size, + attention_rank_size=attention_rank_size, + rank=rank, + group_ep=group_ep, + aiv_num=aiv_num, + compute_gate=1, + ) + ( + expand_x, + simulate_expert_ids, + simulate_expert_scales, + atten_batch_size, + x_active_mask_out, + ) = a2e_output + + # Attention side: receive data back from MOE side via E2A. + e2a_output = torch.ops.afd_ascend.e2a( + expand_x=x, + atten_batch_size=atten_batch_size, + batch_size=batch_size, + hidden_size=hidden_size, + topk=topk, + expert_rank_size=expert_rank_size, + attention_rank_size=attention_rank_size, + rank=rank, + group_ep=group_ep, + aiv_num=aiv_num, + ) + + return ( + e2a_output, + expand_x, + simulate_expert_ids, + simulate_expert_scales, + atten_batch_size, + x_active_mask_out, + ) + + dummy_x = torch.empty(0, hidden_size, dtype=x.dtype, device=x.device) + dummy_expert_ids = torch.empty(0, topk, dtype=torch.int32, device=x.device) + dummy_scales = torch.empty(0, topk, dtype=torch.float, device=x.device) + + a2e_output = torch.ops.afd_ascend.a2e( + x=dummy_x, + expert_ids=dummy_expert_ids, + scales=dummy_scales, + batch_size=batch_size, + hidden_size=hidden_size, + topk=topk, + expert_rank_size=expert_rank_size, + attention_rank_size=attention_rank_size, + rank=rank, + group_ep=group_ep, + aiv_num=aiv_num, + compute_gate=1, + ) + ( + expand_x, + simulate_expert_ids, + simulate_expert_scales, + atten_batch_size, + x_active_mask_out, + ) = a2e_output + + e2a_output = torch.ops.afd_ascend.e2a( + expand_x=expand_x, + atten_batch_size=atten_batch_size, + batch_size=batch_size, + hidden_size=hidden_size, + topk=topk, + expert_rank_size=expert_rank_size, + attention_rank_size=attention_rank_size, + rank=rank, + group_ep=group_ep, + aiv_num=aiv_num, + ) + + return ( + e2a_output, + expand_x, + simulate_expert_ids, + simulate_expert_scales, + atten_batch_size, + x_active_mask_out, + ) + + +def gen_x(rank, batch_size, hidden_size): + """Generate input tensor data.""" + return [ + rank * batch_size + i + 1 for i in range(batch_size) for _ in range(hidden_size) + ] + + +def gen_expert_ids(rank, batch_size, topk, expert_rank_size): + """Generate expert indices data.""" + arr = [0] * (batch_size * topk) + for i in range(batch_size): + for j in range(topk): + arr[i * topk + j] = (rank + i + j) % expert_rank_size + return arr + + +def gen_scales(batch_size, topk): + """Generate scaling factors data.""" + return [1.0 / topk] * (batch_size * topk) + + +def run_once(local_rank_id, ep_world_size): + """Single run test function (1A1F by default; scale via env).""" + devices = [ + int(d) for d in os.environ.get("AFD_TEST_DEVICES", "2,3").split(",") if d != "" + ] + assert len(devices) >= ep_world_size, ( + f"Need at least {ep_world_size} devices via AFD_TEST_DEVICES, got {devices}" + ) + physical_dev = devices[local_rank_id] + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = os.environ.get("AFD_TEST_PORT", "29601") + rank = local_rank_id + world_size = ep_world_size + + torch.npu.set_device(physical_dev) + dist.init_process_group(backend="hccl", rank=rank, world_size=world_size) + + # Test parameters. Default is 1A1F (expert_rank_size=1, attention_rank_size=1). + batch_size = 16 + hidden_size = 512 + topk = 2 + expert_rank_size = int(os.environ.get("AFD_TEST_E", "1")) + attention_rank_size = int(os.environ.get("AFD_TEST_A", expert_rank_size)) + aiv_num = 4 + data_type = torch.bfloat16 + + is_attention_side = rank >= expert_rank_size + + # Communication domain: a single group of both ranks (1 EP + 1 Attn). + ep_ranks_list = list(range(world_size)) + ep_group = dist.new_group(backend="hccl", ranks=ep_ranks_list) + ep_hcomm_info = ep_group._get_backend(torch.device("npu")).get_hccl_comm_name(rank) + torch.npu.synchronize() + + if is_attention_side: + x_data = np.array(gen_x(rank, batch_size, hidden_size)).reshape( + batch_size, hidden_size + ) + x_tensor = torch.tensor(x_data, dtype=data_type, device="npu") + + expert_ids_data = np.array( + gen_expert_ids(rank, batch_size, topk, expert_rank_size) + ).reshape(batch_size, topk) + expert_ids_tensor = torch.tensor( + expert_ids_data, dtype=torch.int32, device="npu" + ) + + scales_data = np.array(gen_scales(batch_size, topk)).reshape(batch_size, topk) + scales_tensor = torch.tensor(scales_data, dtype=torch.float, device="npu") + else: + x_tensor = torch.empty(0, hidden_size, dtype=data_type, device="npu") + expert_ids_tensor = torch.empty(0, topk, dtype=torch.int32, device="npu") + scales_tensor = torch.empty(0, topk, dtype=torch.float, device="npu") + + mod = A2eE2aModule().npu() + + ( + e2a_output, + expand_x, + simulate_expert_ids, + simulate_expert_scales, + atten_batch_size, + x_active_mask_out, + ) = mod( + x=x_tensor, + expert_ids=expert_ids_tensor, + scales=scales_tensor, + batch_size=batch_size, + hidden_size=hidden_size, + topk=topk, + expert_rank_size=expert_rank_size, + attention_rank_size=attention_rank_size, + rank=rank, + group_ep=ep_hcomm_info, + aiv_num=aiv_num, + ) + + torch.npu.synchronize() + + if is_attention_side: + print( + f"Attention Side Rank {rank} (dev{physical_dev}): " + "A2E-E2A sample run completed!" + ) + print(f" Input shape: {x_tensor.shape}") + print(f" E2A output shape: {e2a_output.shape}") + assert e2a_output.shape == x_tensor.shape, ( + f"E2A output shape mismatch: {e2a_output.shape} vs {x_tensor.shape}" + ) + assert torch.allclose(e2a_output, x_tensor, atol=1e-3), ( + "E2A output does not match input x" + ) + print(" Input and output are consistent!") + else: + print( + f"MOE Side Rank {rank} (dev{physical_dev}): A2E-E2A sample run completed!" + ) + print(f" A2E expand_x shape: {expand_x.shape}") + print(f" A2E simulate_expert_ids shape: {simulate_expert_ids.shape}") + print(f" A2E simulate_expert_scales shape: {simulate_expert_scales.shape}") + print(f" A2E atten_batch_size shape: {atten_batch_size.shape}") + print(f" A2E x_active_mask_out shape: {x_active_mask_out.shape}") + + dist.destroy_process_group() + + +if __name__ == "__main__": + e = int(os.environ.get("AFD_TEST_E", "1")) + a = int(os.environ.get("AFD_TEST_A", e)) + ep_world_size = int(os.environ.get("AFD_TEST_WORLD", e + a)) # default 1A1F + devices = os.environ.get("AFD_TEST_DEVICES", "2,3") + if len([d for d in devices.split(",") if d]) < ep_world_size: + print( + f"Need >= {ep_world_size} devices for world_size={ep_world_size}; " + "set AFD_TEST_DEVICES." + ) + raise SystemExit(1) + print("A2E-E2A A5 sample started!") + print(f"Running with {ep_world_size} ranks (E={e}, A={a}), devices={devices}") + mp.spawn(run_once, args=(ep_world_size,), nprocs=ep_world_size, join=True) + print("A2E-E2A A5 sample completed successfully!") diff --git a/tests/unit/compat/npu/test_ops.py b/tests/unit/compat/npu/test_ops.py index 5967746c..bf57b058 100644 --- a/tests/unit/compat/npu/test_ops.py +++ b/tests/unit/compat/npu/test_ops.py @@ -9,6 +9,10 @@ AFD_RUN_ASCEND_OP_RUNTIME=1 torchrun --standalone --nproc_per_node=2 \ tests/unit/compat/npu/test_ops.py + +The A5 1A1F mp.spawn hardware sample is: + + python tests/unit/compat/npu/a2e_e2a_a5.py """ from __future__ import annotations @@ -281,8 +285,8 @@ def _run_as_script() -> int: for name, reason in skipped: print(f"SKIP {name}: {reason}", flush=True) - for name, exc in failures: - print(f"FAIL {name}: {exc!r}", flush=True) + for name, error in failures: + print(f"FAIL {name}: {error!r}", flush=True) return 1 if failures else 0