[ROCm][Installation] Add mooncake package to image using public wheels (#55002)
Signed-off-by: Giuseppe Grossi <ggrossi@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Giuseppe Grossi committed
Sep 2, 2026 at 22:48 UTC
e47356c63e4ab013bfe2af166b0f3952c30ed30b
4 files changed
+7
-168
.buildkite/scripts/ci-bake-rocm.sh
+3
-40
@@ -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 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 lmcache_source build_lmcache build_rocshmem build_deepep mooncake_source build_mooncake package_mooncake export_mooncake 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.
@@ -1684,12 +1684,9 @@ ci_base_metadata_pairs() {
1684
metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
1685
metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
1686
metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
1687
- metadata_pair "vllm.rocm.mooncake_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "MOONCAKE_REPO")"
1688
- metadata_pair "vllm.rocm.mooncake_commit" "${MOONCAKE_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "MOONCAKE_BRANCH")}"
1687
metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
1688
metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
1689
metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
1692
- metadata_pair "vllm.rocm.mooncake_cache_key" "${MOONCAKE_CACHE_KEY:-}"
1690
}
1691
1692
write_ci_base_metadata_annotations() {
@@ -2271,7 +2268,7 @@ extract_dependency_pins() {
2268
return 0
2269
fi
2270
2274
- for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH MOONCAKE_BRANCH; do
2271
+ for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
2272
if [[ -n "${!var:-}" ]]; then
2273
echo "Using provided ${var}: ${!var}"
2274
continue
@@ -2295,11 +2292,9 @@ compute_dependency_cache_keys() {
2292
local ucx_branch=""
2293
local rocshmem_branch=""
2294
local deepep_branch=""
2298
- local mooncake_branch=""
2295
local nixl_material=""
2296
local rocshmem_material=""
2297
local deepep_material=""
2302
- local mooncake_material=""
2298
2299
bake_dir=$(dirname "${VLLM_BAKE_FILE}")
2300
dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
@@ -2307,7 +2302,6 @@ compute_dependency_cache_keys() {
2302
ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
2303
rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
2304
deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
2310
- mooncake_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "MOONCAKE_BRANCH")
2305
2306
if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
2307
nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
@@ -2341,19 +2335,6 @@ compute_dependency_cache_keys() {
2335
export DEEPEP_CACHE_KEY
2336
echo "DeepEP dependency cache key: ${DEEPEP_CACHE_KEY}"
2337
fi
2344
-
2345
- if [[ -n "${mooncake_branch}" ]]; then
2346
- mooncake_material=$(compose_stage_cache_material \
2347
- "${dockerfile_rocm}" \
2348
- "base mooncake_source build_mooncake package_mooncake export_mooncake")
2349
- MOONCAKE_CACHE_KEY=$(
2350
- compose_dependency_cache_key \
2351
- "${mooncake_branch}" \
2352
- "${mooncake_material}"
2353
- )
2354
- export MOONCAKE_CACHE_KEY
2355
- echo "Mooncake dependency cache key: ${MOONCAKE_CACHE_KEY}"
2356
- fi
2338
}
2339
2340
compose_stage_cache_material() {
@@ -2402,13 +2383,6 @@ dependency_cache_ref_for_target() {
2383
printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH:-}"
2384
fi
2385
;;
2405
- mooncake-rocm-ci)
2406
- if [[ -n "${MOONCAKE_CACHE_KEY:-}" ]]; then
2407
- printf '%s\n' "${cache_repo}:mooncake-rocm-${MOONCAKE_CACHE_KEY}"
2408
- elif [[ -n "${MOONCAKE_BRANCH:-}" ]]; then
2409
- printf '%s\n' "${cache_repo}:mooncake-rocm-${MOONCAKE_BRANCH}"
2410
- fi
2411
- ;;
2386
esac
2387
}
2388
@@ -2426,14 +2400,13 @@ resolve_ci_base_dependency_targets() {
2400
local nixl_ref=""
2401
local rocshmem_ref=""
2402
local deepep_ref=""
2429
- local mooncake_ref=""
2403
2404
[[ "${TARGET}" == "ci-base-rocm-ci-with-deps" ]] || return 0
2405
2406
case "${mode}" in
2407
always)
2408
echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
2436
- for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci mooncake-rocm-ci; do
2409
+ for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
2410
if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
2411
add_dependency_cache_target "${target}"
2412
fi
@@ -2483,16 +2456,6 @@ resolve_ci_base_dependency_targets() {
2456
fi
2457
fi
2458
2486
- if [[ "${mode}" != "always" && -n "${MOONCAKE_CACHE_KEY:-}" ]]; then
2487
- mooncake_ref=$(dependency_cache_ref_for_target "mooncake-rocm-ci")
2488
- if dependency_cache_ref_exists "${mooncake_ref}"; then
2489
- echo "Mooncake dependency cache exists: ${mooncake_ref}"
2490
- else
2491
- echo "Mooncake dependency cache missing; will seed: ${mooncake_ref}"
2492
- add_dependency_cache_target "mooncake-rocm-ci"
2493
- fi
2494
- fi
2495
-
2459
# DeepEP inherits from ROCShmem. If ROCShmem is being seeded, seed DeepEP too
2460
# so the pair stays consistent for future ci_base rebuilds.
2461
if printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "rocshmem-rocm-ci" \
docker/Dockerfile.rocm
+1
-103
@@ -7,10 +7,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
-ARG MOONCAKE_BRANCH=df508641d76be20803349bf715d3bfc5205a7b58
11
-ARG MOONCAKE_REPO=https://github.com/kvcache-ai/Mooncake.git
12
-ARG MOONCAKE_GO_VERSION=1.25.10
13
-ARG MOONCAKE_GO_SHA256=42d4f7a32316aa66591eca7e89867256057a4264451aca10570a715b3637ba70
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.
@@ -647,82 +643,6 @@ RUN /bin/bash -lc 'set -euo pipefail; \
643
*) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \
644
esac'
645
650
-FROM base AS mooncake_source
651
-ARG MOONCAKE_GO_VERSION
652
-ARG MOONCAKE_GO_SHA256
653
-ENV PATH=/usr/local/go/bin:${PATH}
654
-ENV GOTOOLCHAIN=local
655
-RUN test "$(uname -m)" = "x86_64" \
656
- && test ! -e /usr/local/go \
657
- && curl --fail --location --retry 5 \
658
- --output /tmp/go.tar.gz \
659
- "https://go.dev/dl/go${MOONCAKE_GO_VERSION}.linux-amd64.tar.gz" \
660
- && echo "${MOONCAKE_GO_SHA256} /tmp/go.tar.gz" | sha256sum --check - \
661
- && tar -C /usr/local -xzf /tmp/go.tar.gz \
662
- && rm /tmp/go.tar.gz \
663
- && test "$(go env GOVERSION)" = "go${MOONCAKE_GO_VERSION}"
664
-
665
-# Build a pinned ROCm Mooncake wheel outside the long-lived ROCm base image.
666
-ARG MOONCAKE_BRANCH
667
-ARG MOONCAKE_REPO
668
-RUN git init mooncake \
669
- && cd mooncake \
670
- && git remote add origin "${MOONCAKE_REPO}" \
671
- && git fetch --depth=1 origin "${MOONCAKE_BRANCH}" \
672
- && git checkout --detach FETCH_HEAD \
673
- && test "$(git rev-parse HEAD)" = "${MOONCAKE_BRANCH}" \
674
- && grep -qx 'GOVER=1.25.9' dependencies.sh \
675
- && sed -i "s/^GOVER=.*/GOVER=${MOONCAKE_GO_VERSION}/" dependencies.sh \
676
- && bash dependencies.sh -y
677
-
678
-FROM mooncake_source AS build_mooncake
679
-ARG USE_SCCACHE
680
-ENV HIPCC_COMPILE_FLAGS_APPEND=--gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/11
681
-RUN --mount=type=cache,id=vllm-rocm-mooncake-go-mod,target=/root/go/pkg/mod \
682
- --mount=type=cache,id=vllm-rocm-mooncake-go-build,target=/root/.cache/go-build \
683
- cd mooncake \
684
- && cmake -G Ninja -B build \
685
- -DUSE_HIP=ON -DUSE_CUDA=OFF -DWITH_EP=OFF \
686
- -DUSE_HTTP=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON \
687
- -DBUILD_UNIT_TESTS=OFF -DBUILD_BENCHMARK=OFF \
688
- -DENABLE_DEBUG_SYMBOLS=OFF \
689
- -DENABLE_SCCACHE="${USE_SCCACHE:-OFF}" \
690
- -DCMAKE_BUILD_TYPE=Release \
691
- -DPython3_EXECUTABLE="$(command -v python3)" \
692
- && cmake --build build -j"${MAX_JOBS:-$(nproc)}" \
693
- && cmake --install build
694
-
695
-FROM build_mooncake AS package_mooncake
696
-# Pin both the outer packaging tools and the isolated PEP 517 build environment.
697
-# patchelf 0.14.3 from dependencies.sh corrupts auditwheel RPATHs.
698
-ENV PIP_CONSTRAINT=/tmp/mooncake-build-constraints.txt
699
-ENV PIP_BUILD_CONSTRAINT=/tmp/mooncake-build-constraints.txt
700
-RUN --mount=type=cache,id=vllm-rocm-mooncake-pip,target=/root/.cache/pip \
701
- printf '%s\n' \
702
- 'pip==26.2.1' \
703
- 'build==1.5.0' \
704
- 'setuptools==79.0.1' \
705
- 'wheel==0.48.0' \
706
- 'auditwheel==6.8.1' \
707
- 'packaging==26.3' \
708
- 'pyproject-hooks==1.2.0' \
709
- 'pyelftools==0.33' \
710
- 'patchelf==0.19.1.0' \
711
- > "${PIP_CONSTRAINT}" \
712
- && python3 -m pip install --upgrade \
713
- pip build setuptools wheel auditwheel patchelf \
714
- && cd mooncake \
715
- && export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib" \
716
- && HIP_BUILD=1 OUTPUT_DIR=dist PLATFORM_TAG=manylinux_2_35_x86_64 \
717
- ./scripts/build_wheel.sh \
718
- && test "$(find mooncake-wheel/dist -maxdepth 1 -type f \
719
- -name 'mooncake_transfer_engine_rocm-*.whl' | wc -l)" -eq 1 \
720
- && mkdir -p /app/install \
721
- && cp mooncake-wheel/dist/mooncake_transfer_engine_rocm-*.whl /app/install/
722
-
723
-FROM scratch AS export_mooncake
724
-COPY --from=package_mooncake /app/install/ /
725
-
646
# -----------------------
647
# vLLM wheel release build stage (for building distributable wheels)
648
# This stage pins dependencies to custom ROCm wheel versions and handles version detection
@@ -953,15 +873,6 @@ RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_instal
873
/lmcache_install/*.whl \
874
&& python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
875
956
-RUN --mount=type=bind,from=export_mooncake,src=/,target=/mooncake_install,ro \
957
- --mount=type=cache,target=/root/.cache/uv \
958
- uv pip install --system --no-deps \
959
- /mooncake_install/mooncake_transfer_engine_rocm-*.whl \
960
- && python3 -c "import aiohttp, msgpack, requests" \
961
- && python3 -c "from mooncake.engine import SUPPORT_HIP, TransferEngine; assert SUPPORT_HIP" \
962
- && python3 -c "from mooncake.store import MooncakeDistributedStore" \
963
- && mooncake_master --version
964
-
876
# -----------------------
877
# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base.
878
FROM ${CI_BASE_IMAGE} AS test
@@ -1041,21 +952,10 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
952
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
953
uv pip install --system /nixl_install/*.whl
954
1044
-RUN --mount=type=bind,from=export_mooncake,src=/,target=/mooncake_install,ro \
1045
- --mount=type=cache,target=/root/.cache/uv \
1046
- uv pip install --system --no-deps \
1047
- /mooncake_install/mooncake_transfer_engine_rocm-*.whl \
1048
- && python3 -c "import aiohttp, msgpack, requests" \
1049
- && python3 -c "from mooncake.engine import SUPPORT_HIP, TransferEngine; assert SUPPORT_HIP" \
1050
- && python3 -c "from mooncake.store import MooncakeDistributedStore" \
1051
- && mooncake_master --version
1052
-
955
ARG COMMON_WORKDIR
956
ARG BASE_IMAGE
957
ARG NIC_BACKEND
958
ARG AINIC_VERSION
1057
-ARG MOONCAKE_BRANCH
1058
-ARG MOONCAKE_REPO
959
960
# Copy over the benchmark scripts as well
961
COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
@@ -1084,9 +984,7 @@ RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
984
ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
985
RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
986
&& echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
1087
- && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt \
1088
- && echo "MOONCAKE_BRANCH=${MOONCAKE_BRANCH}" >> ${COMMON_WORKDIR}/versions.txt \
1089
- && echo "MOONCAKE_REPO=${MOONCAKE_REPO}" >> ${COMMON_WORKDIR}/versions.txt
987
+ && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
988
989
CMD ["/bin/bash"]
990
docker/ci-rocm.hcl
+1
-25
@@ -91,10 +91,6 @@ variable "DEEPEP_BRANCH" {
91
default = ""
92
}
93
94
-variable "MOONCAKE_BRANCH" {
95
- default = ""
96
-}
97
-
94
variable "NIXL_CACHE_KEY" {
95
default = ""
96
}
@@ -107,10 +103,6 @@ variable "DEEPEP_CACHE_KEY" {
103
default = ""
104
}
105
110
-variable "MOONCAKE_CACHE_KEY" {
111
- default = ""
112
-}
113
-
106
# Docker Hub registry cache for AMD builds.
107
#
108
# A separate repo (rocm/vllm-ci-cache) is used for BuildKit layer cache.
@@ -255,7 +247,6 @@ function "get_cache_from_rocm_deps" {
247
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
248
ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""),
249
DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""),
258
- MOONCAKE_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:mooncake-rocm-${MOONCAKE_CACHE_KEY}" : (MOONCAKE_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:mooncake-rocm-${MOONCAKE_BRANCH}" : ""),
250
])
251
}
252
@@ -280,13 +271,6 @@ function "get_cache_to_rocm_deepep" {
271
])
272
}
273
283
-function "get_cache_to_rocm_mooncake" {
284
- params = []
285
- result = compact([
286
- MOONCAKE_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:mooncake-rocm-${MOONCAKE_CACHE_KEY},mode=min" : (MOONCAKE_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:mooncake-rocm-${MOONCAKE_BRANCH},mode=min" : ""),
287
- ])
288
-}
289
-
274
# CI targets
275
276
target "_ci-rocm" {
@@ -417,14 +401,6 @@ target "deepep-rocm-ci" {
401
output = ["type=cacheonly"]
402
}
403
420
-target "mooncake-rocm-ci" {
421
- inherits = ["_common-rocm", "_ci-rocm"]
422
- target = "export_mooncake"
423
- cache-from = get_cache_from_rocm_deps()
424
- cache-to = get_cache_to_rocm_mooncake()
425
- output = ["type=cacheonly"]
426
-}
427
-
404
# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
405
# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
406
# files drifts from the remote image label. Per-PR builds then pull the result
@@ -455,5 +431,5 @@ target "ci-base-rocm-ci" {
431
# Group for ci_base builds -- exports dependency stage caches alongside the
432
# ci_base image so future rebuilds can reuse them independently.
433
group "ci-base-rocm-ci-with-deps" {
458
- targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "mooncake-rocm-ci", "ci-base-rocm-ci"]
434
+ targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
435
}
requirements/rocm.txt
+2
@@ -27,3 +27,5 @@ tilelang==0.1.10
27
apache-tvm-ffi==0.1.10
28
# Required for faster safetensors model loading
29
fastsafetensors >= 0.3.3
30
+# Mooncake transfer engine for KV cache offloading/transfer (ROCm build, published on PyPI)
31
+mooncake-transfer-engine-rocm >= 0.3.13