feat: support qwen3-32b and qwen3-30b-a3b on mlu device. (#227)

Co-authored-by: guoxueting <[email protected]>

Author: a120092009

CMakeLists.txt

@@ -332,7 +332,24 @@ endif()
 
 if(USE_MLU)
   add_definitions(-DUSE_MLU)
-# TODO(mlu): set mlu environment variables
+  set(CMAKE_VERBOSE_MAKEFILE ON)
+  include_directories(
+      $ENV{PYTHON_INCLUDE_PATH}
+      $ENV{PYTORCH_INSTALL_PATH}/include
+      $ENV{PYTORCH_INSTALL_PATH}/include/torch/csrc/api/include
+      $ENV{PYTORCH_MLU_INSTALL_PATH}
+      $ENV{PYTORCH_MLU_INSTALL_PATH}/../
+      $ENV{PYTORCH_MLU_INSTALL_PATH}/csrc
+      $ENV{NEUWARE_HOME}/include
+  )
+
+  link_directories(
+    $ENV{PYTHON_LIB_PATH}
+    $ENV{PYTORCH_INSTALL_PATH}/lib
+    $ENV{PYTORCH_MLU_INSTALL_PATH}/csrc/lib
+    $ENV{PYTORCH_MLU_INSTALL_PATH}
+    $ENV{NEUWARE_HOME}/lib64
+  )
 endif()
 
 # check if USE_CXX11_ABI is set correctly

setup.py

@@ -106,6 +106,14 @@ def get_torch_root_path():
     except ImportError:
         return None
 
+def get_torch_mlu_root_path():
+    try:
+        import torch_mlu
+        import os
+        return os.path.dirname(os.path.abspath(torch_mlu.__file__))
+    except ImportError:
+        return None
+
 
 def set_npu_envs():
     PYTORCH_NPU_INSTALL_PATH = os.getenv("PYTORCH_NPU_INSTALL_PATH")
@@ -196,11 +204,13 @@ def set_npu_envs():
     os.environ["LCCL_DETERMINISTIC"] = "0"
     os.environ["LCCL_PARALLEL"] = "0"
 
-# TODO(mlu): set mlu environment variables
+
 def set_mlu_envs():
     os.environ["PYTHON_INCLUDE_PATH"] = get_python_include_path()
     os.environ["PYTHON_LIB_PATH"] =  get_torch_root_path()
     os.environ["LIBTORCH_ROOT"] = get_torch_root_path()
+    os.environ["PYTORCH_INSTALL_PATH"] = get_torch_root_path()
+    os.environ["PYTORCH_MLU_INSTALL_PATH"] = get_torch_mlu_root_path()
 
 class CMakeExtension(Extension):
     def __init__(self, name: str, path: str, sourcedir: str = "") -> None:

xllm/core/common/CMakeLists.txt

@@ -34,6 +34,7 @@ cc_library(
     Boost::serialization
     cpprest
     etcd-cpp-api
+    $<$<BOOL:${USE_MLU}>:torch_mlu>
 )
 
 cc_library(

xllm/core/distributed_runtime/CMakeLists.txt

@@ -31,6 +31,7 @@ cc_library(
     proto::xllm_proto
     absl::flat_hash_set
     :parallel_state
+    :collective_service
 )
 
 cc_library(

xllm/core/distributed_runtime/collective_service.cpp

@@ -44,7 +44,6 @@ void CollectiveService::Sync(::google::protobuf::RpcController* controller,
                              const proto::AddressInfo* request,
                              proto::CommUniqueIdList* response,
                              ::google::protobuf::Closure* done) {
-#if defined(USE_NPU)
   brpc::ClosureGuard done_guard(done);
 
   std::string address = request->address();
@@ -53,10 +52,9 @@ void CollectiveService::Sync(::google::protobuf::RpcController* controller,
     std::lock_guard<std::mutex> lock(mutex_);
     addrs_map_[global_rank] = address;
   }
-
+#if defined(USE_NPU)
   to_proto_list(root_infos_, response);
 #endif
-  return;
 }
 
 std::unordered_map<int32_t, std::string> CollectiveService::wait() {

xllm/core/distributed_runtime/worker_server.cpp

@@ -91,6 +91,9 @@ void WorkerServer::create_server(const runtime::Options& options,
 
   CollectiveCommunicator comm(worker_global_rank, world_size, dp_size, ep_size);
   const ParallelArgs* parallel_args = comm.parallel_args();
+#if defined(USE_MLU)
+  comm.create_process_groups_cncl(master_node_addr, device);
+#endif
 
   WorkerType worker_type =
       (options.task_type() == "generate") ? WorkerType::LLM : WorkerType::ELM;

xllm/core/framework/CMakeLists.txt

@@ -22,7 +22,7 @@ add_subdirectory(xtensor)
 add_subdirectory(dit_cache)
 
 cc_library(
-  NAME 
+  NAME
     parallel_state
   HDRS
     mapping_npu.h
@@ -33,30 +33,33 @@ cc_library(
   DEPS
     :common
     torch
-    hccl
+    $<$<BOOL:${USE_MLU}>:torch_mlu>
+    $<$<BOOL:${USE_NPU}>:hccl>
     glog::glog
 )
 
 
-cc_test(
-  NAME 
-    mapping_npu_test
-  SRCS
-    mapping_npu_test.cpp
-  DEPS
-    parallel_state
-    absl::synchronization
-    absl::time
-    GTest::gtest_main
-    xllm_kernels
-    ascendcl
-    atb
-    c_sec
-    spdlog::spdlog
-)
+if(USE_NPU)
+  cc_test(
+    NAME
+      mapping_npu_test
+    SRCS
+      mapping_npu_test.cpp
+    DEPS
+      parallel_state
+      absl::synchronization
+      absl::time
+      GTest::gtest_main
+      xllm_kernels
+      ascendcl
+      atb
+      c_sec
+      spdlog::spdlog
+  )
+endif()
 
 cc_library(
-  NAME 
+  NAME
     model_loader
   HDRS
     hf_model_loader.h
@@ -78,7 +81,7 @@ cc_library(
 )
 
 cc_library(
-  NAME 
+  NAME
     model_context
   HDRS
     model_context.h

xllm/core/framework/block/CMakeLists.txt

@@ -28,24 +28,26 @@ cc_library(
     torch
 )
 
-set(TEST_SRCS
-  block_manager_test.cpp
-)
+if(USE_NPU)
+  set(TEST_SRCS
+    block_manager_test.cpp
+  )
 
-cc_test(
-  NAME
-    block_test
-  SRCS
-    ${TEST_SRCS}
-  DEPS
-    :block
-    :flags
-    :kv_cache
-    :prefix_cache
-    absl::random_random
-    Boost::serialization
-    GTest::gtest_main
-)
+  cc_test(
+    NAME
+      block_test
+    SRCS
+      ${TEST_SRCS}
+    DEPS
+      :block
+      :flags
+      :kv_cache
+      :prefix_cache
+      absl::random_random
+      Boost::serialization
+      GTest::gtest_main
+  )
 
-target_link_libraries(block_test PRIVATE brpc OpenSSL::SSL OpenSSL::Crypto ascendcl Folly::folly)
-add_dependencies(block_test brpc-static)
+  target_link_libraries(block_test PRIVATE brpc OpenSSL::SSL OpenSSL::Crypto ascendcl Folly::folly)
+  add_dependencies(block_test brpc-static)
+endif()

xllm/core/framework/dit_model_context.cpp

@@ -72,6 +72,7 @@ const QuantArgs& DiTModelContext::get_quant_args(
   }
 }
 
+#if defined(USE_NPU)
 ModelContext DiTModelContext::get_model_context(
     const std::string& component) const {
   return ModelContext(parallel_args_,
@@ -80,5 +81,6 @@ ModelContext DiTModelContext::get_model_context(
                       tensor_options_,
                       context_);
 }
+#endif
 
 }  // namespace xllm

xllm/core/framework/dit_model_context.h

@@ -42,7 +42,9 @@ class DiTModelContext {
 
   const QuantArgs& get_quant_args(const std::string& component) const;
 
+#if defined(USE_NPU)
   ModelContext get_model_context(const std::string& component) const;
+#endif
 
   const ParallelArgs& get_parallel_args() const { return parallel_args_; }
 
@@ -52,7 +54,9 @@ class DiTModelContext {
 
   const std::string& model_type() const { return model_type_; }
 
+#if defined(USE_NPU)
   const atb::Context* get_atb_context() const { return context_; }
+#endif
 
  private:
   std::unordered_map<std::string, ModelArgs> model_args_;