Support Moore Threads GPU

Signed-off-by: Xiaodong Ye <[email protected]>

Author: yeahdongcn

CMakeLists.txt

@@ -136,3 +136,33 @@ if(CMAKE_HIP_COMPILER)
         endforeach()
     endif()
 endif()
+
+if(NOT EXISTS $ENV{MUSA_PATH})
+    if (NOT EXISTS /opt/musa)
+        set(MUSA_PATH /usr/local/musa)
+    else()
+        set(MUSA_PATH /opt/musa)
+    endif()
+else()
+    set(MUSA_PATH $ENV{MUSA_PATH})
+endif()
+list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake")
+find_package(MUSAToolkit)
+if(MUSAToolkit_FOUND)
+    if(NOT MUSA_ARCHITECTURES)
+        set(MUSA_ARCHITECTURES "21;22;31")
+    endif()
+
+    if(MUSA_ARCHITECTURES)
+        add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/ml/backend/ggml/ggml/src/ggml-musa)
+        set(OLLAMA_MUSA_INSTALL_DIR ${OLLAMA_INSTALL_DIR}/musa_v4)
+        install(TARGETS ggml-musa
+            RUNTIME_DEPENDENCIES
+                DIRECTORIES ${MUSAToolkit_BIN_DIR} ${MUSAToolkit_LIBRARY_DIR}
+                PRE_INCLUDE_REGEXES mublas musart musa
+                PRE_EXCLUDE_REGEXES ".*"
+            RUNTIME DESTINATION ${OLLAMA_MUSA_INSTALL_DIR} COMPONENT MUSA
+            LIBRARY DESTINATION ${OLLAMA_MUSA_INSTALL_DIR} COMPONENT MUSA
+        )
+    endif()
+endif()

CMakePresets.json

@@ -53,6 +53,17 @@
         "CMAKE_HIP_FLAGS": "-parallel-jobs=4",
         "AMDGPU_TARGETS": "gfx900;gfx940;gfx941;gfx942;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102;gfx1151;gfx1200;gfx1201;gfx906:xnack-;gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-"
       }
+    },
+    {
+      "name": "MUSA",
+      "inherits": [ "Default" ]
+    },
+    {
+      "name": "MUSA 4",
+      "inherits": [ "MUSA" ],
+      "cacheVariables": {
+        "MUSA_ARCHITECTURES": "21;22;31"
+      }
     }
   ],
   "buildPresets": [
@@ -95,6 +106,16 @@
       "name": "ROCm 6",
       "inherits": [ "ROCm" ],
       "configurePreset": "ROCm 6"
+    },
+    {
+      "name": "MUSA",
+      "configurePreset": "MUSA",
+      "targets": [ "ggml-musa" ]
+    },
+    {
+      "name": "MUSA 4",
+      "inherits": [ "MUSA" ],
+      "configurePreset": "MUSA 4"
     }
   ]
 }

Dockerfile

@@ -6,6 +6,7 @@ ARG ROCMVERSION=6.3.3
 ARG JETPACK5VERSION=r35.4.1
 ARG JETPACK6VERSION=r36.4.0
 ARG CMAKEVERSION=3.31.2
+ARG MUSAVERSION=rc4.0.1
 
 # We require gcc v10 minimum.  v10.3 has regressions, so the rockylinux 8.5 AppStream has the latest compatible version
 FROM --platform=linux/amd64 rocm/dev-almalinux-8:${ROCMVERSION}-complete AS base-amd64
@@ -100,21 +101,49 @@ COPY --from=jetpack-6 dist/lib/ollama /lib/ollama/cuda_jetpack6
 FROM scratch AS rocm
 COPY --from=rocm-6 dist/lib/ollama /lib/ollama
 
+# Moore Threads (MUSA) build stages
+FROM --platform=linux/amd64 mthreads/musa:${MUSAVERSION}-mudnn-devel-ubuntu22.04 AS musa-4
+RUN apt-get update \
+    && apt-get install -y curl \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+ARG CMAKEVERSION
+RUN curl -fsSL https://github.com/Kitware/CMake/releases/download/v${CMAKEVERSION}/cmake-${CMAKEVERSION}-linux-$(uname -m).tar.gz | tar xz -C /usr/local --strip-components 1
+COPY CMakeLists.txt CMakePresets.json .
+COPY ml/backend/ggml/ggml ml/backend/ggml/ggml
+ENV LDFLAGS=-s
+RUN --mount=type=cache,target=/root/.ccache \
+    cmake --preset 'MUSA 4' \
+        && cmake --build --parallel --preset 'MUSA 4' \
+        && cmake --install build --component MUSA --strip --parallel 8
+
+FROM scratch AS musa
+COPY --from=musa-4 dist/lib/ollama/musa_v4 /lib/ollama/musa_v4
+
 FROM ${FLAVOR} AS archive
 COPY --from=cpu dist/lib/ollama /lib/ollama
 COPY --from=build /bin/ollama /bin/ollama
 
-FROM ubuntu:24.04
+FROM ubuntu:22.04
 RUN apt-get update \
     && apt-get install -y ca-certificates \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
+ARG FLAVOR
+RUN if [ "$FLAVOR" = "musa" ]; then \
+    apt-get update \
+    && apt-get install -y libelf1 libnuma1 \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*; \
+fi
 COPY --from=archive /bin /usr/bin
 ENV PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
 COPY --from=archive /lib/ollama /usr/lib/ollama
 ENV LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
 ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
 ENV NVIDIA_VISIBLE_DEVICES=all
+ENV MTHREADS_DRIVER_CAPABILITIES=compute,utility
+ENV MTHREADS_VISIBLE_DEVICES=all
 ENV OLLAMA_HOST=0.0.0.0:11434
 EXPOSE 11434
 ENTRYPOINT ["/bin/ollama"]

LICENSE

@@ -1,3 +1,14 @@
+The ollama-musa from Moore Threads is licensed under the MIT License listed below.
+Copyright (c) 2025 Moore Threads Technology Co., Ltd ("Moore Threads"). All rights reserved.
+
+-------------------------------------------------------------------------
+The following copyright statements and licenses apply to various open source software
+packages (or portions thereof) that are distributed with this ollama-musa. ollama-musa that
+includes this file does not necessarily use all the open source software packages referred
+to below and may also only use portions of a given package. Some open source software
+packages referred to below may have been modified by Moore Threads Technology Co., Ltd.
+-------------------------------------------------------------------------
+
 MIT License
 
 Copyright (c) Ollama