Merge pull request #1525 from cen121212/csx-main-fix

【WIP】main分支适配NPU

Author: ovowei

config.json

@@ -1,3 +1,23 @@
+option(KTRANSFORMERS_USE_NPU                 "ktransformers: use NPU"                           OFF)
+if(KTRANSFORMERS_USE_NPU)
+    add_definitions(-DKTRANSFORMERS_USE_NPU=1)
+endif()
+
+if(KTRANSFORMERS_USE_NPU)
+    set(ASCEND_HOME_PATH "$ENV{ASCEND_HOME_PATH}")
+    message(STATUS "ASCEND_HOME_PATH is ${ASCEND_HOME_PATH}")
+    include_directories(${ASCEND_HOME_PATH}/include)
+    
+    link_directories(${TORCH_INSTALL_PREFIX}/../torch.libs)
+    # find torch_npu
+    execute_process(
+            COMMAND python -c "import torch; import torch_npu; print(torch_npu.__path__[0])"
+            OUTPUT_VARIABLE TORCH_NPU_PATH
+            OUTPUT_STRIP_TRAILING_WHITESPACE
+    )
+    message(STATUS "Found PTA at: ${TORCH_NPU_PATH}")
+    find_library(PTA_LIBRARY torch_npu PATH "${TORCH_NPU_PATH}/lib")
+endif()
 
 cmake_minimum_required(VERSION 3.21)
 find_program(GCC_COMPILER NAMES g++-13 g++-12 g++-11 g++ REQUIRED)

csrc/balance_serve/CMakeLists.txt

@@ -73,8 +73,19 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../../third_party)
 
 # include_directories(/usr/include/tbb)
 # link_directories(/usr/lib64)
+option(KTRANSFORMERS_USE_NPU                 "ktransformers: use NPU"                           OFF)
+if(KTRANSFORMERS_USE_NPU)
+    add_definitions(-DKTRANSFORMERS_USE_NPU=1)
+endif()
 find_package(TBB REQUIRED)
-find_package(CUDA REQUIRED)
+
+if(KTRANSFORMERS_USE_NPU)
+    # NPU 构建
+    # find_package(CUDA REQUIRED)  # NPU 情况不需要 CUDA
+else()
+    # GPU 构建
+    find_package(CUDA REQUIRED)
+endif()
 
 # find_package(prometheus-cpp CONFIG REQUIRED)
 if(NOT TARGET prometheus-cpp::pull)
@@ -83,18 +94,29 @@ else()
     message(STATUS "prometheus Found!")
 endif()
 
-if(CUDA_FOUND)
-    message(STATUS "CUDA Found!")
-    message(STATUS "CUDA Version: ${CUDA_VERSION_STRING}")
-    message(STATUS "CUDA Toolkit Root: ${CUDA_TOOLKIT_ROOT_DIR}")
+if(KTRANSFORMERS_USE_NPU)
+    # NPU 情况下不检查 CUDA
 else()
-    message(FATAL_ERROR "CUDA not found!")
+    if(CUDA_FOUND)
+        message(STATUS "CUDA Found!")
+        message(STATUS "CUDA Version: ${CUDA_VERSION_STRING}")
+        message(STATUS "CUDA Toolkit Root: ${CUDA_TOOLKIT_ROOT_DIR}")
+    else()
+        message(FATAL_ERROR "CUDA not found!")
+    endif()
 endif()
 
 add_subdirectory(src)
 
 if(BUILD_TEST)
-    add_subdirectory(test)
+    if(KTRANSFORMERS_USE_NPU)
+        message(STATUS "Build test...")
+        set(THIRD_PARTY_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party)
+        add_subdirectory(${THIRD_PARTY_DIR}/spdlog ${CMAKE_CURRENT_SOURCE_DIR}/../build/third_party/spdlog)
+        add_subdirectory(test)
+    else()
+        add_subdirectory(test)
+    endif()
 endif()
 
 message(STATUS "BUILD_PYTHON_EXT: ${BUILD_PYTHON_EXT}")
@@ -126,7 +148,12 @@ endif()
 
 set(PHOTON_CXX_STANDARD 14 CACHE INTERNAL "C++ standard")
 
-set(CMAKE_CXX_FLAGS "-O3 -march=native")
+if(KTRANSFORMERS_USE_NPU)
+    set(CMAKE_CXX_FLAGS "-O3 -march=armv8.2-a")
+else()
+    set(CMAKE_CXX_FLAGS "-O3 -march=native")
+endif()
+
 message(STATUS "CMAKE_CXX_FLAGS of PhotonLibOS: ${CMAKE_CXX_FLAGS}")
 
 add_subdirectory(${THIRD_PARTY_DIR}/PhotonLibOS ${THIRD_PARTY_BUILD_DIR}/PhotonLibOS)

