feat: add additional musa build supports.

CMakeLists.txt

@@ -256,7 +256,30 @@ set(CPPREST_EXCLUDE_WEBSOCKETS ON CACHE BOOL "Exclude websockets functionality."
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-format-truncation")
 
 list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_LIST_DIR}/cmake)
-if(USE_CUDA)
+set(_XLLM_MUSIFY_COPTS "")
+if(USE_MUSA)
+  if(NOT USE_CUDA)
+    message(FATAL_ERROR "USE_MUSA requires USE_CUDA=ON")
+  endif()
+
+  if(DEFINED ENV{MUSAMAPPING_PATH} AND NOT "$ENV{MUSAMAPPING_PATH}" STREQUAL "")
+    set(_XLLM_MUSAMAPPING_PATH "$ENV{MUSAMAPPING_PATH}")
+  elseif(DEFINED ENV{MUSA_HOME} AND NOT "$ENV{MUSA_HOME}" STREQUAL "")
+    set(_XLLM_MUSAMAPPING_PATH "$ENV{MUSA_HOME}/tools/musamapping")
+  else()
+    set(_XLLM_MUSAMAPPING_PATH "/usr/local/musa/tools/musamapping")
+  endif()
+  set(_XLLM_MCC_WRAPPER "${_XLLM_MUSAMAPPING_PATH}/mcc_wrapper")
+  set(_XLLM_MUSIFY_COPTS
+    "-fplugin=${_XLLM_MUSAMAPPING_PATH}/libMusaMapping.so"
+    "-x" "musa"
+    "-I."
+  )
+  set(CMAKE_CXX_COMPILER "${_XLLM_MCC_WRAPPER}" CACHE FILEPATH
+      "C++ compiler for the MUSA compatibility path" FORCE)
+  project("xllm" LANGUAGES C CXX CUDA)
+  add_compile_options(-Wno-c++11-narrowing)
+elseif(USE_CUDA)
   project("xllm" LANGUAGES C CXX CUDA)
   find_package(CUDAToolkit REQUIRED)
 elseif(USE_DCU)
@@ -268,11 +291,6 @@ elseif(USE_DCU)
     list(APPEND CMAKE_PREFIX_PATH "/opt/dtk")
     set(DCU_ROOT "/opt/dtk")
   endif()
-elseif(USE_MUSA)
-  project("xllm" LANGUAGES C CXX MUSA)
-  add_compile_options(
-        -Wno-c++11-narrowing
-      )
 else()
   project("xllm" LANGUAGES C CXX)
 endif()
@@ -515,67 +533,99 @@ endif()
 
 if(USE_MUSA)
   add_definitions(-DUSE_MUSA)
+  add_definitions(-DUSE_CUDA)
   add_compile_definitions(TORCH_CUDA=1)
-
   if(NOT DEFINED MUSA_PATH)
-    set(MUSA_PATH /usr/local/musa)
+    if(DEFINED ENV{MUSA_HOME} AND NOT "$ENV{MUSA_HOME}" STREQUAL "")
+      set(MUSA_PATH "$ENV{MUSA_HOME}")
+    else()
+      set(MUSA_PATH /usr/local/musa)
+    endif()
   endif()
   list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake")
   find_package(MUSA REQUIRED)
-
-  if(NOT DEFINED ENV{MTT_OPLIB_PATH})
-    set(ENV{MTT_OPLIB_PATH} $ENV{MUSA_HOME}/tools/MTTOplib)
+  # MUSA builds use the FlashInfer/Mate TVM-FFI path; MTTOplib is not needed.
+  if(NOT DEFINED ENV{TORCH_MUSA_PYTHONPATH} OR
+     "$ENV{TORCH_MUSA_PYTHONPATH}" STREQUAL "")
+    message(FATAL_ERROR
+            "TORCH_MUSA_PYTHONPATH is required for USE_MUSA builds")
   endif()
 
-  message(STATUS  "using MTT Oplib at: $ENV{MTT_OPLIB_PATH}")
-  set(MTTOplib_DIR $ENV{MTT_OPLIB_PATH}/cmake)
-  find_package(MTTOplib REQUIRED)
-
   find_package(Python COMPONENTS Interpreter REQUIRED)
-  list(APPEND CMAKE_PREFIX_PATH $ENV{TORCH_MUSA_PYTHONPATH})
+  list(APPEND CMAKE_PREFIX_PATH "$ENV{TORCH_MUSA_PYTHONPATH}")
   find_package(TorchMusa REQUIRED CONFIG)
   list(POP_BACK CMAKE_PREFIX_PATH)
