InfiniTensor · voltjia · Jun 3, 2026
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -38,6 +38,49 @@ option(GENERATE_OPERATOR_CALL_INSTANTIATIONS
     "Generate explicit operator call instantiations" ON)
 option(GENERATE_PYTHON_BINDINGS "Generate Python bindings" OFF)
 
+set(INFINI_OPS_PLUGINS "" CACHE STRING
+    "Comma- or semicolon-separated infini_ops build-time plugins to enable")
+
+function(_infini_ops_enable_legacy_options_from_plugins)
+    if(NOT INFINI_OPS_PLUGINS)
+        return()
+    endif()
+
+    set(_requested_plugins "${INFINI_OPS_PLUGINS}")
+    string(REPLACE "," ";" _requested_plugins "${_requested_plugins}")
+    foreach(_plugin IN LISTS _requested_plugins)
+        string(STRIP "${_plugin}" _plugin)
+        if(_plugin STREQUAL "")
+            continue()
+        elseif(_plugin STREQUAL "cpu")
+            set(WITH_CPU ON CACHE BOOL "Enable CPU backend" FORCE)
+        elseif(_plugin STREQUAL "nvidia")
+            set(WITH_NVIDIA ON CACHE BOOL "Enable CUDA backend" FORCE)
+        elseif(_plugin STREQUAL "iluvatar")
+            set(WITH_ILUVATAR ON CACHE BOOL "Enable Iluvatar GPU backend" FORCE)
+        elseif(_plugin STREQUAL "hygon")
+            set(WITH_HYGON ON CACHE BOOL "Enable Hygon GPU backend" FORCE)
+        elseif(_plugin STREQUAL "metax")
+            set(WITH_METAX ON CACHE BOOL "Enable MetaX backend" FORCE)
+        elseif(_plugin STREQUAL "cambricon")
+            set(WITH_CAMBRICON ON CACHE BOOL "Enable Cambricon backend" FORCE)
+        elseif(_plugin STREQUAL "moore")
+            set(WITH_MOORE ON CACHE BOOL "Enable Moore backend" FORCE)
+        elseif(_plugin STREQUAL "ascend")
+            set(WITH_ASCEND ON CACHE BOOL "Enable Ascend backend" FORCE)
+        elseif(_plugin STREQUAL "cuda-common")
+            # Shared dependency plugin; no legacy device option to set.
+        else()
+            message(FATAL_ERROR
+                "Unknown infini_ops plugin `${_plugin}`. v1 supports built-in plugins: "
+                "`cpu`, `nvidia`, `iluvatar`, `hygon`, `metax`, `moore`, `cambricon`, `ascend`, `cuda-common`")
+        endif()
+    endforeach()
+endfunction()
+
+_infini_ops_enable_legacy_options_from_plugins()
+
+
 set(_DEFAULT_HYGON_DTK_ROOT "/opt/dtk")
 
 function(_infiniops_find_hygon_cuda_root out_var dtk_root)

