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

Co-authored-by: zhiguo.qin <zhiguo.qin@mthreads.com>
Co-authored-by: Kangyan-Zhou <zky314343421@gmail.com>

docker/musa.Dockerfile

@@ -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"]

docs/docs/hardware-platforms/mthreads_gpu.mdx

@@ -4,7 +4,20 @@ metatags:
     description: "Run SGLang on Moore Threads GPUs."
 ---
 
-This document describes how run SGLang on Moore Threads GPUs. If you encounter issues or have questions, please [open an issue](https://github.com/sgl-project/sglang/issues).
+This document describes how to run SGLang on Moore Threads GPUs. If you
+encounter issues or have questions, please [open an issue](https://github.com/sgl-project/sglang/issues).
+
+SGLang uses the MUSA runtime through
+[`torch_musa`](https://github.com/MooreThreads/torch_musa) and
+[`torchada`](https://github.com/MooreThreads/torchada). The
+`python[all_musa]` extra installs the MUSA torch, Triton, TileLang, MATE, and
+runtime Python stack from the Moore Threads package index, with public PyPI as
+the fallback for common Python dependencies.
+
+## Prerequisites
+
+Install the Moore Threads driver and MUSA toolkit before installing SGLang. For
+MTT S5000, follow the [official driver installation guide](https://docs.mthreads.com/driver-linux-server/driver-linux-server-doc-online/MTT_S5000/install_guide).
 
 ## Install SGLang
 
@@ -12,18 +25,92 @@ You can install SGLang using one of the methods below.
 
 ### Install from Source
 
-```bash
+Start from an environment with the prerequisites above installed.
-# Use the default branch
+
+```bash Command
 git clone https://github.com/sgl-project/sglang.git
 cd sglang
 
-# Compile sgl-kernel
+python -m pip install --upgrade pip "setuptools<82" wheel
-pip install --upgrade pip
+cp python/pyproject_other.toml python/pyproject.toml
-cd python/sglang/kernels/aot
+python -m pip install -e "python[all_musa]" \
-python setup_musa.py install
+  --index-url https://dl.mthreads.com/repo/api/pypi/pypi/simple \
+  --extra-index-url https://pypi.org/simple \
+  --trusted-host dl.mthreads.com \
+  --no-build-isolation
 
-# Install sglang python package along with diffusion support
+# Ensure the MUSA Triton wheel (with triton.backends.mtgpu) is installed.
-cd ../../../..
+python -m pip install --no-cache-dir --force-reinstall --no-deps \
-rm -f python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml
+  --index-url https://dl.mthreads.com/repo/api/pypi/pypi/simple \
-pip install -e "python[all_musa]"
+  --trusted-host dl.mthreads.com \
+  triton==3.2.0
+python -c 'import triton.backends.mtgpu'
+
+cd python/sglang/kernels/aot
+cp pyproject_musa.toml pyproject.toml
+MTGPU_TARGET=mp_31 python setup_musa.py install
+```
+
+### Install Using Docker
+
+The Dockerfile installs the MUSA runtime and Python stack declared by
+`python/pyproject_other.toml`.
+
+```bash Command
+git clone https://github.com/sgl-project/sglang.git
+cd sglang
+
+docker build -f docker/musa.Dockerfile -t sglang:main-musa .
+```
+
+Run the image with MUSA devices exposed by the host container toolkit.
+
+```bash Command
+alias drun='docker run -it --rm --network=host \
+  --env MTHREADS_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+  --env MTHREADS_DRIVER_CAPABILITIES=all \
+  --shm-size=32g \
+  -v ~/.cache/huggingface:/root/.cache/huggingface'
+
+drun sglang:main-musa \
+  python3 -c 'import triton.backends.mtgpu; print("MUSA Triton backend OK")'
+```
+
+Launch the server:
+
+```bash Command
+drun sglang:main-musa \
+  python3 -m sglang.launch_server \
+  --model-path Qwen/Qwen3-8B \
+  --host 0.0.0.0 \
+  --port 30000
+```
+
+When the server displays `The server is fired up and ready to roll!`, startup is
+successful.
+
+## Verify
+
+You can send an OpenAI-compatible request to the engine:
+
+```bash Command
+curl http://127.0.0.1:30000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen3-8B",
+    "messages": [{"role": "user", "content": "Hello!"}],
+    "max_tokens": 32
+  }'
+```
+
+Or run a benchmark:
+
+```bash Command
+drun sglang:main-musa \
+  python3 -m sglang.bench_serving \
+  --backend sglang \
+  --dataset-name random \
+  --num-prompts 4000 \
+  --random-input 128 \
+  --random-output 128
 ```

docs/docs/hardware-platforms/overview.mdx

@@ -10,3 +10,4 @@ description: Platform-specific guides for running SGLang on GPUs, TPUs, NPUs, CP
 - [NVIDIA Jetson Orin](./nvidia_jetson)
 - [TPU](./tpu)
 - [XPU](./xpu)
+- [Moore Threads GPUs](./mthreads_gpu)

docs/index.mdx

@@ -50,7 +50,7 @@ import { PopularModels } from "/src/snippets/_popular_models.jsx";
 
 <Card title="Extensive Hardware Support" icon="microchip">
   Native support across <a href="./docs/hardware-platforms/overview">Hardware Platforms</a>
-  including NVIDIA, AMD, Intel Xeon, Google TPU, and Ascend NPU accelerators.
+  including NVIDIA, AMD, Intel Xeon, Google TPU, Ascend NPU, and Moore Threads MUSA accelerators.
 </Card>
 
   <Card title="Community & Training" icon="users">

python/pyproject_other.toml

@@ -144,15 +144,26 @@ srt_hpu = ["sglang[runtime_common]"]
 
 # https://docs.sglang.io/platforms/mthreads_gpu.md
 srt_musa = [
-  "deep-gemm>=0.1.3",
+  "apache-tvm-ffi==0.1.11.post1+musa.1",
-  "flash_attn_3>=0.1.4",
+  "deep_ep==1.1.0+musa5.2.0torch2.11.0.post1",
-  "mate>=0.2.0",
+  "deep-gemm==0.2.6+musa",
+  "flash_attn_3==0.2.6+musa",
+  "flash_mla==0.2.6+musa",
+  "flashinfer-python==0.2.6+musa",
+  "mate==0.2.6",
+  "mate-mubin==0.2.6",
   "mthreads-ml-py",
   "numpy<2.0",
+  "sageattention==0.2.6+musa",
   "sglang[runtime_common]",
-  "torch",
+  "tilelang_musa==0.1.12+musa.2",
-  "torch_musa",
+  "torch==2.11.0.post1+musa5.2.0",
-  "torchada>=0.1.74",
+  "torch-c-dlpack-ext==0.1.5",
+  "torch_musa==2.11.0.post1+musa5.2.0",
+  "torchada>=0.1.86",
+  "torchaudio==2.11.0+musa5.2.0",
+  "torchvision==0.26.0.post1+musa5.2.0",
+  "triton==3.2.0",
 ]
 
 diffusion_musa = [

python/sglang/kernels/aot/pyproject_musa.toml

@@ -3,7 +3,7 @@ requires = [
   "setuptools>=75.0",
   "scikit-build-core>=0.10",
   "torch",
-  "torchada>=0.1.74",
+  "torchada>=0.1.86",
   "wheel",
 ]
 build-backend = "setuptools.build_meta"