From b1284e89c06a95cded08102edc49e4c3b3c850cb Mon Sep 17 00:00:00 2001 From: gongchensu Date: Mon, 17 Aug 2026 07:52:00 +0000 Subject: [PATCH 1/2] feat(iluvatar): route Qwen3 ops through InfiniOps --- src/infinicore/nn/rope.cc | 20 ++++++++++++++++++- .../ops/random_sample/random_sample.cc | 18 +++++++++++++---- .../rotary_embedding_infiniops.cc | 8 +++++++- 3 files changed, 40 insertions(+), 6 deletions(-) diff --git a/src/infinicore/nn/rope.cc b/src/infinicore/nn/rope.cc index 6c7e52ad4..5f3f4f1cd 100644 --- a/src/infinicore/nn/rope.cc +++ b/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 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(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); diff --git a/src/infinicore/ops/random_sample/random_sample.cc b/src/infinicore/ops/random_sample/random_sample.cc index 2ff505c5a..fb2874c55 100644 --- a/src/infinicore/ops/random_sample/random_sample.cc +++ b/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; } diff --git a/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc b/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc index 5c1c5cb19..435f33e89 100644 --- a/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc +++ b/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; }(); From 91d27017fb2297b6d5b7d3ff44569de15273dee7 Mon Sep 17 00:00:00 2001 From: gongchensu Date: Mon, 17 Aug 2026 07:52:00 +0000 Subject: [PATCH 2/2] fix(iluvatar): build CoreX ATen InfiniOps --- xmake.lua | 12 ++++++++++-- xmake/iluvatar.lua | 14 +++++++++----- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/xmake.lua b/xmake.lua index b6806dc72..79b22ba46 100644 --- a/xmake.lua +++ b/xmake.lua @@ -416,7 +416,9 @@ 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 @@ -424,6 +426,12 @@ local function build_infiniops_external(xmake_os) 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 diff --git a/xmake/iluvatar.lua b/xmake/iluvatar.lua index 38602357a..09e7043f4 100644 --- a/xmake/iluvatar.lua +++ b/xmake/iluvatar.lua @@ -45,7 +45,7 @@ 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", @@ -53,12 +53,12 @@ target("infinicore_cpp_api") }):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)