[MUSA] Add installation guide and Dockerfile (#36709)

Co-authored-by: zhiguo.qin <zhiguo.qin@mthreads.com> Co-authored-by: Kangyan-Zhou <zky314343421@gmail.com>
2026-09-06 20:13:53 -05:00
parent 707da81e84
commit 39a80354aa
6 changed files with 244 additions and 19 deletions
@@ -0,0 +1,126 @@
+# syntax=docker/dockerfile:1.7
+
+# Usage:
+#   docker build -f docker/musa.Dockerfile -t sglang:main-musa .
+#   docker run --rm -it --network=host \
+#     --env MTHREADS_VISIBLE_DEVICES=0 \
+#     --env MTHREADS_DRIVER_CAPABILITIES=all \
+#     --shm-size=32g \
+#     sglang:main-musa
+
+ARG BASE_IMAGE=ubuntu:22.04
+
+FROM ${BASE_IMAGE}
+
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+
+ARG DEBIAN_FRONTEND=noninteractive
+ARG MUSA_APT_SOURCE=https://dl.mthreads.com/repo/repository/ubuntu2204/
+ARG MUSA_PIP_INDEX_URL=https://dl.mthreads.com/repo/api/pypi/pypi/simple
+ARG PYPI_INDEX_URL=https://pypi.org/simple
+
+ENV MUSA_HOME=/usr/local/musa \
+    MATE_MUSA_ARCH_LIST=3.1 \
+    PATH=/usr/local/musa/bin:/usr/local/musa/mudnn/bin:${PATH} \
+    LD_LIBRARY_PATH=/usr/local/musa/lib:/usr/local/musa/mudnn/lib:/usr/local/mtshmem/lib:/usr/local/lib \
+    TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions \
+    TRITON_CACHE_DIR=/root/.triton/cache
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends \
+        bash \
+        build-essential \
+        ca-certificates \
+        cmake \
+        curl \
+        ffmpeg \
+        g++ \
+        gcc \
+        git \
+        libmkl-core \
+        libmkl-def \
+        libmkl-gnu-thread \
+        libmkl-intel-lp64 \
+        libmkl-vml-def \
+        libomp-dev \
+        libopenmpi3 \
+        libsndfile1 \
+        ninja-build \
+        pkg-config \
+        python-is-python3 \
+        python3 \
+        python3-dev \
+        python3-pip \
+        sox \
+    && true
+
+RUN echo "deb [trusted=true] ${MUSA_APT_SOURCE} jammy main" \
+        > /etc/apt/sources.list.d/mthreads-musa.list \
+    && apt-get update \
+    && apt-get install -y --no-install-recommends \
+        libmthreads-compute \
+        libmthreads-mtml \
+        libmublas-5-2 \
+        libmublaslt-5-2 \
+        libmudnn3-dev-musa-5-2 \
+        libmudnn3-musa-5-2 \
+        libmufft-5-2 \
+        libmupp-5-2 \
+        libmurand-5-2 \
+        libmusolver-5-2 \
+        libmusparse-5-2 \
+        mccl-s5000 \
+        mccl-s5000-dev \
+        mtcc-5-2 \
+        musa-mualg-5-2 \
+        musa-mupti-5-2 \
+        musa-musart-5-2 \
+        musa-muthrust-5-2 \
+        musa-toolkit-5-2 \
+        musa-toolkit-5-2-config-common \
+    && printf '%s\n' \
+        "${MUSA_HOME}/lib" \
+        "${MUSA_HOME}/mudnn/lib" \
+        "/usr/local/mtshmem/lib" \
+        "/usr/lib/x86_64-linux-gnu" \
+        > /etc/ld.so.conf.d/musa-runtime.conf \
+    && ln -sf /usr/lib/x86_64-linux-gnu/libmkl_intel_lp64.so /usr/lib/x86_64-linux-gnu/libmkl_intel_lp64.so.2 \
+    && ln -sf /usr/lib/x86_64-linux-gnu/libmkl_core.so /usr/lib/x86_64-linux-gnu/libmkl_core.so.2 \
+    && ln -sf /usr/lib/x86_64-linux-gnu/libmkl_gnu_thread.so /usr/lib/x86_64-linux-gnu/libmkl_gnu_thread.so.2 \
+    && ldconfig
+
+RUN python -m pip install --upgrade pip "setuptools<82" wheel
+
+WORKDIR /workspace/sglang
+COPY . .
+
+# Pip does not prioritize --index-url over --extra-index-url for equal-version
+# candidates. Reinstall Triton from the MUSA index in this same layer.
+RUN cp python/pyproject_other.toml python/pyproject.toml \
+    && python -m pip install -e "python[all_musa]" \
+        --index-url "${MUSA_PIP_INDEX_URL}" \
+        --extra-index-url "${PYPI_INDEX_URL}" \
+        --trusted-host dl.mthreads.com \
+        --no-build-isolation \
+    && python -m pip install --no-cache-dir --force-reinstall --no-deps \
+        --index-url "${MUSA_PIP_INDEX_URL}" \
+        --trusted-host dl.mthreads.com \
+        "triton==3.2.0" \
+    && python -c "import triton.backends.mtgpu" \
+    && ! python -m pip freeze | grep -E '^(nvidia-|cuda-)'
+
+RUN cd python/sglang/kernels/aot \
+    && cp pyproject_musa.toml pyproject.toml \
+    && MTGPU_TARGET=mp_31 python setup_musa.py install
+
+# Keep this check in a single shell command: the legacy Docker builder does not
+# reliably preserve Dockerfile heredocs and can turn the check into a silent
+# `python -` EOF success.
+RUN python -c \
+    "import torch; \
+assert getattr(torch.version, 'musa', None), torch.__version__; \
+assert hasattr(torch, 'musa'); \
+import torchada, triton, tilelang, sglang; \
+import triton.backends.mtgpu"
+
+CMD ["/bin/bash"]