diff --git a/cmake/infini_ops_plugins.cmake b/cmake/infini_ops_plugins.cmake
@@ -0,0 +1,253 @@
+include_guard(GLOBAL)
+
+set(INFINI_OPS_PLUGINS "" CACHE STRING
+    "Comma- or semicolon-separated infini_ops build-time plugins to enable")
+set(INFINI_OPS_PLUGIN_ROOT "${PROJECT_SOURCE_DIR}/plugins" CACHE PATH
+    "Directory containing infini_ops build-time plugins")
+set(INFINI_OPS_PLUGIN_CONTRACT_VERSION 1)
+
+set(_INFINI_OPS_KNOWN_DEVICE_PLUGINS
+    cpu nvidia iluvatar hygon metax moore cambricon ascend)
+
+function(_infini_ops_append_unique_global property_name)
+    get_property(_values GLOBAL PROPERTY "${property_name}")
+    foreach(_value ${ARGN})
+        if("${_value}" STREQUAL "")
+            continue()
+        endif()
+
+        list(FIND _values "${_value}" _index)
+        if(_index EQUAL -1)
+            set_property(GLOBAL APPEND PROPERTY "${property_name}" "${_value}")
+        endif()
+    endforeach()
+endfunction()
+
+function(infini_ops_register_plugin)
+    set(_one_value_args NAME KIND CONTRACT_VERSION CMAKE_ENTRY)
+    set(_multi_value_args
+        DEVICES
+        DEPENDS
+        SOURCE_ROOTS
+        OPERATOR_ROOTS
+        DEVICE_HEADERS
+        TEST_DEVICES)
+    cmake_parse_arguments(ARG "" "${_one_value_args}" "${_multi_value_args}" ${ARGN})
+
+    foreach(_required NAME KIND CONTRACT_VERSION CMAKE_ENTRY)
+        if(NOT ARG_${_required})
+            message(FATAL_ERROR "`infini_ops_register_plugin` missing `${_required}`")
+        endif()
+    endforeach()
+
+    if(NOT ARG_KIND STREQUAL "shared" AND NOT ARG_KIND STREQUAL "device")
+        message(FATAL_ERROR "infini_ops plugin `${ARG_NAME}` has invalid `kind`: `${ARG_KIND}`")
+    endif()
+
+    if(NOT "${ARG_CONTRACT_VERSION}" STREQUAL "${INFINI_OPS_PLUGIN_CONTRACT_VERSION}")
+        message(FATAL_ERROR
+            "infini_ops plugin `${ARG_NAME}` uses contract `${ARG_CONTRACT_VERSION}`; "
+            "expected `${INFINI_OPS_PLUGIN_CONTRACT_VERSION}`")
+    endif()
+
+    foreach(_device IN LISTS ARG_DEVICES)
+        list(FIND _INFINI_OPS_KNOWN_DEVICE_PLUGINS "${_device}" _known_index)
+        if(_known_index EQUAL -1)
+            message(FATAL_ERROR
+                "infini_ops plugin `${ARG_NAME}` declares unknown device `${_device}`")
+        endif()
+    endforeach()
+
+    if(ARG_KIND STREQUAL "device" AND NOT ARG_DEVICES)
+        message(FATAL_ERROR "infini_ops device plugin `${ARG_NAME}` must declare `DEVICES`")
+    endif()
+
+    if(ARG_KIND STREQUAL "shared" AND ARG_DEVICES)
+        message(FATAL_ERROR "infini_ops shared plugin `${ARG_NAME}` must not declare `DEVICES`")
+    endif()
+
+    _infini_ops_append_unique_global(INFINI_OPS_PLUGIN_NAMES "${ARG_NAME}")
+    _infini_ops_append_unique_global(INFINI_OPS_PLUGIN_DEVICES ${ARG_DEVICES})
+    _infini_ops_append_unique_global(INFINI_OPS_PLUGIN_SOURCE_ROOTS ${ARG_SOURCE_ROOTS})
+    _infini_ops_append_unique_global(INFINI_OPS_PLUGIN_OPERATOR_ROOTS ${ARG_OPERATOR_ROOTS})
+    _infini_ops_append_unique_global(INFINI_OPS_PLUGIN_DEVICE_HEADERS ${ARG_DEVICE_HEADERS})
+    _infini_ops_append_unique_global(INFINI_OPS_PLUGIN_TEST_DEVICES ${ARG_TEST_DEVICES})
+endfunction()
+
+function(infini_ops_register_device)
+    infini_ops_register_plugin(
+        KIND device
+        CONTRACT_VERSION ${INFINI_OPS_PLUGIN_CONTRACT_VERSION}
+        ${ARGN})
+endfunction()
+
+function(infini_ops_enable_plugin name)
+    get_property(_loaded GLOBAL PROPERTY INFINI_OPS_PLUGIN_LOADED)
+    list(FIND _loaded "${name}" _loaded_index)
+    if(NOT _loaded_index EQUAL -1)
+        return()
+    endif()
+
+    get_property(_loading GLOBAL PROPERTY INFINI_OPS_PLUGIN_LOADING_STACK)
+    list(FIND _loading "${name}" _loading_index)
+    if(NOT _loading_index EQUAL -1)
+        list(APPEND _loading "${name}")
+        string(REPLACE ";" " -> " _cycle "${_loading}")
+        message(FATAL_ERROR "infini_ops plugin dependency cycle detected: `${_cycle}`")
+    endif()
+
+    set(_entry_path "${INFINI_OPS_PLUGIN_ROOT}/${name}/plugin.cmake")
+    if(NOT EXISTS "${_entry_path}")
+        message(FATAL_ERROR "infini_ops plugin `${name}` `CMake` entry not found: `${_entry_path}`")
+    endif()
+
+    set_property(GLOBAL APPEND PROPERTY INFINI_OPS_PLUGIN_LOADING_STACK "${name}")
+    include("${_entry_path}")
+    get_property(_loading GLOBAL PROPERTY INFINI_OPS_PLUGIN_LOADING_STACK)
+    list(REMOVE_ITEM _loading "${name}")
+    set_property(GLOBAL PROPERTY INFINI_OPS_PLUGIN_LOADING_STACK "${_loading}")
+
+    get_property(_registered GLOBAL PROPERTY INFINI_OPS_PLUGIN_NAMES)
+    list(FIND _registered "${name}" _registered_index)
+    if(_registered_index EQUAL -1)
+        message(FATAL_ERROR "infini_ops plugin `${name}` did not call `infini_ops_register_plugin`")
+    endif()
+
+    set_property(GLOBAL APPEND PROPERTY INFINI_OPS_PLUGIN_LOADED "${name}")
+endfunction()
+
+function(infini_ops_enable_requested_plugins)
+    set(_requested)
+
+    if(INFINI_OPS_PLUGINS)
+        set(_raw_plugins "${INFINI_OPS_PLUGINS}")
+        string(REPLACE "," ";" _raw_plugins "${_raw_plugins}")
+        foreach(_plugin IN LISTS _raw_plugins)
+            string(STRIP "${_plugin}" _plugin)
+            if(NOT _plugin STREQUAL "")
+                list(APPEND _requested "${_plugin}")
+            endif()
+        endforeach()
+    endif()
+
+    if(WITH_CPU)
+        list(APPEND _requested cpu)
+    endif()
+    if(WITH_NVIDIA)
+        list(APPEND _requested nvidia)
+    endif()
+    if(WITH_ILUVATAR)
+        list(APPEND _requested iluvatar)
+    endif()
+    if(WITH_HYGON)
+        list(APPEND _requested hygon)
+    endif()
+    if(WITH_METAX)
+        list(APPEND _requested metax)
+    endif()
+    if(WITH_MOORE)
+        list(APPEND _requested moore)
+    endif()
+    if(WITH_CAMBRICON)
+        list(APPEND _requested cambricon)
+    endif()
+    if(WITH_ASCEND)
+        list(APPEND _requested ascend)
+    endif()
+
+    if(_requested)
+        list(REMOVE_DUPLICATES _requested)
+    else()
+        list(APPEND _requested cpu)
+        set(WITH_CPU ON CACHE BOOL "Enable CPU backend" FORCE)
+    endif()
+
+    foreach(_plugin IN LISTS _requested)
+        infini_ops_enable_plugin("${_plugin}")
+    endforeach()
+endfunction()
+
+function(infini_ops_get_enabled_devices out_var)
+    get_property(_devices GLOBAL PROPERTY INFINI_OPS_PLUGIN_DEVICES)
+    if(NOT _devices)
+        set(_devices)
+    endif()
+    set(${out_var} ${_devices} PARENT_SCOPE)
+endfunction()
+
+function(_infini_ops_json_escape value out_var)
+    string(REPLACE "\\" "\\\\" _escaped "${value}")
+    string(REPLACE "\"" "\\\"" _escaped "${_escaped}")
+    set(${out_var} "${_escaped}" PARENT_SCOPE)
+endfunction()
+
+function(_infini_ops_append_json_array path field trailing_comma)
+    file(APPEND "${path}" "  \"${field}\": [")
+    set(_first TRUE)
+    foreach(_value ${ARGN})
+        if(_first)
+            set(_first FALSE)
+        else()
+            file(APPEND "${path}" ", ")
+        endif()
+        _infini_ops_json_escape("${_value}" _escaped)
+        file(APPEND "${path}" "\"${_escaped}\"")
+    endforeach()
+    file(APPEND "${path}" "]")
+    if(trailing_comma)
+        file(APPEND "${path}" ",")
+    endif()
+    file(APPEND "${path}" "\n")
+endfunction()
+
+function(_infini_ops_append_json_map path field trailing_comma)
+    file(APPEND "${path}" "  \"${field}\": {")
+    set(_first TRUE)
+    foreach(_entry ${ARGN})
+        string(FIND "${_entry}" "=" _equals)
+        if(_equals EQUAL -1)
+            message(FATAL_ERROR "Invalid infini_ops plugin map entry `${_entry}`")
+        endif()
+        string(SUBSTRING "${_entry}" 0 ${_equals} _key)
+        math(EXPR _value_start "${_equals} + 1")
+        string(SUBSTRING "${_entry}" ${_value_start} -1 _value)
+
+        if(_first)
+            set(_first FALSE)
+        else()
+            file(APPEND "${path}" ",")
+        endif()
+        _infini_ops_json_escape("${_key}" _escaped_key)
+        _infini_ops_json_escape("${_value}" _escaped_value)
+        file(APPEND "${path}" "\n    \"${_escaped_key}\": \"${_escaped_value}\"")
+    endforeach()
+    if(NOT _first)
+        file(APPEND "${path}" "\n  ")
+    endif()
+    file(APPEND "${path}" "}")
+    if(trailing_comma)
+        file(APPEND "${path}" ",")
+    endif()
+    file(APPEND "${path}" "\n")
+endfunction()
+
+function(infini_ops_write_plugin_registry path)
+    get_property(_plugins GLOBAL PROPERTY INFINI_OPS_PLUGIN_NAMES)
+    get_property(_devices GLOBAL PROPERTY INFINI_OPS_PLUGIN_DEVICES)
+    get_property(_source_roots GLOBAL PROPERTY INFINI_OPS_PLUGIN_SOURCE_ROOTS)
+    get_property(_operator_roots GLOBAL PROPERTY INFINI_OPS_PLUGIN_OPERATOR_ROOTS)
+    get_property(_device_headers GLOBAL PROPERTY INFINI_OPS_PLUGIN_DEVICE_HEADERS)
+    get_property(_test_devices GLOBAL PROPERTY INFINI_OPS_PLUGIN_TEST_DEVICES)
+
+    file(WRITE "${path}" "{\n")
+    _infini_ops_append_json_array("${path}" "plugins" TRUE ${_plugins})
+    _infini_ops_append_json_array("${path}" "devices" TRUE ${_devices})
+    _infini_ops_append_json_array("${path}" "source_roots" TRUE ${_source_roots})
+    _infini_ops_append_json_array("${path}" "operator_roots" TRUE ${_operator_roots})
+    _infini_ops_append_json_map("${path}" "device_headers" TRUE ${_device_headers})
+    _infini_ops_append_json_map("${path}" "test_devices" FALSE ${_test_devices})
+    file(APPEND "${path}" "}\n")
+
+    message(STATUS "infini_ops plugins: `${_plugins}`")
+    message(STATUS "infini_ops plugin devices: `${_devices}`")
+endfunction()
diff --git a/plugins/ascend/plugin.cmake b/plugins/ascend/plugin.cmake
@@ -0,0 +1,84 @@
+infini_ops_register_device(
+    NAME ascend
+    CMAKE_ENTRY plugin.cmake
+    DEVICES ascend
+    SOURCE_ROOTS src/native/ascend
+    OPERATOR_ROOTS src/native/ascend/ops
+    DEVICE_HEADERS ascend=native/ascend/device_.h
+    TEST_DEVICES ascend=npu)
+
+# ASCEND_HOME is set by the top-level CMakeLists.txt.
+file(GLOB_RECURSE ASCEND_SOURCES CONFIGURE_DEPENDS
+    "${INFINI_OPS_SRC_DIR}/native/ascend/*.cc"
+    "${INFINI_OPS_SRC_DIR}/native/ascend/*.cpp")
+# Exclude `kernel_impl.cpp`: `AscendC` device code, not compiled by the host C++ compiler.
+list(FILTER ASCEND_SOURCES EXCLUDE REGEX ".*kernel_impl\\.cpp$")
+# Exclude custom/: standalone PyTorch extension, built separately.
+list(FILTER ASCEND_SOURCES EXCLUDE REGEX ".*/custom/.*")
+
+target_compile_definitions(infiniops PUBLIC WITH_ASCEND=1)
+target_sources(infiniops PRIVATE ${ASCEND_SOURCES})
+
+# Resolve the driver lib dir two levels above the toolkit root.
+get_filename_component(ASCEND_ROOT "${ASCEND_HOME}/../.." ABSOLUTE)
+
+# Prefer the real driver HAL; fall back to the toolkit stub for build-only
+# environments (e.g., Docker CI images without hardware drivers installed).
+# CANN <= 8.0: stub at runtime/lib64/stub/; CANN >= 8.5: devlib/<arch>-linux/devlib/.
+set(ASCEND_HAL_REAL   "${ASCEND_ROOT}/driver/lib64/driver/libascend_hal.so")
+set(ASCEND_HAL_STUB   "${ASCEND_HOME}/runtime/lib64/stub/libascend_hal.so")
+set(ASCEND_HAL_DEVLIB "${ASCEND_HOME}/${CMAKE_SYSTEM_PROCESSOR}-linux/devlib/libascend_hal.so")
+if(EXISTS "${ASCEND_HAL_REAL}")
+    set(ASCEND_HAL_LIB "${ASCEND_HAL_REAL}")
+elseif(EXISTS "${ASCEND_HAL_STUB}")
+    set(ASCEND_HAL_LIB "${ASCEND_HAL_STUB}")
+    message(STATUS "ascend_hal: driver not found, using stub for linking")
+elseif(EXISTS "${ASCEND_HAL_DEVLIB}")
+    set(ASCEND_HAL_LIB "${ASCEND_HAL_DEVLIB}")
+    message(STATUS "ascend_hal: driver not found, using devlib for linking")
+else()
+    message(FATAL_ERROR "libascend_hal.so not found (tried ${ASCEND_HAL_REAL}, ${ASCEND_HAL_STUB}, and ${ASCEND_HAL_DEVLIB})")
+endif()
+
+target_include_directories(infiniops PUBLIC
+    "${ASCEND_HOME}/include"
+    "${ASCEND_HOME}/include/aclnn"
+    "${ASCEND_HOME}/include/aclnnop")
+target_link_libraries(infiniops PUBLIC
+    "${ASCEND_HOME}/lib64/libascendcl.so"
+    "${ASCEND_HOME}/lib64/libnnopbase.so"
+    "${ASCEND_HOME}/lib64/libopapi.so"
+    "${ASCEND_HAL_LIB}")
+
+# ATB (Ascend Transformer Boost) provides fused operators like
+# `PagedAttention` and `ReshapeAndCache` that are graph-capture safe.
+set(ATB_HOME_DIR "$ENV{ATB_HOME_PATH}")
+if(NOT ATB_HOME_DIR)
+    # Default search path under CANN nnal directory.
+    file(GLOB ATB_SEARCH_DIRS "/usr/local/Ascend/nnal/atb/*/atb/cxx_abi_1")
+    if(ATB_SEARCH_DIRS)
+        list(SORT ATB_SEARCH_DIRS ORDER DESCENDING)
+        list(GET ATB_SEARCH_DIRS 0 ATB_HOME_DIR)
+    endif()
+endif()
+
+if(ATB_HOME_DIR AND EXISTS "${ATB_HOME_DIR}/include/atb/operation.h")
+    message(STATUS "ATB found: ${ATB_HOME_DIR}")
+    target_compile_definitions(infiniops PUBLIC INFINI_HAS_ATB=1)
+    target_include_directories(infiniops PUBLIC "${ATB_HOME_DIR}/include")
+    target_link_libraries(infiniops PUBLIC "${ATB_HOME_DIR}/lib/libatb.so")
+else()
+    message(STATUS "ATB not found - ATB-based operators disabled")
+endif()
+
+# Custom `AscendC` kernels (PyTorch extension, requires `torch_npu`).
+if(BUILD_CUSTOM_KERNEL)
+    add_subdirectory(
+        "${INFINI_OPS_SRC_DIR}/native/ascend/custom"
+        "${CMAKE_CURRENT_BINARY_DIR}/native/ascend/custom")
+
+    # Link the compiled `AscendC` kernel objects into `infiniops` so that
+    # custom kernel implementations (e.g. `RmsNorm` index 1) can call
+    # them via the generated launch functions.
+    target_compile_definitions(infiniops PUBLIC INFINI_HAS_CUSTOM_KERNELS=1)
+endif()