Skip to content

feat(cambricon): integrate InfiniOps backend #1522

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
voltjia merged 3 commits into from
Aug 25, 2026
Merged

feat(cambricon): integrate InfiniOps backend #1522

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
38af5fe
feat(cambricon): integrate InfiniOps backend
baominghelly Aug 24, 2026
79425c1
refactor(build): simplify Cambricon InfiniOps configuration
baominghelly Aug 24, 2026
5697d86
merge: resolve conflicts with main
baominghelly Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/infinicore/context/runtime/runtime.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
}

Expand Down
1 change: 1 addition & 0 deletions src/infinicore/nn/rope.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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_));
Expand Down
4 changes: 4 additions & 0 deletions src/infinicore/ops/infiniops_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ inline infini::ops::Device toInfiniOpsDevice(const Device &device) {
return infini::ops::Device{infini::ops::Device::Type::kMoore, static_cast<int>(device.getIndex())};
case Device::Type::ILUVATAR:
return infini::ops::Device{infini::ops::Device::Type::kIluvatar, static_cast<int>(device.getIndex())};
case Device::Type::CAMBRICON:
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())};
default:
Expand All @@ -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:
Expand All @@ -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);
}

Expand Down
28 changes: 16 additions & 12 deletions src/infinicore/ops/random_sample/random_sample.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand All @@ -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
Expand Down Expand Up @@ -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;
Expand Down
4 changes: 4 additions & 0 deletions src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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);
Expand Down
71 changes: 67 additions & 4 deletions xmake.lua
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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")
Expand All @@ -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
Expand Down Expand Up @@ -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")
Expand All @@ -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")
Expand Down Expand Up @@ -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
Expand Down
Loading