feat: add musa device compatibility.

xllm/core/framework/parallel_state/collective_communicator.cpp

@@ -23,12 +23,12 @@ limitations under the License.
 #include "xllm_atb_layers/core/include/atb_speed/utils/singleton.h"
 #elif defined(USE_MLU)
 #include "mlu_process_group.h"
+#elif defined(USE_MUSA)
+#include "musa_process_group.h"
 #elif defined(USE_CUDA) || defined(USE_DCU)
 #include "cuda_process_group.h"
 #elif defined(USE_ILU)
 #include "ilu_process_group.h"
-#elif defined(USE_MUSA)
-#include "musa_process_group.h"
 #endif
 #include "core/framework/config/eplb_config.h"
 #include "core/framework/config/kernel_config.h"

xllm/core/framework/parallel_state/dit_collective_communicator.cpp

@@ -21,12 +21,12 @@ limitations under the License.
 #include "npu_process_group.h"
 #elif defined(USE_MLU)
 #include "mlu_process_group.h"
+#elif defined(USE_MUSA)
+#include "musa_process_group.h"
 #elif defined(USE_CUDA) || defined(USE_DCU)
 #include "cuda_process_group.h"
 #elif defined(USE_ILU)
 #include "ilu_process_group.h"
-#elif defined(USE_MUSA)
-#include "musa_process_group.h"
 #endif
 #include "common/global_flags.h"
 #include "parallel_args.h"

xllm/core/framework/parallel_state/process_group.cpp

@@ -19,10 +19,10 @@ limitations under the License.
 #include "npu_process_group.h"
 #elif defined(USE_MLU)
 #include "mlu_process_group.h"
-#elif defined(USE_CUDA) || defined(USE_DCU)
-#include "cuda_process_group.h"
 #elif defined(USE_MUSA)
 #include "musa_process_group.h"
+#elif defined(USE_CUDA) || defined(USE_DCU)
+#include "cuda_process_group.h"
 #elif defined(USE_ILU)
 #include "ilu_process_group.h"
 #endif

xllm/core/platform/CMakeLists.txt