-
+  get_filename_component(
+    _TORCH_MUSA_ROOT
+    "$ENV{TORCH_MUSA_PYTHONPATH}/../.."
+    ABSOLUTE)
+
+  if(DEFINED ENV{MKLROOT} AND NOT "$ENV{MKLROOT}" STREQUAL "" AND
+     NOT DEFINED MKL_DIR)
+    set(MKL_DIR "$ENV{MKLROOT}/lib/cmake/mkl")
+  endif()
   find_package(MKL REQUIRED)
+
+  set(_MUSA_MKL_LIBRARY_DIRS ${MKL_LIB_PATH})
+  if(DEFINED ENV{MKLROOT} AND NOT "$ENV{MKLROOT}" STREQUAL "")
+    foreach(_mkl_library_dir
+            "$ENV{MKLROOT}/lib/intel64"
+            "$ENV{MKLROOT}/lib")
+      if(IS_DIRECTORY "${_mkl_library_dir}")
+        list(APPEND _MUSA_MKL_LIBRARY_DIRS "${_mkl_library_dir}")
+      endif()
+    endforeach()
+  endif()
+  list(REMOVE_DUPLICATES _MUSA_MKL_LIBRARY_DIRS)
   execute_process(
-    COMMAND tvm-ffi-config --includedir
-    OUTPUT_VARIABLE tvm_ffi_INCLUDE_DIR
-    OUTPUT_STRIP_TRAILING_WHITESPACE
-  )
-  execute_process(
-    COMMAND tvm-ffi-config --libdir
-    OUTPUT_VARIABLE tvm_ffi_LIB_DIR
+    COMMAND "${Python_EXECUTABLE}" -c
+            "from sysconfig import get_paths; print(get_paths()['platlib'])"
+    OUTPUT_VARIABLE _PYTHON_PLATLIB
     OUTPUT_STRIP_TRAILING_WHITESPACE
+    RESULT_VARIABLE _PYTHON_PLATLIB_RESULT
   )
-  string(REGEX MATCH "[^\n]*/include[^\n]*" tvm_ffi_INCLUDE_DIR "${tvm_ffi_INCLUDE_DIR}")
-  string(REGEX MATCH "[^\n]*/lib[^\n]*" tvm_ffi_LIB_DIR "${tvm_ffi_LIB_DIR}")
+  if(NOT _PYTHON_PLATLIB_RESULT EQUAL 0 OR "${_PYTHON_PLATLIB}" STREQUAL "")
+    message(FATAL_ERROR "Failed to obtain Python platlib from sysconfig.get_paths()")
+  endif()
+  set(_TVM_FFI_ROOT "${_PYTHON_PLATLIB}/tvm_ffi")
+  set(tvm_ffi_INCLUDE_DIR "${_TVM_FFI_ROOT}/include")
+  set(tvm_ffi_LIB_DIR "${_TVM_FFI_ROOT}/lib")
+  if(NOT IS_DIRECTORY "${tvm_ffi_INCLUDE_DIR}" OR
+     NOT IS_DIRECTORY "${tvm_ffi_LIB_DIR}")
+    message(FATAL_ERROR
+            "tvm_ffi is not installed under Python platlib: ${_PYTHON_PLATLIB}")
+  endif()
 
 
   include_directories(
       ${MUSA_INCLUDE_DIRS}
       ${TorchMusa_INCLUDE_DIRS}
       ${Python_SITELIB}/torch_musa_compiled/share/torch_musa_codegen
-
+      ${tvm_ffi_INCLUDE_DIR}
+      ${MUSA_PATH}/include
+      ${_TORCH_MUSA_ROOT}/share/torch_musa_codegen
+      ${_TORCH_MUSA_ROOT}/share/generated_cuda_compatible/include
+      ${_TORCH_MUSA_ROOT}/share/generated_cuda_compatible/include/torch/csrc/api/include
+      ${_TORCH_MUSA_ROOT}/share/generated_cuda_compatible/include/torch/csrc/musa
       $ENV{PYTHON_INCLUDE_PATH}
       $ENV{PYTORCH_INSTALL_PATH}/include
       $ENV{PYTORCH_INSTALL_PATH}/include/torch/csrc/api/include
-      ${MTTOplib_INCLUDE_DIRS}
-      ${tvm_ffi_INCLUDE_DIR}
   )
 
   link_directories(
     ${MUSA_LIB_PATH}
+    ${MUSA_PATH}/lib
+    ${_TORCH_MUSA_ROOT}/lib
+    ${_MUSA_MKL_LIBRARY_DIRS}
     $ENV{PYTHON_LIB_PATH}
     $ENV{PYTORCH_INSTALL_PATH}/lib
     ${TorchMusa_LIB_PATH}
-    ${MKL_LIB_PATH}
     ${tvm_ffi_LIB_DIR}
     ${TorchMusa_INSTALL_PREFIX}/lib
-    )
-  set(MUSA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} -O3)
-endif()
-
-if(USE_CUDA)
+  )
+  set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} -O3)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=parentheses")
+elseif(USE_CUDA)
   message(STATUS "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}")
   add_definitions(-DUSE_CUDA)
   add_compile_definitions(TORCH_CUDA=1)

