fix(iluvatar): enable InfiniOps-backed Qwen3 inference

src/infinicore/nn/rope.cc

@@ -78,7 +78,10 @@ void RoPE::initialize_cache() {
     INFINICORE_NN_BUFFER_INIT(cos_cache, ({max_seq_len_, cache_dim}, dtype_, device_));
 
 #ifdef ENABLE_INFINIOPS_API
-    if ((device_.getType() == Device::Type::NVIDIA || device_.getType() == Device::Type::METAX) && !mrope_section_) {
+    if ((device_.getType() == Device::Type::NVIDIA
+         || device_.getType() == Device::Type::METAX
+         || device_.getType() == Device::Type::ILUVATAR)
+        && !mrope_section_) {
         INFINICORE_NN_BUFFER_INIT(cos_sin_cache, ({max_seq_len_, rotary_dim_}, dtype_, device_));
     }
 #endif
@@ -269,6 +272,21 @@ std::pair<Tensor, Tensor> RoPE::forward(const Tensor &q_out,
                                         const Tensor &k,
                                         const Tensor &positions) const {
     if (!mrope_section_.has_value()) {
+#ifdef ENABLE_INFINIOPS_API
+        if (cos_sin_cache_) {
+            auto mutable_q_out = q_out;
+            auto mutable_k_out = k_out;
+            mutable_q_out->copy_from(q);
+            mutable_k_out->copy_from(k);
+            op::rotary_embedding_(positions,
+                                  mutable_q_out,
+                                  mutable_k_out,
+                                  cos_sin_cache_,
+                                  static_cast<int64_t>(head_dim_),
+                                  algo_ == Algo::GPT_NEOX);
+            return {q_out, k_out};
+        }
+#endif
         auto apply_standard = [this, &positions](Tensor out, const Tensor &in) {
             if (rotary_dim_ < head_dim_) {
                 out->copy_from(in);

src/infinicore/ops/random_sample/random_sample.cc

@@ -2,7 +2,10 @@
 
 #include "../../utils.hpp"
 
-#ifdef ENABLE_INFINIOPS_API
+#if defined(ENABLE_INFINIOPS_API)    \
+    && (defined(ENABLE_NVIDIA_API)   \
+        || defined(ENABLE_METAX_API) \
+        || (defined(ENABLE_ILUVATAR_API) && defined(ENABLE_ATEN)))
 #include "../infiniops_impl.hpp"
 
 #include "base/argmax.h"
@@ -11,12 +14,16 @@
 namespace infinicore::op {
 namespace {
 
-#ifdef ENABLE_INFINIOPS_API
+#if defined(ENABLE_INFINIOPS_API)    \
+    && (defined(ENABLE_NVIDIA_API)   \
+        || defined(ENABLE_METAX_API) \
+        || (defined(ENABLE_ILUVATAR_API) && defined(ENABLE_ATEN)))
 bool tryGreedyWithInfiniOps(Tensor indices, Tensor logits, int topk) {
     const auto dtype = logits->dtype();
     const auto device_type = logits->device().getType();
     if ((device_type != Device::Type::NVIDIA
-         && device_type != Device::Type::METAX)
+         && device_type != Device::Type::METAX
+         && device_type != Device::Type::ILUVATAR)
         || topk != 1
         || logits->ndim() != 1
         || logits->numel() == 0
@@ -56,7 +63,10 @@ void RandomSample::execute(
     float random_val, float topp, int topk, float temperature) {
     INFINICORE_ASSERT_TENSORS_SAME_DEVICE(indices, logits);
     infinicore::context::setDevice(logits->device());
-#ifdef ENABLE_INFINIOPS_API
+#if defined(ENABLE_INFINIOPS_API)    \
+    && (defined(ENABLE_NVIDIA_API)   \
+        || defined(ENABLE_METAX_API) \
+        || (defined(ENABLE_ILUVATAR_API) && defined(ENABLE_ATEN)))
     if (tryGreedyWithInfiniOps(indices, logits, topk)) {
         return;
     }

src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc

@@ -31,7 +31,10 @@ void *plan(const Tensor &positions,
            bool is_neox,
            int64_t rope_dim_offset,
            bool inverse) {
-    INFINICORE_ASSERT(query->device().getType() == Device::Type::NVIDIA || query->device().getType() == Device::Type::METAX);
+    const auto device_type = query->device().getType();
+    INFINICORE_ASSERT(device_type == Device::Type::NVIDIA
+                      || device_type == Device::Type::METAX
+                      || device_type == Device::Type::ILUVATAR);
     return new PlannedMeta{
         TensorMeta(positions),
         TensorMeta(query),
@@ -80,6 +83,9 @@ static bool registered = []() {
     RotaryEmbedding::plan_dispatcher().registerDevice(Device::Type::METAX, &plan);
     RotaryEmbedding::run_dispatcher().registerDevice(Device::Type::METAX, &run);
     RotaryEmbedding::cleanup_dispatcher().registerDevice(Device::Type::METAX, &cleanup);
+    RotaryEmbedding::plan_dispatcher().registerDevice(Device::Type::ILUVATAR, &plan);
+    RotaryEmbedding::run_dispatcher().registerDevice(Device::Type::ILUVATAR, &run);
+    RotaryEmbedding::cleanup_dispatcher().registerDevice(Device::Type::ILUVATAR, &cleanup);
     return true;
 }();
 

xmake.lua

@@ -416,14 +416,22 @@ local function build_infiniops_external(xmake_os)
         "-DGENERATE_PYTHON_BINDINGS=OFF",
         "-DCMAKE_BUILD_TYPE=Release"
     }
-    if has_config("nv-gpu") or has_config("metax-gpu") then
+    if has_config("nv-gpu")
+        or has_config("metax-gpu")
+        or (has_config("iluvatar-gpu") and has_config("aten")) then
         table.insert(cmake_config_args, "-DWITH_TORCH=ON")
         table.insert(cmake_config_args, "-DINFINI_OPS_TORCH_OPS=argmax")
     end
     if has_config("iluvatar-gpu") and has_config("aten") then
         table.insert(cmake_config_args, "-DTORCH_CXX11_ABI=0")
         table.insert(cmake_config_args, "-DCMAKE_CXX_FLAGS=-D_GLIBCXX_USE_CXX11_ABI=0")
     end
+    if has_config("iluvatar-gpu") then
+        local iluvatar_arch = get_config("iluvatar-arch")
+        if iluvatar_arch and iluvatar_arch ~= "" then
+            table.insert(cmake_config_args, "-DILUVATAR_ARCH=" .. iluvatar_arch)
+        end
+    end
     local infiniops_ops, with_linked_flash_attn_with_kvcache, with_linked_flash_attn_varlen_func = configure_infiniops_ops(os.getenv("INFINI_OPS_OPS"))
     if with_linked_flash_attn_with_kvcache or with_linked_flash_attn_varlen_func then
         table.insert(cmake_config_args, "-DWITH_LINKED=ON")
@@ -706,7 +714,7 @@ target("infinicore_cpp_api")
         add_defines("CHAR_BIT=8", "INT_MIN=(-2147483647 - 1)", "INT_MAX=2147483647", "UINT_MAX=4294967295U")
     end
     add_includedirs(INFINI_ROOT.."/include", { public = true })
-    if has_config("nv-gpu") or has_config("ali-ppu") then
+    if has_config("nv-gpu") or has_config("iluvatar-gpu") or has_config("ali-ppu") then
         local cuda_root = os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or get_config("cuda") or "/usr/local/cuda"
         add_includedirs(cuda_root .. "/include")
     end

xmake/iluvatar.lua

@@ -45,20 +45,20 @@ target("infinicore_cpp_api")
         add_defines("ENABLE_ILUVATAR_VENDOR_OPS")
     end
 
-    if iluvatar_flash_attn_enabled then
+    if has_config("aten") then
         before_link(function (target)
             local torch_dir = os.iorunv("python3", {
                 "-c",
                 "import torch, os; print(os.path.dirname(torch.__file__))",
             }):trim()
             local torch_lib_dir = path.join(torch_dir, "lib")
             if not os.isdir(torch_lib_dir) then
-                raise("Iluvatar Flash Attention: torch library directory not found: " .. torch_lib_dir)
+                raise("Iluvatar ATen: torch library directory not found: " .. torch_lib_dir)
             end
 
-            -- Flash Attention and vendor extensions use ATen symbols and may be
-            -- loaded after InfiniCore. Keep the complete Torch runtime discoverable
-            -- even when users import infinicore before importing torch.
+            -- ATen-backed InfiniOps operators, Flash Attention, and vendor
+            -- extensions use Torch symbols. Keep the complete Torch runtime
+            -- discoverable even when users import infinicore before importing torch.
             target:add(
                 "shflags",
                 "-Wl,--no-as-needed",
@@ -74,6 +74,10 @@ target("infinicore_cpp_api")
                 {force = true}
             )
 
+            if not iluvatar_flash_attn_enabled then
+                return
+            end
+
             local attention_so = iluvatar_attention_so_path(os.iorunv)
             print("Iluvatar attention extension: " .. attention_so)
             local attention_dir = path.directory(attention_so)