From 38af5fe9473406a471c99ffc5cb0fb87af611d72 Mon Sep 17 00:00:00 2001 From: baominghelly <41820386+baominghelly@users.noreply.github.com> Date: Mon, 24 Aug 2026 15:43:46 +0800 Subject: [PATCH 1/2] feat(cambricon): integrate InfiniOps backend --- src/infinicore/context/runtime/runtime.cc | 3 + src/infinicore/nn/rope.cc | 3 +- src/infinicore/ops/infiniops_impl.hpp | 4 + .../ops/random_sample/random_sample.cc | 22 +++-- .../rotary_embedding_infiniops.cc | 6 +- xmake.lua | 90 ++++++++++++++++++- 6 files changed, 113 insertions(+), 15 deletions(-) diff --git a/src/infinicore/context/runtime/runtime.cc b/src/infinicore/context/runtime/runtime.cc index 65d244f41..1808b4d9c 100644 --- a/src/infinicore/context/runtime/runtime.cc +++ b/src/infinicore/context/runtime/runtime.cc @@ -111,6 +111,9 @@ void Runtime::memcpyH2D(void *dst, const void *src, size_t size, bool async) { } void Runtime::memcpyD2H(void *dst, const void *src, size_t size) { + if (device_.getType() == Device::Type::CAMBRICON) { + syncStream(); + } INFINICORE_CHECK_ERROR(infinirtMemcpy(dst, src, size, INFINIRT_MEMCPY_D2H)); } diff --git a/src/infinicore/nn/rope.cc b/src/infinicore/nn/rope.cc index 5f3f4f1cd..c92c09a36 100644 --- a/src/infinicore/nn/rope.cc +++ b/src/infinicore/nn/rope.cc @@ -80,7 +80,8 @@ void RoPE::initialize_cache() { #ifdef ENABLE_INFINIOPS_API if ((device_.getType() == Device::Type::NVIDIA || device_.getType() == Device::Type::METAX - || device_.getType() == Device::Type::ILUVATAR) + || device_.getType() == Device::Type::ILUVATAR + || device_.getType() == Device::Type::CAMBRICON) && !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 befb830cd..583a9ffdb 100644 --- a/src/infinicore/ops/infiniops_impl.hpp +++ b/src/infinicore/ops/infiniops_impl.hpp @@ -55,6 +55,8 @@ inline infini::ops::Device toInfiniOpsDevice(const Device &device) { return infini::ops::Device{infini::ops::Device::Type::kMoore, static_cast(device.getIndex())}; case Device::Type::ILUVATAR: return infini::ops::Device{infini::ops::Device::Type::kIluvatar, static_cast(device.getIndex())}; + case Device::Type::CAMBRICON: + return infini::ops::Device{infini::ops::Device::Type::kCambricon, static_cast(device.getIndex())}; default: throw std::runtime_error("InfiniOps backend does not support this device type."); } @@ -66,6 +68,7 @@ inline bool isSupportedDevice(Device::Type device_type) { case Device::Type::METAX: case Device::Type::MOORE: case Device::Type::ILUVATAR: + case Device::Type::CAMBRICON: return true; default: return false; @@ -78,6 +81,7 @@ void registerSupportedDevices(Dispatcher &dispatcher, Function function) { dispatcher.registerDevice(Device::Type::METAX, function); dispatcher.registerDevice(Device::Type::MOORE, function); dispatcher.registerDevice(Device::Type::ILUVATAR, function); + dispatcher.registerDevice(Device::Type::CAMBRICON, function); } struct TensorMeta { diff --git a/src/infinicore/ops/random_sample/random_sample.cc b/src/infinicore/ops/random_sample/random_sample.cc index fb2874c55..b3eb4b9c0 100644 --- a/src/infinicore/ops/random_sample/random_sample.cc +++ b/src/infinicore/ops/random_sample/random_sample.cc @@ -2,9 +2,10 @@ #include "../../utils.hpp" -#if defined(ENABLE_INFINIOPS_API) \ - && (defined(ENABLE_NVIDIA_API) \ - || defined(ENABLE_METAX_API) \ +#if defined(ENABLE_INFINIOPS_API) \ + && (defined(ENABLE_NVIDIA_API) \ + || defined(ENABLE_METAX_API) \ + || (defined(ENABLE_CAMBRICON_API) && defined(ENABLE_ATEN)) \ || (defined(ENABLE_ILUVATAR_API) && defined(ENABLE_ATEN))) #include "../infiniops_impl.hpp" @@ -14,15 +15,17 @@ namespace infinicore::op { namespace { -#if defined(ENABLE_INFINIOPS_API) \ - && (defined(ENABLE_NVIDIA_API) \ - || defined(ENABLE_METAX_API) \ +#if defined(ENABLE_INFINIOPS_API) \ + && (defined(ENABLE_NVIDIA_API) \ + || defined(ENABLE_METAX_API) \ + || (defined(ENABLE_CAMBRICON_API) && defined(ENABLE_ATEN)) \ || (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::CAMBRICON && device_type != Device::Type::ILUVATAR) || topk != 1 || logits->ndim() != 1 @@ -63,9 +66,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()); -#if defined(ENABLE_INFINIOPS_API) \ - && (defined(ENABLE_NVIDIA_API) \ - || defined(ENABLE_METAX_API) \ +#if defined(ENABLE_INFINIOPS_API) \ + && (defined(ENABLE_NVIDIA_API) \ + || defined(ENABLE_METAX_API) \ + || (defined(ENABLE_CAMBRICON_API) && defined(ENABLE_ATEN)) \ || (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 435f33e89..33adf5c76 100644 --- a/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc +++ b/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc @@ -34,7 +34,8 @@ void *plan(const Tensor &positions, const auto device_type = query->device().getType(); INFINICORE_ASSERT(device_type == Device::Type::NVIDIA || device_type == Device::Type::METAX - || device_type == Device::Type::ILUVATAR); + || device_type == Device::Type::ILUVATAR + || device_type == Device::Type::CAMBRICON); return new PlannedMeta{ TensorMeta(positions), TensorMeta(query), @@ -86,6 +87,9 @@ static bool registered = []() { RotaryEmbedding::plan_dispatcher().registerDevice(Device::Type::ILUVATAR, &plan); RotaryEmbedding::run_dispatcher().registerDevice(Device::Type::ILUVATAR, &run); RotaryEmbedding::cleanup_dispatcher().registerDevice(Device::Type::ILUVATAR, &cleanup); + RotaryEmbedding::plan_dispatcher().registerDevice(Device::Type::CAMBRICON, &plan); + RotaryEmbedding::run_dispatcher().registerDevice(Device::Type::CAMBRICON, &run); + RotaryEmbedding::cleanup_dispatcher().registerDevice(Device::Type::CAMBRICON, &cleanup); return true; }(); diff --git a/xmake.lua b/xmake.lua index d16d186b1..88c9de43f 100644 --- a/xmake.lua +++ b/xmake.lua @@ -439,6 +439,27 @@ local function configure_infiniops_ops(infiniops_ops, xmake_os, json) with_linked_flash_attn_varlen_func, nil end +local infiniops_adapter_dependencies = { + add_rms_norm = {"copy", "fused_add_rms_norm"}, + conv2d = {"convolution"}, + gelutanh = {"gelu"}, + paged_attention = {"flash_attn_with_kvcache"}, + paged_caching = {"reshape_and_cache_flash"}, + rearrange = {"copy"}, + topksoftmax = {"topk_softmax"} +} + +local function is_infiniops_adapter_selected(adapter_name, selected_ops) + local dependencies = infiniops_adapter_dependencies[adapter_name] + or {adapter_name} + for _, dependency in ipairs(dependencies) do + if not selected_ops[dependency] then + return false + end + end + return true +end + local function get_infiniops_backend_cmake_arg() local enabled = {} local function add_backend(config, cmake_arg) @@ -450,8 +471,9 @@ local function get_infiniops_backend_cmake_arg() add_backend("metax-gpu", "-DWITH_METAX=ON") add_backend("iluvatar-gpu", "-DWITH_ILUVATAR=ON") add_backend("moore-gpu", "-DWITH_MOORE=ON") + add_backend("cambricon-mlu", "-DWITH_CAMBRICON=ON") if #enabled == 0 then - raise("InfiniOps integration requires one of --nv-gpu, --metax-gpu, --iluvatar-gpu, or --moore-gpu") + raise("InfiniOps integration requires one of --nv-gpu, --metax-gpu, --iluvatar-gpu, --moore-gpu, or --cambricon-mlu") end if #enabled > 1 then raise("InfiniOps can build only one GPU backend at a time") @@ -476,9 +498,34 @@ local function build_infiniops_external(xmake_os, json) "-DGENERATE_PYTHON_BINDINGS=OFF", "-DCMAKE_BUILD_TYPE=Release" } + if has_config("cambricon-mlu") then + -- Avoid stale finder and optional-provider state when reusing an + -- InfiniOps build tree previously configured for another backend. + table.insert(cmake_config_args, "-UINFINI_RT_INCLUDE_DIRS") + table.insert(cmake_config_args, "-UINFINI_RT_LIBRARY") + table.insert(cmake_config_args, "-U_INFINI_RT_INCLUDE_DIR") + table.insert(cmake_config_args, "-U_INFINI_RT_LIBRARY") + table.insert(cmake_config_args, "-DWITH_LINKED=OFF") + if has_config("aten") then + local torch_cxx11_abi = xmake_os.iorunv(PYTHON, { + "-c", + "import torch; print(1 if torch._C._GLIBCXX_USE_CXX11_ABI else 0)", + }):trim() + table.insert( + cmake_config_args, + "-DCMAKE_CXX_FLAGS=-D_GLIBCXX_USE_CXX11_ABI=" .. torch_cxx11_abi) + else + table.insert(cmake_config_args, "-DWITH_TORCH=OFF") + end + + -- Cambricon headers define host half helpers with external linkage, + -- so generated call instantiations must stay in one translation unit. + xmake_os.setenv("INFINI_OPS_DISPATCH_BATCH_SIZE", "64") + end if has_config("nv-gpu") or has_config("metax-gpu") - or (has_config("iluvatar-gpu") and has_config("aten")) then + or (has_config("iluvatar-gpu") and has_config("aten")) + or (has_config("cambricon-mlu") and has_config("aten")) then table.insert(cmake_config_args, "-DWITH_TORCH=ON") local torch_ops = "argmax" if has_config("nv-gpu") or has_config("metax-gpu") then @@ -823,7 +870,7 @@ target("infinicore_cpp_api") add_rpathdirs(INFINI_ROOT .. "/lib") on_load(function (target) local json = import("core.base.json") - local _, with_linked_flash_attn_with_kvcache, with_linked_flash_attn_varlen_func = configure_infiniops_ops( + local selected_ops, with_linked_flash_attn_with_kvcache, with_linked_flash_attn_varlen_func = configure_infiniops_ops( os.getenv("INFINI_OPS_OPS"), os, json) if with_linked_flash_attn_with_kvcache then target:add("defines", "ENABLE_INFINIOPS_LINKED_FLASH_ATTN_WITH_KVCACHE") @@ -832,6 +879,40 @@ target("infinicore_cpp_api") target:add("defines", "ENABLE_INFINIOPS_LINKED_FLASH_ATTN_VARLEN_FUNC") end build_infiniops_external(os, json) + if has_config("cambricon-mlu") and has_config("aten") then + local torch_mlu_dir = os.iorunv(PYTHON, { + "-c", + "import torch_mlu, os; print(os.path.dirname(torch_mlu.__file__))", + }):trim() + local torch_cxx11_abi = os.iorunv(PYTHON, { + "-c", + "import torch; print(1 if torch._C._GLIBCXX_USE_CXX11_ABI else 0)", + }):trim() + + target:add( + "defines", + "_GLIBCXX_USE_CXX11_ABI=" .. torch_cxx11_abi) + target:add( + "includedirs", + path.join(torch_mlu_dir, "csrc"), + path.join(torch_mlu_dir, "csrc", "include"), + {public = true}) + + if selected_ops then + for _, adapter_file in ipairs(os.files("src/infinicore/ops/*/*_infiniops.cc")) do + local adapter_name = + path.filename(adapter_file):match("^(.-)_infiniops%.cc$") + if not is_infiniops_adapter_selected(adapter_name, + selected_ops) then + target:remove("files", adapter_file) + end + end + else + target:remove( + "files", + "src/infinicore/ops/paged_attention/paged_attention_infiniops.cc") + end + end end) after_install(function (target) local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") @@ -1086,7 +1167,8 @@ target("infinicore_cpp_api") else remove_files("src/infinicore/ops/*/hygon/*.cc") end - if has_config("infiniops") and not has_config("nv-gpu") then + if has_config("infiniops") and not has_config("nv-gpu") + and not has_config("cambricon-mlu") then remove_files("src/infinicore/ops/paged_attention/paged_attention_infiniops.cc") end if has_config("mutual-awareness") then From 79425c11a3811ee4f3eb8faa96bc0ed9dcaecfc4 Mon Sep 17 00:00:00 2001 From: baominghelly <41820386+baominghelly@users.noreply.github.com> Date: Mon, 24 Aug 2026 22:17:52 +0800 Subject: [PATCH 2/2] refactor(build): simplify Cambricon InfiniOps configuration --- xmake.lua | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/xmake.lua b/xmake.lua index 88c9de43f..b7aef1632 100644 --- a/xmake.lua +++ b/xmake.lua @@ -499,25 +499,6 @@ local function build_infiniops_external(xmake_os, json) "-DCMAKE_BUILD_TYPE=Release" } if has_config("cambricon-mlu") then - -- Avoid stale finder and optional-provider state when reusing an - -- InfiniOps build tree previously configured for another backend. - table.insert(cmake_config_args, "-UINFINI_RT_INCLUDE_DIRS") - table.insert(cmake_config_args, "-UINFINI_RT_LIBRARY") - table.insert(cmake_config_args, "-U_INFINI_RT_INCLUDE_DIR") - table.insert(cmake_config_args, "-U_INFINI_RT_LIBRARY") - table.insert(cmake_config_args, "-DWITH_LINKED=OFF") - if has_config("aten") then - local torch_cxx11_abi = xmake_os.iorunv(PYTHON, { - "-c", - "import torch; print(1 if torch._C._GLIBCXX_USE_CXX11_ABI else 0)", - }):trim() - table.insert( - cmake_config_args, - "-DCMAKE_CXX_FLAGS=-D_GLIBCXX_USE_CXX11_ABI=" .. torch_cxx11_abi) - else - table.insert(cmake_config_args, "-DWITH_TORCH=OFF") - end - -- Cambricon headers define host half helpers with external linkage, -- so generated call instantiations must stay in one translation unit. xmake_os.setenv("INFINI_OPS_DISPATCH_BATCH_SIZE", "64")