scripts/build_support/env.py

@@ -74,15 +74,6 @@ def get_dcu_root_path() -> Optional[str]:
         return None
 
 
-def get_torch_musa_root_path() -> Optional[str]:
-    try:
-        import torch_musa
-        import os
-        return os.path.dirname(os.path.abspath(torch_musa.__file__))
-    except ImportError:
-        return None
-
-
 def _find_dcu_so(package: str, pattern: str) -> Optional[str]:
     try:
         import glob
@@ -279,13 +270,45 @@ def set_ilu_envs() -> None:
 
 
 def set_musa_envs() -> None:
+    """Configure MUSA through mcc_wrapper and the CUDA compatibility path."""
+    from sysconfig import get_paths
     set_common_envs()
-    os.environ["PYTORCH_MUSA_INSTALL_PATH"] = get_torch_musa_root_path() or ""
     import torch_musa
-    from torch_musa.utils.musa_extension import MUSA_HOME
-    os.environ["TORCH_MUSA_PYTHONPATH"] = torch_musa.core.cmake_prefix_path
-    os.environ["MUSA_TOOLKIT_ROOT_DIR"] = MUSA_HOME
-    os.environ["MKL_DIR"] = "/opt/intel/oneapi/mkl/lib/cmake/mkl"
-    os.environ["MKLROOT"] = "/opt/intel/oneapi/mkl"
-    os.environ["TorchMusa_DIR"] = torch_musa.core.cmake_prefix_path + "/TorchMusa"
-    os.environ["MUSAMAPPING_PATH"] = MUSA_HOME + "/tools/musamapping"
+    from torch_musa.utils.musa_extension import MUSA_HOME as _MUSA_HOME
+    musa_home = os.getenv("MUSA_HOME") or _MUSA_HOME or "/usr/local/musa"
+    os.environ["MUSA_HOME"] = musa_home
+    os.environ["CUDA_HOME"] = musa_home
+    os.environ["CUDAToolkit_ROOT"] = musa_home
+    os.environ["CUDA_TOOLKIT_ROOT_DIR"] = musa_home
+    os.environ["MUSAMAPPING_PATH"] = os.path.join(
+        musa_home, "tools", "musamapping"
+    )
+
+    cmake_prefix = torch_musa.core.cmake_prefix_path
+    os.environ["TORCH_MUSA_PYTHONPATH"] = cmake_prefix
+    os.environ["TorchMusa_DIR"] = os.path.join(cmake_prefix, "TorchMusa")
+
+    torch_musa_root = os.path.abspath(os.path.join(cmake_prefix, "../.."))
+    library_paths: list[str] = [
+        os.path.join(musa_home, "lib"),
+        os.path.join(torch_musa_root, "lib"),
+        os.path.join(get_torch_root_path() or "", "lib"),
+    ]
+    python_platlib = get_paths()["platlib"]
+    library_paths.append(os.path.join(python_platlib, "tvm_ffi", "lib"))
+
+    mkl_root = os.getenv("MKLROOT")
+    if mkl_root:
+        os.environ.setdefault(
+            "MKL_DIR", os.path.join(mkl_root, "lib", "cmake", "mkl")
+        )
+        library_paths.extend(
+            [
+                os.path.join(mkl_root, "lib", "intel64"),
+                os.path.join(mkl_root, "lib"),
+            ]
+        )
+
+    for path in library_paths:
+        if path and os.path.isdir(path):
+            prepend_path_env("LD_LIBRARY_PATH", path)

setup.py

@@ -335,6 +335,20 @@ def build_extension(self, ext: CMakeExtension) -> None:
         elif self.device == "mlu":
             cmake_args += ["-DUSE_MLU=ON"]
             set_mlu_envs()
