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 d293bbbf9..8decfb25e 100644 --- a/src/infinicore/nn/rope.cc +++ b/src/infinicore/nn/rope.cc @@ -81,6 +81,7 @@ void RoPE::initialize_cache() { if ((device_.getType() == Device::Type::NVIDIA || device_.getType() == Device::Type::METAX || device_.getType() == Device::Type::ILUVATAR + || device_.getType() == Device::Type::CAMBRICON || device_.getType() == Device::Type::HYGON) && !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 917047b5f..c450ccfbd 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())}; case Device::Type::HYGON: return infini::ops::Device{infini::ops::Device::Type::kHygon, static_cast(device.getIndex())}; default: @@ -68,6 +70,7 @@ inline bool isSupportedDevice(Device::Type device_type) { case Device::Type::METAX: case Device::Type::MOORE: case Device::Type::ILUVATAR: + case Device::Type::CAMBRICON: case Device::Type::HYGON: return true; default: @@ -81,6 +84,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); dispatcher.registerDevice(Device::Type::HYGON, function); } diff --git a/src/infinicore/ops/random_sample/random_sample.cc b/src/infinicore/ops/random_sample/random_sample.cc index 2041e4302..7502a3bc5 100644 --- a/src/infinicore/ops/random_sample/random_sample.cc +++ b/src/infinicore/ops/random_sample/random_sample.cc @@ -2,10 +2,11 @@ #include "../../utils.hpp" -#if defined(ENABLE_INFINIOPS_API) \ - && (defined(ENABLE_NVIDIA_API) \ - || defined(ENABLE_METAX_API) \ - || defined(ENABLE_HYGON_API) \ +#if defined(ENABLE_INFINIOPS_API) \ + && (defined(ENABLE_NVIDIA_API) \ + || defined(ENABLE_METAX_API) \ + || defined(ENABLE_HYGON_API) \ + || (defined(ENABLE_CAMBRICON_API) && defined(ENABLE_ATEN)) \ || (defined(ENABLE_ILUVATAR_API) && defined(ENABLE_ATEN))) #include "../infiniops_impl.hpp" @@ -15,16 +16,18 @@ namespace infinicore::op { namespace { -#if defined(ENABLE_INFINIOPS_API) \ - && (defined(ENABLE_NVIDIA_API) \ - || defined(ENABLE_METAX_API) \ - || defined(ENABLE_HYGON_API) \ +#if defined(ENABLE_INFINIOPS_API) \ + && (defined(ENABLE_NVIDIA_API) \ + || defined(ENABLE_METAX_API) \ + || defined(ENABLE_HYGON_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 && device_type != Device::Type::HYGON) || topk != 1 @@ -68,10 +71,11 @@ 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) \ - || defined(ENABLE_HYGON_API) \ +#if defined(ENABLE_INFINIOPS_API) \ + && (defined(ENABLE_NVIDIA_API) \ + || defined(ENABLE_METAX_API) \ + || defined(ENABLE_HYGON_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 2c08c9ec7..9d64f7f61 100644 --- a/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc +++ b/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc @@ -35,6 +35,7 @@ void *plan(const Tensor &positions, INFINICORE_ASSERT(device_type == Device::Type::NVIDIA || device_type == Device::Type::METAX || device_type == Device::Type::ILUVATAR + || device_type == Device::Type::CAMBRICON || device_type == Device::Type::HYGON); return new PlannedMeta{ TensorMeta(positions), @@ -87,6 +88,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); RotaryEmbedding::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); RotaryEmbedding::run_dispatcher().registerDevice(Device::Type::HYGON, &run); RotaryEmbedding::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); diff --git a/xmake.lua b/xmake.lua index 68f73a0ec..7ddfbf710 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,9 +471,10 @@ 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") add_backend("hygon-dcu", "-DWITH_HYGON=ON") if #enabled == 0 then - raise("InfiniOps integration requires one of --nv-gpu, --metax-gpu, --iluvatar-gpu, --moore-gpu, or --hygon-dcu") + raise("InfiniOps integration requires one of --nv-gpu, --metax-gpu, --iluvatar-gpu, --moore-gpu, --cambricon-mlu, or --hygon-dcu") end if #enabled > 1 then raise("InfiniOps can build only one GPU backend at a time") @@ -477,10 +499,16 @@ local function build_infiniops_external(xmake_os, json) "-DGENERATE_PYTHON_BINDINGS=OFF", "-DCMAKE_BUILD_TYPE=Release" } + if has_config("cambricon-mlu") then + -- 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("hygon-dcu") - 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") or has_config("hygon-dcu") then @@ -825,7 +853,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") @@ -834,6 +862,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") @@ -1088,7 +1150,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