[ROCm][AMD][Installation] Add mooncake build to rocm base image (#52650)
Signed-off-by: Giuseppe Grossi <ggrossi@amd.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Giuseppe Grossi committed
Sep 2, 2026 at 01:14 UTC
56b54953772d3ec53fe5036449613ec02256d842
3 files changed
+168
-5
.buildkite/scripts/ci-bake-rocm.sh
+40
-3
@@ -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 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 mooncake_source build_mooncake package_mooncake export_mooncake 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,9 +1684,12 @@ 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")}"
1689
metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
1690
metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
1691
metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
1692
+ metadata_pair "vllm.rocm.mooncake_cache_key" "${MOONCAKE_CACHE_KEY:-}"
1693
}
1694
1695
write_ci_base_metadata_annotations() {
@@ -2268,7 +2271,7 @@ extract_dependency_pins() {
2271
return 0
2272
fi
2273
2271
- for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
2274
+ for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH MOONCAKE_BRANCH; do
2275
if [[ -n "${!var:-}" ]]; then
2276
echo "Using provided ${var}: ${!var}"
2277
continue
@@ -2292,9 +2295,11 @@ compute_dependency_cache_keys() {
2295
local ucx_branch=""
2296
local rocshmem_branch=""
2297
local deepep_branch=""
2298
+ local mooncake_branch=""
2299
local nixl_material=""
2300
local rocshmem_material=""
2301
local deepep_material=""
2302
+ local mooncake_material=""
2303
2304
bake_dir=$(dirname "${VLLM_BAKE_FILE}")
2305
dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
@@ -2302,6 +2307,7 @@ compute_dependency_cache_keys() {
2307
ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
2308
rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
2309
deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
2310
+ mooncake_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "MOONCAKE_BRANCH")
2311
2312
if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
2313
nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
@@ -2335,6 +2341,19 @@ compute_dependency_cache_keys() {
2341
export DEEPEP_CACHE_KEY
2342
echo "DeepEP dependency cache key: ${DEEPEP_CACHE_KEY}"
2343
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
2357
}
2358
2359
compose_stage_cache_material() {
@@ -2383,6 +2402,13 @@ dependency_cache_ref_for_target() {
2402
printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH:-}"
2403
fi
2404
;;
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
+ ;;
2412
esac
2413
}
2414
@@ -2400,13 +2426,14 @@ resolve_ci_base_dependency_targets() {
2426
local nixl_ref=""
2427
local rocshmem_ref=""
2428
local deepep_ref=""
2429
+ local mooncake_ref=""
2430
2431
[[ "${TARGET}" == "ci-base-rocm-ci-with-deps" ]] || return 0
2432
2433
case "${mode}" in
2434
always)
2435
echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
2409
- for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
2436
+ for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci mooncake-rocm-ci; do
2437
if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
2438
add_dependency_cache_target "${target}"
2439
fi
@@ -2456,6 +2483,16 @@ resolve_ci_base_dependency_targets() {
2483
fi
2484
fi
2485
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
+
2496
# DeepEP inherits from ROCShmem. If ROCShmem is being seeded, seed DeepEP too
2497
# so the pair stays consistent for future ci_base rebuilds.
2498
if printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "rocshmem-rocm-ci" \
docker/Dockerfile.rocm
+103
-1
@@ -7,6 +7,10 @@ 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
14
# NIC backend for MoRI RDMA support.
15
# By default (all), drivers and userspace libraries for all supported NIC types
16
# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
@@ -643,6 +647,82 @@ RUN /bin/bash -lc 'set -euo pipefail; \
647
*) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \
648
esac'
649
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
+
726
# -----------------------
727
# vLLM wheel release build stage (for building distributable wheels)
728
# This stage pins dependencies to custom ROCm wheel versions and handles version detection
@@ -873,6 +953,15 @@ RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_instal
953
/lmcache_install/*.whl \
954
&& python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
955
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
+
965
# -----------------------
966
# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base.
967
FROM ${CI_BASE_IMAGE} AS test
@@ -952,10 +1041,21 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
1041
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
1042
uv pip install --system /nixl_install/*.whl
1043
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
+
1053
ARG COMMON_WORKDIR
1054
ARG BASE_IMAGE
1055
ARG NIC_BACKEND
1056
ARG AINIC_VERSION
1057
+ARG MOONCAKE_BRANCH
1058
+ARG MOONCAKE_REPO
1059
1060
# Copy over the benchmark scripts as well
1061
COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
@@ -984,7 +1084,9 @@ RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
1084
ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
1085
RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
1086
&& echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
987
- && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${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
1090
1091
CMD ["/bin/bash"]
1092
docker/ci-rocm.hcl
+25
-1
@@ -91,6 +91,10 @@ variable "DEEPEP_BRANCH" {
91
default = ""
92
}
93
94
+variable "MOONCAKE_BRANCH" {
95
+ default = ""
96
+}
97
+
98
variable "NIXL_CACHE_KEY" {
99
default = ""
100
}
@@ -103,6 +107,10 @@ variable "DEEPEP_CACHE_KEY" {
107
default = ""
108
}
109
110
+variable "MOONCAKE_CACHE_KEY" {
111
+ default = ""
112
+}
113
+
114
# Docker Hub registry cache for AMD builds.
115
#
116
# A separate repo (rocm/vllm-ci-cache) is used for BuildKit layer cache.
@@ -247,6 +255,7 @@ function "get_cache_from_rocm_deps" {
255
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}" : ""),
256
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}" : ""),
257
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}" : ""),
259
])
260
}
261
@@ -271,6 +280,13 @@ function "get_cache_to_rocm_deepep" {
280
])
281
}
282
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
+
290
# CI targets
291
292
target "_ci-rocm" {
@@ -401,6 +417,14 @@ target "deepep-rocm-ci" {
417
output = ["type=cacheonly"]
418
}
419
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
+
428
# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
429
# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
430
# files drifts from the remote image label. Per-PR builds then pull the result
@@ -431,5 +455,5 @@ target "ci-base-rocm-ci" {
455
# Group for ci_base builds -- exports dependency stage caches alongside the
456
# ci_base image so future rebuilds can reuse them independently.
457
group "ci-base-rocm-ci-with-deps" {
434
- targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
458
+ targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "mooncake-rocm-ci", "ci-base-rocm-ci"]
459
}