39a80354aa
Co-authored-by: zhiguo.qin <zhiguo.qin@mthreads.com> Co-authored-by: Kangyan-Zhou <zky314343421@gmail.com>
127 lines
4.1 KiB
Docker
127 lines
4.1 KiB
Docker
# 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"]
|