[NPU] add A5 support for A2E and E2A ops

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.

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

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<uint32_t>(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);
 

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<HCCL_GROUP_ID_0>();
+        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<int32_t> tempLocal = tBuf.GetWithOffset<int32_t>(INT32_COUNT_PER_BLOCK, 0);
         tempLocal(0) = 1;
@@ -91,20 +91,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] = 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<QuePosition::VECCALC> tBuf;
     GM_ADDR shareAddrs[CAM_MAX_RANK_SIZE];

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<AscendC::HCCL_GROUP_ID_0>();
+    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<AscendC::HCCL_GROUP_ID_0>();
+    return (GM_ADDR)(ctx->windowsIn[rankId] + A5_MTE_STATE_WIN_SIZE);
+}
+#else
+__aicore__ inline GM_ADDR GetHcclLocalWindowsIn()
+{
+    __gm__ HcclOpResParam *ctx = (__gm__ HcclOpResParam *)AscendC::GetHcclContext<AscendC::HCCL_GROUP_ID_0>();
+    return (GM_ADDR)(ctx->localWindowsIn);
+}
+
+__aicore__ inline GM_ADDR GetHcclRankWindowsIn(int32_t rankId, int32_t curRank)
+{
+    __gm__ HcclOpResParam *ctx = (__gm__ HcclOpResParam *)AscendC::GetHcclContext<AscendC::HCCL_GROUP_ID_0>();
+    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,

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

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

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<uint32_t>(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);
 

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<HCCL_GROUP_ID_0>();
 
-        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<int32_t> tempLocal = tBuf.GetWithOffset<int32_t>(INT32_COUNT_PER_BLOCK, 0);
         tempLocal(0) = 1;
@@ -71,20 +71,15 @@ class E2a {
         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] = 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<QuePosition::VECCALC> tBuf;
     GM_ADDR shareAddrs[CAM_MAX_RANK_SIZE];

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<AscendC::HCCL_GROUP_ID_0>();
+    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<AscendC::HCCL_GROUP_ID_0>();
+    return (GM_ADDR)(ctx->windowsIn[rankId] + A5_MTE_STATE_WIN_SIZE);
+}
+#else
+__aicore__ inline GM_ADDR GetHcclLocalWindowsIn()
+{
+    __gm__ HcclOpResParam *ctx = (__gm__ HcclOpResParam *)AscendC::GetHcclContext<AscendC::HCCL_GROUP_ID_0>();
+    return (GM_ADDR)(ctx->localWindowsIn);
+}
+
+__aicore__ inline GM_ADDR GetHcclRankWindowsIn(int32_t rankId, int32_t curRank)
+{
+    __gm__ HcclOpResParam *ctx = (__gm__ HcclOpResParam *)AscendC::GetHcclContext<AscendC::HCCL_GROUP_ID_0>();
+    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,

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