csrc/balance_serve/kvc2/CMakeLists.txt

@@ -23,24 +23,46 @@ target_link_libraries(cache_entry PUBLIC gpu_cache)
 
 add_library(gpu_cache gpu_cache.cpp)
 add_third_party_includes(gpu_cache)
-target_link_libraries(gpu_cache PUBLIC xxHash::xxhash ${TORCH_LIBRARIES} cuda_stream_manager)
 
+# gpu_cache
+if(KTRANSFORMERS_USE_NPU)
+    target_include_directories(gpu_cache PUBLIC ${TORCH_NPU_PATH}/include)
+    find_package(Python COMPONENTS Interpreter Development REQUIRED)
+    message("python include location: " ${Python_INCLUDE_DIRS} " lib location: " ${Python_LIBRARIES})
+    target_link_libraries(gpu_cache PUBLIC xxHash::xxhash "${TORCH_LIBRARIES}"  "${TORCH_PYTHON_LIBRARY}" "${PTA_LIBRARY}" ${Python_LIBRARIES} cuda_stream_manager)
+else()
+    target_link_libraries(gpu_cache PUBLIC xxHash::xxhash ${TORCH_LIBRARIES} cuda_stream_manager)
+endif()
+
+# kvc2
 add_library(kvc2 prefix.cpp)
 target_include_directories(kvc2 PRIVATE ${THIRD_PARTY_DIR}/nlohmann/single_include)
 add_third_party_includes(kvc2)
-target_link_libraries(kvc2 PUBLIC TBB::tbb xxHash::xxhash cache_entry cuda_stream_manager page_aligned_memory_pool ${TORCH_LIBRARIES} prometheus-cpp::pull kvc2_metrics)
+if(KTRANSFORMERS_USE_NPU)
+    target_link_libraries(kvc2 PUBLIC TBB::tbb xxHash::xxhash cache_entry cuda_stream_manager page_aligned_memory_pool prometheus-cpp::pull kvc2_metrics)
+else()
+    target_link_libraries(kvc2 PUBLIC TBB::tbb xxHash::xxhash cache_entry cuda_stream_manager page_aligned_memory_pool ${TORCH_LIBRARIES} prometheus-cpp::pull kvc2_metrics)
+endif()
 
 message(STATUS "CMAKE_SOURCE_DIR: " ${CMAKE_SOURCE_DIR})
+
+# async_store
 add_library(async_store async_store.cpp)
 target_include_directories(async_store PRIVATE ${THIRD_PARTY_DIR}/nlohmann/single_include)
-target_include_directories(async_store PRIVATE ${THIRD_PARTY_DIR}/PhotonLibOS/include)
 target_include_directories(async_store PRIVATE ${THIRD_PARTY_DIR}/spdlog/include)
+target_include_directories(async_store PRIVATE ${THIRD_PARTY_DIR}/PhotonLibOS/include)
 target_link_libraries(async_store PUBLIC photon_static pthread)
 
-
-
+# cuda_stream_manager
 add_library(cuda_stream_manager cuda_stream_manager.cpp)
 target_include_directories(cuda_stream_manager PUBLIC ${THIRD_PARTY_DIR}/nlohmann/single_include)
 target_include_directories(cuda_stream_manager PUBLIC ${THIRD_PARTY_DIR}/spdlog/include)
-target_include_directories(cuda_stream_manager  PUBLIC ${CUDAToolkit_INCLUDE_DIRS})
-target_link_libraries(cuda_stream_manager PUBLIC CUDA::cudart)
+if(KTRANSFORMERS_USE_NPU)
+    set(ASCEND_HOME_PATH "$ENV{ASCEND_HOME_PATH}")
+    target_include_directories(cuda_stream_manager PUBLIC ${ASCEND_HOME_PATH}/include)
+    target_link_directories(cuda_stream_manager PUBLIC ${ASCEND_HOME_PATH}/lib64)
+    target_link_libraries(cuda_stream_manager PUBLIC ascendcl)
+else()
+    target_include_directories(cuda_stream_manager PUBLIC ${CUDAToolkit_INCLUDE_DIRS})
+    target_link_libraries(cuda_stream_manager PUBLIC CUDA::cudart)
+endif()