Skip to content

[NPU] add A5/ascend950 support for A2E and E2A ops #324

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wangyibo1005 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[NPU] add A5/ascend950 support for A2E and E2A ops #324

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions csrc/npu/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}
)

Expand Down Expand Up @@ -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
$<BUILD_INTERFACE:intf_pub>
-Wl,--whole-archive
Expand Down
3 changes: 2 additions & 1 deletion csrc/npu/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
1 change: 1 addition & 0 deletions csrc/npu/a2e/op_host/a2e_def.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ class A2e : public OpDef {

this->MC2().HcclGroup({"group_ep"});
this->AICore().AddConfig("ascend910_93");
this->AICore().AddConfig("ascend950");
}
};

Expand Down
11 changes: 11 additions & 0 deletions csrc/npu/a2e/op_host/a2e_tiling.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
64 changes: 43 additions & 21 deletions csrc/npu/a2e/op_kernel/a2e.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"); \
Expand All @@ -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:
Expand Down Expand Up @@ -73,10 +78,14 @@ class A2e {

pipe.InitBuffer(tBuf, UB_SINGLE_TOTAL_SIZE_MAX);

#ifdef AFD_ARCH_A5
epWinContextA5_ = (__gm__ Moe::HcclCombinOpParam *)AscendC::GetHcclContext<HCCL_GROUP_ID_0>();
#else
epWinContext_ = (__gm__ HcclOpResParam *)AscendC::GetHcclContext<HCCL_GROUP_ID_0>();
#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<int32_t> tempLocal = tBuf.GetWithOffset<int32_t>(INT32_COUNT_PER_BLOCK, 0);
tempLocal(0) = 1;
Expand All @@ -91,20 +100,15 @@ class A2e {
PipeBarrier<PIPE_ALL>();

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;
}
}

Expand Down Expand Up @@ -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<uint32_t> 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;
}
Expand All @@ -170,7 +189,7 @@ class A2e {
flagLt.SetValue(0, mergeMagicWithValue(magic, 0));

GlobalTensor<uint64_t> shareFlagGt;
shareFlagGt.SetGlobalBuffer((__gm__ uint64_t *)(shareAddrs[rank % expertRankSize]) +
shareFlagGt.SetGlobalBuffer((__gm__ uint64_t *)(shareAddrs[rank % expertRankSize]) +
(1 * attnToMoeRatio + sendOffset) * flagUnitInt64Num);

AscendC::SetFlag<HardEvent::S_MTE3>(EVENT_ID0);
Expand All @@ -193,7 +212,7 @@ class A2e {
flagLt.SetValue(0, mergeMagicWithValue(magic, 0));

GlobalTensor<uint64_t> 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<HardEvent::S_MTE3>(EVENT_ID0);
Expand All @@ -220,7 +239,7 @@ class A2e {
flagLt.SetValue(0, mergeMagicWithValue(magic, 0));

GlobalTensor<uint64_t> 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<HardEvent::S_MTE3>(EVENT_ID0);
Expand All @@ -236,9 +255,9 @@ class A2e {
__aicore__ inline void sendX(int sendOffset, int xOffset) {
GlobalTensor<T> 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<HardEvent::MTE3_S>(EVENT_ID0);
Expand All @@ -250,7 +269,7 @@ class A2e {
flagLt.SetValue(0, mergeMagicWithValue(magic, 0));

GlobalTensor<uint64_t> 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<HardEvent::S_MTE3>(EVENT_ID0);
Expand Down Expand Up @@ -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);
Expand All @@ -308,7 +327,7 @@ class A2e {
GlobalTensor<float> shareExpertScalesGt;
shareExpertScalesGt.SetGlobalBuffer((__gm__ float *)(shareAddrs[sendRank] + IPC_DATA_OFFSET +
expertScalesOffset));

copyGmToGmWithBlocks(simulateExpertScalesGt[expertScalesOutputOffset], shareExpertScalesGt,
recvBatchSize * topk, blockNum - 1, blockIdx);
expertScalesOutputOffset += recvBatchSize * topk;
Expand Down Expand Up @@ -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<QuePosition::VECCALC> tBuf;
GM_ADDR shareAddrs[CAM_MAX_RANK_SIZE];
Expand Down
38 changes: 38 additions & 0 deletions csrc/npu/a2e/op_kernel/comm_args.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
15 changes: 12 additions & 3 deletions csrc/npu/build_aclnn.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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}"
1 change: 1 addition & 0 deletions csrc/npu/e2a/op_host/e2a_def.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"});
}
};
Expand Down
11 changes: 11 additions & 0 deletions csrc/npu/e2a/op_host/e2a_tiling.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
38 changes: 38 additions & 0 deletions csrc/npu/e2a/op_kernel/comm_args.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading
Loading