wehub-resource-sync
127 lines
4.6 KiB
Docker
127 lines
4.6 KiB
Docker
# The LMCache ROCm Dockerfile builds an LMCache image integrated with
|
|
# vLLM-ROCm for AMD Instinct GPUs (MI300X/MI325X, MI350X/MI355X).
|
|
#
|
|
# Please update any changes made here to:
|
|
# docs/source/developer_guide/docker_file.rst
|
|
# docs/source/getting_started/installation.rst
|
|
# docs/source/production/docker_deployment.rst
|
|
#
|
|
# Key differences from the CUDA Dockerfile:
|
|
# - Uses ROCm base image instead of CUDA
|
|
# - Builds HIP extensions (BUILD_WITH_HIP=1, CXX=hipcc)
|
|
# - Skips flashinfer (not available on ROCm)
|
|
# - Skips NIXL (not yet available on ROCm)
|
|
|
|
ARG ROCM_VERSION=7.0
|
|
ARG BASE_IMAGE=rocm/dev-ubuntu-24.04:${ROCM_VERSION}-complete
|
|
|
|
#################### BASE BUILD IMAGE ####################
|
|
|
|
FROM ${BASE_IMAGE} AS base
|
|
|
|
ARG PYTHON_VERSION=3.12
|
|
ENV DEBIAN_FRONTEND=noninteractive
|
|
ENV PATH="/opt/venv/bin:${PATH}"
|
|
|
|
# Install Python and other dependencies
|
|
RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
|
|
&& echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \
|
|
&& apt-get update -y \
|
|
&& apt-get install -y --no-install-recommends \
|
|
ccache software-properties-common git curl sudo \
|
|
python3 python3-dev python3-venv python3-pip tzdata \
|
|
&& curl -LsSf https://astral.sh/uv/install.sh | sh \
|
|
&& mv ~/.local/bin/uv /usr/local/bin/ \
|
|
&& mv ~/.local/bin/uvx /usr/local/bin/ \
|
|
&& uv venv /opt/venv \
|
|
&& . /opt/venv/bin/activate \
|
|
&& python3 --version
|
|
|
|
WORKDIR /workspace
|
|
|
|
# ROCm GPU architectures — use PYTORCH_ROCM_ARCH (comma-separated)
|
|
# which is what PyTorch's extension builder reads for ROCm targets.
|
|
ARG PYTORCH_ROCM_ARCH="gfx942,gfx950"
|
|
ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
|
|
|
|
#################### vLLM-ROCm & LMCache (Build) ##########################
|
|
|
|
FROM base AS image-build
|
|
|
|
ARG ROCM_VERSION
|
|
|
|
COPY ./requirements/build.txt build.txt
|
|
|
|
ARG max_jobs=2
|
|
ENV MAX_JOBS=${max_jobs}
|
|
|
|
ARG VLLM_VERSION=nightly
|
|
|
|
RUN --mount=type=cache,target=/root/.cache/pip \
|
|
. /opt/venv/bin/activate && \
|
|
uv pip install -r build.txt
|
|
|
|
ARG LMCACHE_COMMIT_ID=1
|
|
|
|
COPY . /workspace/LMCache
|
|
WORKDIR /workspace/LMCache
|
|
|
|
# Build LMCache with HIP extensions and install vLLM-ROCm.
|
|
# - BUILD_WITH_HIP=1: compile HIP kernels instead of CUDA
|
|
# - CXX=hipcc: required by setup.py for HIP compilation
|
|
# - --no-build-isolation: use the ROCm torch installed by vLLM
|
|
# - vLLM ROCm wheels from vllm.ai index (dynamic ROCM_TAG)
|
|
# - Skip flashinfer (CUDA-only)
|
|
RUN --mount=type=cache,target=/root/.cache/ccache,id=ccache \
|
|
--mount=type=cache,target=/root/.cache/uv,id=uv-cache,sharing=locked \
|
|
. /opt/venv/bin/activate && \
|
|
ROCM_TAG=rocm$(echo ${ROCM_VERSION} | tr -d '.') && \
|
|
if [ "$VLLM_VERSION" = "nightly" ]; then \
|
|
uv pip install --prerelease=allow \
|
|
'vllm[runai,tensorizer]' \
|
|
--extra-index-url https://wheels.vllm.ai/nightly/${ROCM_TAG} \
|
|
--index-strategy unsafe-best-match ; \
|
|
else \
|
|
uv pip install --prerelease=allow \
|
|
"vllm[runai,tensorizer]==${VLLM_VERSION}" \
|
|
--extra-index-url https://wheels.vllm.ai/${ROCM_TAG} \
|
|
--index-strategy unsafe-best-match ; \
|
|
fi && \
|
|
python3 -c 'import torch; print("TORCH=", torch.__version__, "HIP=", torch.version.hip)' && \
|
|
BUILD_WITH_HIP=1 CXX=hipcc python3 setup.py bdist_wheel --dist-dir=dist_lmcache && \
|
|
uv pip install ./dist_lmcache/*.whl --verbose
|
|
|
|
WORKDIR /workspace
|
|
ENTRYPOINT ["/opt/venv/bin/vllm", "serve"]
|
|
|
|
#################### vLLM-ROCm & LMCache (Release) #######################
|
|
|
|
FROM base AS image-release
|
|
|
|
ARG ROCM_VERSION
|
|
|
|
# Install vLLM-ROCm and build LMCache from source with HIP extensions.
|
|
# PyPI lmcache>=0.4 ships only binary wheels (no sdist), so --no-binary
|
|
# cannot force a source build. Cloning from git ensures HIP extensions
|
|
# are compiled against the ROCm torch installed by vLLM.
|
|
# Build deps are installed from the cloned source's requirements/build.txt
|
|
# to avoid version mismatches with a local COPY.
|
|
ARG max_jobs=2
|
|
ENV MAX_JOBS=${max_jobs}
|
|
|
|
ARG LMCACHE_RELEASE_VERSION=dev
|
|
RUN . /opt/venv/bin/activate && \
|
|
ROCM_TAG=rocm$(echo ${ROCM_VERSION} | tr -d '.') && \
|
|
uv pip install --prerelease=allow 'vllm[runai,tensorizer]' \
|
|
--extra-index-url https://wheels.vllm.ai/${ROCM_TAG} \
|
|
--index-strategy unsafe-best-match && \
|
|
git clone --depth 1 -b ${LMCACHE_RELEASE_VERSION} \
|
|
https://github.com/LMCache/LMCache.git /tmp/lmcache && \
|
|
cd /tmp/lmcache && \
|
|
uv pip install -r requirements/build.txt && \
|
|
BUILD_WITH_HIP=1 CXX=hipcc uv pip install --no-build-isolation . --verbose && \
|
|
rm -rf /tmp/lmcache
|
|
|
|
WORKDIR /workspace
|
|
ENTRYPOINT ["/opt/venv/bin/vllm", "serve"]
|