@@ -34,7 +34,7 @@ cc_library(
     $<$<BOOL:${USE_DCU}>:dcu/dcu_layer_synchronizer.h>
     $<$<BOOL:${USE_DCU}>:dcu/dcu_tensor_alloc.h>
     $<$<BOOL:${USE_CUDA}>:cuda/cuda_utils.h>
-    $<$<BOOL:${USE_CUDA}>:cuda_profiler.h>
+    $<$<AND:$<BOOL:${USE_CUDA}>,$<NOT:$<BOOL:${USE_MUSA}>>>:cuda_profiler.h>
     $<$<BOOL:${USE_CUDA}>:torch_profiler.h>
     $<$<OR:$<BOOL:${USE_CUDA}>,$<BOOL:${USE_MLU}>,$<BOOL:${USE_DCU}>>:numa_utils.h>
   SRCS
@@ -49,7 +49,7 @@ cc_library(
     layer_synchronizer.cpp
     $<$<BOOL:${USE_MLU}>:mlu/mlu_layer_synchronizer.cpp>
     $<$<BOOL:${USE_MLU}>:mlu/mlu_tensor_alloc.cpp>
-    $<$<BOOL:${USE_CUDA}>:cuda_profiler.cpp>
+    $<$<AND:$<BOOL:${USE_CUDA}>,$<NOT:$<BOOL:${USE_MUSA}>>>:cuda_profiler.cpp>
     $<$<BOOL:${USE_CUDA}>:torch_profiler.cpp>
     $<$<BOOL:${USE_DCU}>:dcu/dcu_layer_synchronizer.cpp>
     $<$<BOOL:${USE_DCU}>:dcu/dcu_tensor_alloc.cpp>
@@ -66,8 +66,8 @@ cc_library(
     $<$<BOOL:${USE_MLU}>:torch_mlu>
     $<$<BOOL:${USE_MLU}>:cnrt>
     $<$<BOOL:${USE_MLU}>:cndrv>
-    $<$<OR:$<BOOL:${USE_CUDA}>,$<BOOL:${USE_ILU}>>:cuda>
-    $<$<OR:$<BOOL:${USE_CUDA}>,$<BOOL:${USE_ILU}>>:cudart>
+    $<$<AND:$<OR:$<BOOL:${USE_CUDA}>,$<BOOL:${USE_ILU}>>,$<NOT:$<BOOL:${USE_MUSA}>>>:cuda>
+    $<$<AND:$<OR:$<BOOL:${USE_CUDA}>,$<BOOL:${USE_ILU}>>,$<NOT:$<BOOL:${USE_MUSA}>>>:cudart>
     $<$<OR:$<BOOL:${USE_CUDA}>,$<BOOL:${USE_MLU}>,$<BOOL:${USE_ILU}>,$<BOOL:${USE_DCU}>>:numa>
     $<$<BOOL:${USE_MUSA}>:musa>
     $<$<BOOL:${USE_MUSA}>:musart>

xllm/core/platform/cuda/cuda_utils.h

@@ -15,47 +15,84 @@ limitations under the License.
 
 #pragma once
 
+#if defined(USE_MUSA)
+#include <musa_runtime.h>
+#else
 #include <cuda.h>
 #include <cuda_runtime.h>
+#endif
 #include <glog/logging.h>
+#if !defined(USE_MUSA)
 #include <nvtx3/nvToolsExt.h>
+#endif
 #include <torch/torch.h>
 
 #include <utility>
 
 namespace xllm::cuda {
 
+#if !defined(USE_MUSA)
 class NvtxRange {
  public:
   NvtxRange(const std::string& name) { nvtxRangePush(name.c_str()); }
 
   ~NvtxRange() { nvtxRangePop(); }
 };
+#else
+class NvtxRange {
+ public:
+  explicit NvtxRange(const std::string& name) { (void)name; }
+  ~NvtxRange() = default;
+};
+#endif
 
 inline int32_t get_device_sm_count(int32_t device_id) {
+#if defined(USE_MUSA)
+  musaDeviceProp prop;
+  musaError_t err = musaGetDeviceProperties(&prop, device_id);
+  if (err != musaSuccess) {
+    LOG(FATAL) << "Failed to get device properties for device " << device_id;
+  }
+#else
   cudaDeviceProp prop;
   cudaError_t err = cudaGetDeviceProperties(&prop, device_id);
   if (err != cudaSuccess) {
     LOG(FATAL) << "Failed to get device properties for device " << device_id;
   }
+#endif
   return prop.multiProcessorCount;
 }
 
 inline std::pair<int32_t, int32_t> get_compute_capability(int32_t device_id) {
+#if defined(USE_MUSA)
+  musaDeviceProp prop;
+  musaError_t err = musaGetDeviceProperties(&prop, device_id);
+  if (err != musaSuccess) {
+    LOG(FATAL) << "Failed to get compute capability for device " << device_id;
+  }
+#else
   cudaDeviceProp prop;
   cudaError_t err = cudaGetDeviceProperties(&prop, device_id);
   if (err != cudaSuccess) {
     LOG(FATAL) << "Failed to get compute capability for device " << device_id;
   }
+#endif
   return std::make_pair(prop.major, prop.minor);
 }
 
 inline int32_t get_cuda_version() {
   int32_t version;
+#if defined(USE_MUSA)
+  musaError_t err = musaRuntimeGetVersion(&version);
+  if (err != musaSuccess) {
+    LOG(FATAL) << "Failed to get MUSA runtime version!";
+  }
+#else
   cudaError_t err = cudaRuntimeGetVersion(&version);
   if (err != cudaSuccess) {
     LOG(FATAL) << "Failed to get CUDA version!";
   }
+#endif
   return version;
 }
 
@@ -93,4 +130,4 @@ inline bool support_sm120a(int32_t device_id) {
   return (major == 12) && (minor == 0) && cuda_version_at_least(12, 8);
 }
 
-}  // namespace xllm::cuda
+}  // namespace xllm::cuda

xllm/core/platform/device.cpp

@@ -25,13 +25,13 @@ limitations under the License.
 #include <framework/core/device.h>
 #include <framework/core/device_utils.h>
 #include <framework/generator/generator_impl.h>
+#elif defined(USE_MUSA)
+#include <c10/musa/MUSAGuard.h>
+#include <musa.h>
 #elif defined(USE_CUDA) || defined(USE_ILU)
 #include <c10/cuda/CUDACachingAllocator.h>
 #include <c10/cuda/CUDAStream.h>
 #include <cuda.h>
-#elif defined(USE_MUSA)
-#include <c10/musa/MUSAGuard.h>
-#include <musa.h>
 #elif defined(USE_DCU)
 #include <c10/hip/HIPCachingAllocator.h>
 #include <c10/hip/HIPStream.h>
@@ -41,14 +41,14 @@ limitations under the License.
 namespace xllm {
 
 Device::Device(const torch::Device& device) : device_(device) {
-#if defined(USE_CUDA)
+#if defined(USE_MUSA) || defined(USE_CUDA)
   Platform::init_capabilities(device.index());
 #endif
 }
 
 Device::Device(const int32_t device_index)
     : device_(torch::Device(Platform::type_torch(), device_index)) {
-#if defined(USE_CUDA)
+#if defined(USE_MUSA) || defined(USE_CUDA)
   Platform::init_capabilities(device_index);
 #endif
 }
@@ -60,10 +60,10 @@ void Device::set_device() const {
   c10_npu::set_device(index());
 #elif defined(USE_MLU)
   torch_mlu::setDevice(index());
-#elif defined(USE_CUDA) || defined(USE_ILU)
-  c10::cuda::set_device(index());
 #elif defined(USE_MUSA)
   c10::musa::set_device(index());
+#elif defined(USE_CUDA) || defined(USE_ILU)
+  c10::cuda::set_device(index());
 #elif defined(USE_DCU)
   c10::hip::set_device(index());
 #endif
@@ -80,10 +80,10 @@ void Device::set_seed(uint64_t seed) const {
     std::lock_guard<std::mutex> lock(gen.mutex());
     gen.set_current_seed(seed);
   }
-#elif defined(USE_CUDA) || defined(USE_DCU)
-  torch::cuda::manual_seed(seed);
 #elif defined(USE_MUSA)
   torch::manual_seed(seed);
+#elif defined(USE_CUDA) || defined(USE_DCU)
+  torch::cuda::manual_seed(seed);
 #endif
 }
 
@@ -107,10 +107,10 @@ Device::DeviceMem Device::get_device_mem() const {
   aclrtGetMemInfo(ACL_HBM_MEM, &free_memory, &total_memory);
 #elif defined(USE_MLU)
   cnrtMemGetInfo(&free_memory, &total_memory);
-#elif defined(USE_CUDA) || defined(USE_ILU)
-  cudaMemGetInfo(&free_memory, &total_memory);
 #elif defined(USE_MUSA)
   musaMemGetInfo(&free_memory, &total_memory);
+#elif defined(USE_CUDA) || defined(USE_ILU)
+  cudaMemGetInfo(&free_memory, &total_memory);
 #elif defined(USE_DCU)
   hipMemGetInfo(&free_memory, &total_memory);
 #endif
@@ -128,10 +128,10 @@ void Device::empty_cache(int32_t device_index) {
   c10_npu::NPUCachingAllocator::FreeDeviceCachedMemory(device_index);
 #elif defined(USE_MLU)
   torch_mlu::MLUCachingAllocator::emptyCache();
-#elif defined(USE_CUDA) || defined(USE_ILU)
-  c10::cuda::CUDACachingAllocator::emptyCache();
 #elif defined(USE_MUSA)
   c10::musa::MUSACachingAllocator::emptyCache();
+#elif defined(USE_CUDA) || defined(USE_ILU)
+  c10::cuda::CUDACachingAllocator::emptyCache();
 #elif defined(USE_DCU)
   c10::hip::HIPCachingAllocator::emptyCache();
 #endif
@@ -144,10 +144,10 @@ int Device::synchronize_default_stream() {
   return aclrtSynchronizeStream(c10_npu::getCurrentNPUStream(index()).stream());
 #elif defined(USE_MLU)
   torch_mlu::getCurrentMLUStream(index()).synchronize();
-#elif defined(USE_CUDA) || defined(USE_ILU)
-  c10::cuda::getCurrentCUDAStream().synchronize();
 #elif defined(USE_MUSA)
   c10::musa::getCurrentMUSAStream().synchronize();
+#elif defined(USE_CUDA) || defined(USE_ILU)
+  c10::cuda::getCurrentCUDAStream().synchronize();
 #elif defined(USE_DCU)
   c10::hip::getCurrentHIPStream().synchronize();
 #endif
@@ -163,10 +163,10 @@ std::unique_ptr<Stream> Device::current_stream() const {
   auto current_s = c10_npu::getCurrentNPUStream(index());
 #elif defined(USE_MLU)
   auto current_s = torch_mlu::getCurrentMLUStream(index());
-#elif defined(USE_CUDA) || defined(USE_ILU)
-  auto current_s = c10::cuda::getCurrentCUDAStream(index());
 #elif defined(USE_MUSA)
   auto current_s = c10::musa::getCurrentMUSAStream(index());
+#elif defined(USE_CUDA) || defined(USE_ILU)
+  auto current_s = c10::cuda::getCurrentCUDAStream(index());
 #elif defined(USE_DCU)
   auto current_s = c10::hip::getCurrentHIPStream(index());
 #endif

xllm/core/platform/numa_utils.cpp

@@ -15,7 +15,9 @@ limitations under the License.
 
 #include "numa_utils.h"
 
-#if defined(USE_CUDA)
+#if defined(USE_MUSA)
+#include <musa_runtime.h>
+#elif defined(USE_CUDA)
 #include <cuda_runtime.h>
 #elif defined(USE_MLU)
 #include <cnrt.h>
@@ -181,7 +183,18 @@ int32_t get_device_numa_node(int32_t device_index) {
     return -1;
   }
 
-#if defined(USE_CUDA)
+#if defined(USE_MUSA)
+  char pci_bus_id[32] = {0};
+  musaError_t ret =
+      musaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), device_index);
+  if (ret == musaSuccess) {
+    return get_numa_node_from_sysfs("MUSA", device_index, pci_bus_id);
+  }
+
+  LOG(WARNING) << "Failed to query PCI bus ID for MUSA device " << device_index
+               << " via musaDeviceGetPCIBusId: " << musaGetErrorString(ret)
+               << ", skipping NUMA binding";
+#elif defined(USE_CUDA)
   char pci_bus_id[32] = {0};
   cudaError_t ret =
       cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), device_index);

xllm/core/platform/platform.cpp

@@ -21,12 +21,14 @@ limitations under the License.
 #include <torch_npu/csrc/core/npu/NPUCachingAllocator.h>
 #elif defined(USE_MLU)
 #include <framework/core/device.h>
+#elif defined(USE_MUSA)
+#include <c10/musa/MUSAGuard.h>
+
+#include "core/platform/cuda/cuda_utils.h"
 #elif defined(USE_CUDA) || defined(USE_ILU)
 #include <c10/cuda/CUDACachingAllocator.h>
 
 #include "core/platform/cuda/cuda_utils.h"
-#elif defined(USE_MUSA)
-#include <c10/musa/MUSAGuard.h>
 #elif defined(USE_DCU)
 #include <c10/hip/HIPCachingAllocator.h>
 #endif
@@ -49,12 +51,12 @@ std::string Platform::type_str() {
   return "npu";
 #elif defined(USE_MLU)
   return "mlu";
+#elif defined(USE_MUSA)
+  return "musa";
 #elif defined(USE_CUDA)
   return "cuda";
 #elif defined(USE_ILU)
   return "ilu";
-#elif defined(USE_MUSA)
-  return "musa";
 #elif defined(USE_DCU)
   return "dcu";
 #endif
@@ -63,10 +65,10 @@ std::string Platform::type_str() {
 torch::DeviceType Platform::type_torch() {
 #if defined(USE_NPU) || defined(USE_MLU)
   return torch::kPrivateUse1;
-#elif defined(USE_CUDA) || defined(USE_ILU) || defined(USE_DCU)
-  return torch::kCUDA;
 #elif defined(USE_MUSA)
   return torch::kMUSA;
+#elif defined(USE_CUDA) || defined(USE_ILU) || defined(USE_DCU)
+  return torch::kCUDA;
 #endif
 }
 
@@ -75,10 +77,10 @@ int32_t Platform::device_count() {
   return static_cast<int32_t>(c10_npu::device_count());
 #elif defined(USE_MLU)
   return static_cast<int32_t>(torch_mlu::device_count());
-#elif defined(USE_CUDA) || defined(USE_ILU)
-  return static_cast<int32_t>(c10::cuda::device_count());
 #elif defined(USE_MUSA)
   return static_cast<int32_t>(c10::musa::device_count());
+#elif defined(USE_CUDA) || defined(USE_ILU)
+  return static_cast<int32_t>(c10::cuda::device_count());
 #elif defined(USE_DCU)
   return static_cast<int32_t>(c10::hip::device_count());
 #endif
@@ -89,10 +91,10 @@ int32_t Platform::current_device() {
   return static_cast<int32_t>(c10_npu::current_device());
 #elif defined(USE_MLU)
   return static_cast<int32_t>(torch_mlu::current_device());
-#elif defined(USE_CUDA) || defined(USE_ILU)
-  return static_cast<int32_t>(c10::cuda::current_device());
 #elif defined(USE_MUSA)
   return static_cast<int32_t>(c10::musa::current_device());
+#elif defined(USE_CUDA) || defined(USE_ILU)
+  return static_cast<int32_t>(c10::cuda::current_device());
 #elif defined(USE_DCU)
   return static_cast<int32_t>(c10::hip::current_device());
 #endif
@@ -111,7 +113,7 @@ bool Platform::is_support_sm100f() { return support_sm100f_; }
 bool Platform::is_support_sm120a() { return support_sm120a_; }
 
 void Platform::init_capabilities(int32_t device_index) {
-#if defined(USE_CUDA)
+#if defined(USE_MUSA) || defined(USE_CUDA)
   std::call_once(g_init_flag, [device_index]() {
     sm_count_ = cuda::get_device_sm_count(device_index);
     enable_pdl_ = cuda::support_pdl(device_index);