From 409e800326678999ac1b0bd3a670f096cdaca361 Mon Sep 17 00:00:00 2001
From: baominghelly <41820386+baominghelly@users.noreply.github.com>
Date: Tue, 25 Aug 2026 21:36:51 +0800
Subject: [PATCH] feat(ascend): integrate InfiniOps backend

---
 src/infinicore/nn/rope.cc                                   | 3 ++-
 src/infinicore/ops/infiniops_impl.hpp                       | 4 ++++
 .../ops/rotary_embedding/rotary_embedding_infiniops.cc      | 6 +++++-
 xmake.lua                                                   | 2 ++
 4 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/infinicore/nn/rope.cc b/src/infinicore/nn/rope.cc
index 8decfb25e..85064e792 100644
--- a/src/infinicore/nn/rope.cc
+++ b/src/infinicore/nn/rope.cc
@@ -82,7 +82,8 @@ void RoPE::initialize_cache() {
          || device_.getType() == Device::Type::METAX
          || device_.getType() == Device::Type::ILUVATAR
          || device_.getType() == Device::Type::CAMBRICON
-         || device_.getType() == Device::Type::HYGON)
+         || device_.getType() == Device::Type::HYGON
+         || device_.getType() == Device::Type::ASCEND)
         && !mrope_section_) {
         INFINICORE_NN_BUFFER_INIT(cos_sin_cache, ({max_seq_len_, rotary_dim_}, dtype_, device_));
     }
diff --git a/src/infinicore/ops/infiniops_impl.hpp b/src/infinicore/ops/infiniops_impl.hpp
index c450ccfbd..29b988387 100644
--- a/src/infinicore/ops/infiniops_impl.hpp
+++ b/src/infinicore/ops/infiniops_impl.hpp
@@ -59,6 +59,8 @@ inline infini::ops::Device toInfiniOpsDevice(const Device &device) {
         return infini::ops::Device{infini::ops::Device::Type::kCambricon, static_cast<int>(device.getIndex())};
     case Device::Type::HYGON:
         return infini::ops::Device{infini::ops::Device::Type::kHygon, static_cast<int>(device.getIndex())};
+    case Device::Type::ASCEND:
+        return infini::ops::Device{infini::ops::Device::Type::kAscend, static_cast<int>(device.getIndex())};
     default:
         throw std::runtime_error("InfiniOps backend does not support this device type.");
     }
@@ -72,6 +74,7 @@ inline bool isSupportedDevice(Device::Type device_type) {
     case Device::Type::ILUVATAR:
     case Device::Type::CAMBRICON:
     case Device::Type::HYGON:
+    case Device::Type::ASCEND:
         return true;
     default:
         return false;
@@ -86,6 +89,7 @@ void registerSupportedDevices(Dispatcher &dispatcher, Function function) {
     dispatcher.registerDevice(Device::Type::ILUVATAR, function);
     dispatcher.registerDevice(Device::Type::CAMBRICON, function);
     dispatcher.registerDevice(Device::Type::HYGON, function);
+    dispatcher.registerDevice(Device::Type::ASCEND, function);
 }
 
 struct TensorMeta {
diff --git a/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc b/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc
index 9d64f7f61..76c12a4dd 100644
--- a/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc
+++ b/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc
@@ -36,7 +36,8 @@ void *plan(const Tensor &positions,
                       || device_type == Device::Type::METAX
                       || device_type == Device::Type::ILUVATAR
                       || device_type == Device::Type::CAMBRICON
-                      || device_type == Device::Type::HYGON);
+                      || device_type == Device::Type::HYGON
+                      || device_type == Device::Type::ASCEND);
     return new PlannedMeta{
         TensorMeta(positions),
         TensorMeta(query),
@@ -94,6 +95,9 @@ static bool registered = []() {
     RotaryEmbedding::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan);
     RotaryEmbedding::run_dispatcher().registerDevice(Device::Type::HYGON, &run);
     RotaryEmbedding::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup);
+    RotaryEmbedding::plan_dispatcher().registerDevice(Device::Type::ASCEND, &plan);
+    RotaryEmbedding::run_dispatcher().registerDevice(Device::Type::ASCEND, &run);
+    RotaryEmbedding::cleanup_dispatcher().registerDevice(Device::Type::ASCEND, &cleanup);
     return true;
 }();
 
diff --git a/xmake.lua b/xmake.lua
index 462aebe5d..ebfa5e548 100644
--- a/xmake.lua
+++ b/xmake.lua
@@ -470,6 +470,8 @@ local function configure_infiniops_ops(infiniops_ops, xmake_os, json)
             local implementations = "all"
             if use_linked_implementation then
                 implementations = {16}
+            elseif has_config("ascend-npu") and op == "argmax" then
+                implementations = {0}
             elseif op == "argmax" or op == "index_select" then
                 implementations = {8}
             elseif op == "rms_norm" or op == "silu_and_mul" or op == "topk_softmax" then