musa: upgrade musa sdk to rc4.2.0 (#13)

* musa: upgrade musa sdk to 4.2.0

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

* musa: remove mudnn and apply wa for libmusa.so

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

* musa: restore rc in docker image tag

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

---------

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

Author: yeahdongcn

CMakeLists.txt

@@ -160,7 +160,7 @@ if(MUSAToolkit_FOUND)
         install(TARGETS ggml-musa
             RUNTIME_DEPENDENCIES
                 DIRECTORIES ${MUSAToolkit_BIN_DIR} ${MUSAToolkit_LIBRARY_DIR}
-                PRE_INCLUDE_REGEXES mudnn mublas musart musa
+                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

Dockerfile

@@ -6,7 +6,8 @@ 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
+ARG MUSAVERSION=rc4.2.0
+ARG UBUNTUVERSION=22.04
 
 # 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
@@ -105,7 +106,7 @@ FROM scratch AS rocm
 COPY --from=rocm-6 dist/lib/ollama /lib/ollama
 
 # Moore Threads (MUSA) build stages
-FROM mthreads/musa:${MUSAVERSION}-mudnn-devel-ubuntu22.04 AS musa-4
+FROM mthreads/musa:${MUSAVERSION}-devel-ubuntu${UBUNTUVERSION}-amd64 AS musa-4
 RUN apt-get update \
     && apt-get install -y curl \
     && apt-get clean \
@@ -119,6 +120,14 @@ 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
+# TODO: Remove the following lines in next release
+RUN cd dist/lib/ollama/musa_v4 && \
+    for f in libmusa.so.*; do \
+        if [ -f "$f" ] && [[ "$f" =~ ^libmusa\.so\.[0-9]+$ ]]; then \
+            ln -sf "$f" libmusa.so; \
+            break; \
+        fi; \
+    done
 
 FROM scratch AS musa
 COPY --from=musa-4 dist/lib/ollama/musa_v4 /lib/ollama/musa_v4
@@ -127,7 +136,7 @@ FROM ${FLAVOR} AS archive
 COPY --from=cpu dist/lib/ollama /lib/ollama
 COPY --from=build /bin/ollama /bin/ollama
 
-FROM ubuntu:22.04
+FROM ubuntu:${UBUNTUVERSION}
 RUN apt-get update \
     && apt-get install -y ca-certificates \
     && apt-get clean \

ml/backend/ggml/ggml/src/ggml-musa/CMakeLists.txt

@@ -34,8 +34,12 @@ if (MUSAToolkit_FOUND)
     list(APPEND GGML_SOURCES_MUSA ${SRCS})
     file(GLOB   SRCS "../ggml-cuda/template-instances/mmq*.cu")
     list(APPEND GGML_SOURCES_MUSA ${SRCS})
-    file(GLOB   SRCS "../ggml-musa/*.cu")
-    list(APPEND GGML_SOURCES_MUSA ${SRCS})
+
+    if (GGML_MUSA_MUDNN_COPY)
+        file(GLOB   SRCS "../ggml-musa/*.cu")
+        list(APPEND GGML_SOURCES_MUSA ${SRCS})
+        add_compile_definitions(GGML_MUSA_MUDNN_COPY)
+    endif()
 
     if (GGML_CUDA_FA_ALL_QUANTS)
         file(GLOB   SRCS "../ggml-cuda/template-instances/fattn-vec*.cu")
@@ -97,10 +101,16 @@ if (MUSAToolkit_FOUND)
     endif()
 
     if (GGML_STATIC)
-        # TODO: mudnn has not provided static libraries yet
         target_link_libraries(ggml-musa PRIVATE MUSA::musart_static MUSA::mublas_static)
+        # TODO: mudnn has not provided static libraries yet
+        # if (GGML_MUSA_MUDNN_COPY)
+        #     target_link_libraries(ggml-musa PRIVATE mudnn_static)
+        # endif()
     else()
-        target_link_libraries(ggml-musa PRIVATE MUSA::musart MUSA::mublas mudnn)
+        target_link_libraries(ggml-musa PRIVATE MUSA::musart MUSA::mublas)
+        if (GGML_MUSA_MUDNN_COPY)
+            target_link_libraries(ggml-musa PRIVATE mudnn)
+        endif()
     endif()
 
     if (GGML_CUDA_NO_VMM)