+        elif self.device == "musa":
+            torch_cuda_architectures = os.getenv("TORCH_CUDA_ARCH_LIST")
+            if not torch_cuda_architectures:
+                torch_cuda_architectures = "9.0"
+            cmake_args += [
+                "-DUSE_MUSA=ON",
+                "-DUSE_CUDA=ON",
+                f"-DTORCH_CUDA_ARCH_LIST={torch_cuda_architectures}",
+                "-DCMAKE_CUDA_ARCHITECTURES=90",
+                "-DBUILD_TESTING=OFF",
+            ]
+            set_musa_envs()
+            global BUILD_TEST_FILE
+            BUILD_TEST_FILE = False
         elif self.device == "cuda":
             torch_cuda_architectures = os.getenv("TORCH_CUDA_ARCH_LIST")
             if not torch_cuda_architectures:
@@ -400,11 +414,6 @@ def build_extension(self, ext: CMakeExtension) -> None:
         elif self.device == "ilu":
             cmake_args += ["-DUSE_ILU=ON"]
             set_ilu_envs()
-        elif self.device == "musa":
-            cmake_args += ["-DUSE_MUSA=ON"]
-            set_musa_envs()
-            global BUILD_TEST_FILE
-            BUILD_TEST_FILE = False
         else:
             raise ValueError("Please set --device to npu, mlu, cuda, dcu, ilu, musa or maca.")
 

xllm/CMakeLists.txt

@@ -21,10 +21,19 @@ include_directories("${CMAKE_CURRENT_BINARY_DIR}")
 
 # Set warning-as-error for xllm code (third-party headers are marked as SYSTEM)
 if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang")
-  add_compile_options(
-    $<$<COMPILE_LANGUAGE:CXX>:-Werror>
-    $<$<COMPILE_LANGUAGE:C>:-Werror>
-  )
+  if(USE_MUSA)
+    add_compile_options(
+      $<$<COMPILE_LANGUAGE:CXX>:-Wno-defaulted-function-deleted>
+      $<$<COMPILE_LANGUAGE:CXX>:-Wno-sign-conversion>
+      $<$<COMPILE_LANGUAGE:CXX>:-Wno-inconsistent-missing-override>
+      $<$<COMPILE_LANGUAGE:CXX>:-Wno-conversion>
+    )
+  else()
+    add_compile_options(
+      $<$<COMPILE_LANGUAGE:CXX>:-Werror>
+      $<$<COMPILE_LANGUAGE:C>:-Werror>
+    )
+  endif()
 
   if(CMAKE_CXX_COMPILER_ID MATCHES "Clang")
     add_compile_options(
@@ -107,10 +116,16 @@ if(USE_MSPTI)
   list(APPEND COMMON_LIBS mspti)
 endif()
 target_link_libraries(xllm PUBLIC ${COMMON_LIBS})
-if (USE_MUSA)
+if(USE_MUSA)
   target_link_libraries(xllm PUBLIC atomic musa_python torch_cpu c10)
 endif()
 
+if(USE_MUSA)
+  target_link_options(xllm PRIVATE
+    "LINKER:--whole-archive,$<TARGET_FILE:runtime>,--no-whole-archive"
+  )
+endif()
+
 # install xllm
 install(TARGETS xllm RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
 

xllm/core/kernels/CMakeLists.txt

@@ -9,16 +9,8 @@ if(USE_MLU)
 endif()
 
 if(USE_MUSA)
-  set(CMAKE_CXX_COMPILER ${CMAKE_MUSA_COMPILER})
-  string(REPLACE "-Wno-format-truncation" "" CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}")
-  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=macro-redefined \
-  -Wno-error=defaulted-function-deleted -x musa --offload-arch=mp_31 \
-  -fplugin=$ENV{MUSAMAPPING_PATH}/libMusaMapping.so")
-  set(CMAKE_MUSA_FLAGS "${CMAKE_MUSA_FLAGS} -fplugin=$ENV{MUSAMAPPING_PATH}/libMusaMapping.so")
   add_subdirectory(musa)
-endif()
-
-if(USE_CUDA)
+elseif(USE_CUDA)
   add_subdirectory(cuda)
 endif()
 
@@ -39,13 +31,14 @@ cc_library(
     xllm_torch_ops.h
   SRCS
     ops_api.cpp
+  COPTS
+    ${_XLLM_MUSIFY_COPTS}
   DEPS
     torch
     $<$<BOOL:${USE_NPU}>:triton_adapter>
     $<$<BOOL:${USE_NPU}>:npu_kernels>
     $<$<BOOL:${USE_NPU}>:xllm_ops>
     $<$<BOOL:${USE_MLU}>:mlu_kernels>
-    $<$<BOOL:${USE_MUSA}>:musa_kernels>
     $<$<BOOL:${USE_CUDA}>:cuda_kernels>
     $<$<BOOL:${USE_ILU}>:ilu_kernels>
     $<$<BOOL:${USE_DCU}>:dcu_kernels>