[ROCm][AMD][Installation] add LMCache kv-connector installation and runtime packages to docker image (#51208)
[ROCm][AMD][Installation] add LMCache kv-connector installation and runtime packages to docker image (#51208) Signed-off-by: Hongxia Yang <hongxia.yang@amd.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Hongxia Yang committed
Aug 17, 2026 at 17:58 UTC
e68fb75b250e33f013162582ece39e766a424df8
3 files changed
+88
-2
.buildkite/release-pipeline.yaml
+1
@@ -906,6 +906,7 @@ steps:
906
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
907
--build-arg SCCACHE_REGION_NAME=us-west-2 \
908
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
909
+ --build-arg INSTALL_LMCACHE=true \
910
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
911
--target vllm-openai \
912
--progress plain \
.buildkite/scripts/ci-bake-rocm.sh
+1
-1
@@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm"
17
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
18
DEFAULT_CI_BASE_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils"
19
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
20
-DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
20
+DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl lmcache_source build_lmcache build_rocshmem build_deepep mori_base ci_base"
21
DEFAULT_CI_BASE_METADATA_VERSION="3"
22
# ROCm CI forces REMOTE_VLLM=0, so content identity covers only the selected
23
# local-source stages rather than unreachable remote-fetch alternatives.
docker/Dockerfile.rocm
+86
-1
@@ -6,6 +6,7 @@ ARG COMMON_WORKDIR=/app
6
ARG BASE_IMAGE=rocm/vllm-dev:base
7
ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base
8
ARG ROCM_TRITON_KERNELS_COMMIT=0f380657dbf3ee86eb57558ff71df24f03b5d4e7
9
+ARG INSTALL_LMCACHE=false
10
# NIC backend for MoRI RDMA support.
11
# By default (all), drivers and userspace libraries for all supported NIC types
12
# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
@@ -471,6 +472,71 @@ RUN cd /opt/nixl && \
472
/tmp/nixl_wheels/repaired/*.whl && \
473
cp /tmp/nixl_wheels/repaired/*.whl /app/install
474
475
+# Fetch and verify pinned LMCache source.
476
+FROM base AS lmcache_source
477
+ARG LMCACHE_REPO="https://github.com/LMCache/LMCache.git"
478
+ARG LMCACHE_REF="140819c9d57a975dbc5678a6459a218e544cb58b"
479
+ARG LMCACHE_VERSION="0.5.3"
480
+
481
+# Remove NVIDIA-only dependencies and accept the release image's NumPy version.
482
+RUN printf '%s\n' "${LMCACHE_REF}" | grep -Eq '^[0-9a-f]{40}$' \
483
+ && git clone --depth=1 --branch "v${LMCACHE_VERSION}" \
484
+ "${LMCACHE_REPO}" /app/lmcache \
485
+ && test "$(git -C /app/lmcache rev-parse HEAD)" = "${LMCACHE_REF}" \
486
+ && test "$(grep -Ec '^(cufile-python|nvtx)([<=>[:space:]]|$)' \
487
+ /app/lmcache/requirements/common.txt)" -eq 2 \
488
+ && test "$(grep -Fxc 'numpy<=2.2.6' \
489
+ /app/lmcache/requirements/common.txt)" -eq 1 \
490
+ && sed -i -E \
491
+ -e '/^(cufile-python|nvtx)([<=>[:space:]]|$)/d' \
492
+ -e 's/^numpy<=2\.2\.6$/numpy<=2.3.5/' \
493
+ /app/lmcache/requirements/common.txt \
494
+ && rm -rf /app/lmcache/.git
495
+
496
+FROM base AS build_lmcache
497
+COPY --from=lmcache_source \
498
+ /app/lmcache/requirements/build.txt /tmp/lmcache-build-requirements.txt
499
+RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
500
+ uv pip install --system -r /tmp/lmcache-build-requirements.txt
501
+
502
+ARG LMCACHE_VERSION="0.5.3"
503
+ARG LMCACHE_ROCM_ARCH="gfx942;gfx950"
504
+ARG USE_SCCACHE
505
+
506
+COPY --from=lmcache_source /app/lmcache /app/lmcache
507
+WORKDIR /app/lmcache
508
+
509
+RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
510
+ set -eu; \
511
+ printf '%s\n' "${LMCACHE_ROCM_ARCH}" \
512
+ | grep -Eq '^gfx[0-9a-z]+(;gfx[0-9a-z]+)*$'; \
513
+ offload_flags="$(printf '%s' "${LMCACHE_ROCM_ARCH}" \
514
+ | sed 's/;/ --offload-arch=/g; s/^/--offload-arch=/')"; \
515
+ if [ "${USE_SCCACHE:-0}" = "1" ]; then \
516
+ export HIP_CLANG_PATH=/opt/sccache-wrappers; \
517
+ fi; \
518
+ mkdir -p /app/install; \
519
+ SETUPTOOLS_SCM_PRETEND_VERSION="${LMCACHE_VERSION}" \
520
+ PYTORCH_ROCM_ARCH="${LMCACHE_ROCM_ARCH}" \
521
+ HIPCC_COMPILE_FLAGS_APPEND="${offload_flags}" \
522
+ BUILD_WITH_HIP=1 CXX=hipcc \
523
+ python3 setup.py bdist_wheel --dist-dir=/app/install; \
524
+ test "$(find /app/install -maxdepth 1 -name '*.whl' | wc -l)" -eq 1
525
+
526
+RUN --mount=type=tmpfs,target=/tmp/lmcache-wheel \
527
+ set -eu; \
528
+ python3 -m zipfile -e /app/install/*.whl /tmp/lmcache-wheel; \
529
+ c_ops="$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so')"; \
530
+ test "$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so' | wc -l)" -eq 1; \
531
+ expected="$(printf '%s' "${LMCACHE_ROCM_ARCH}" | tr ';' '\n' | sort -u)"; \
532
+ actual="$(/opt/rocm/llvm/bin/llvm-objdump --offloading "${c_ops}" \
533
+ | grep -oE 'gfx[0-9a-z]+' | sort -u)"; \
534
+ if [ "${actual}" != "${expected}" ]; then \
535
+ printf 'LMCache code objects mismatch\nexpected:\n%s\nactual:\n%s\n' \
536
+ "${expected}" "${actual}" >&2; \
537
+ exit 1; \
538
+ fi
539
+
540
# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
541
# invalidate the slow ROCShmem build.
542
FROM base AS build_rocshmem
@@ -797,6 +863,13 @@ RUN --mount=type=cache,target=/root/.cache/uv \
863
uv pip install --system /tmp/vllm_test_utils \
864
&& rm -rf /tmp/vllm_test_utils
865
866
+RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \
867
+ --mount=type=bind,source=requirements/test/rocm.txt,target=/tmp/lmcache-constraints.txt \
868
+ --mount=type=cache,target=/root/.cache/uv \
869
+ uv pip install --system --constraint /tmp/lmcache-constraints.txt \
870
+ /lmcache_install/*.whl \
871
+ && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
872
+
873
# -----------------------
874
# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base.
875
FROM ${CI_BASE_IMAGE} AS test
@@ -825,7 +898,7 @@ ENV VLLM_GPU_SYNC_CHECK=error
898
899
# -----------------------
900
# Final vLLM image
828
-FROM mori_base AS final
901
+FROM mori_base AS final_common
902
903
RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
904
@@ -901,6 +974,18 @@ RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
974
975
CMD ["/bin/bash"]
976
977
+FROM final_common AS final_lmcache_false
978
+
979
+FROM final_common AS final_lmcache_true
980
+RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \
981
+ --mount=type=cache,target=/root/.cache/uv \
982
+ uv pip install --system /lmcache_install/*.whl \
983
+ && test "$(python3 -c 'import numpy; print(numpy.__version__)')" = "2.3.5" \
984
+ && uv pip check --system \
985
+ && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
986
+
987
+FROM final_lmcache_${INSTALL_LMCACHE} AS final
988
+
989
#Set entrypoint for vllm-openai official images
990
FROM final AS vllm-openai
991
ENTRYPOINT ["vllm", "serve"]