[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>

.buildkite/scripts/ci-bake-rocm.sh

@@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm"
 DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
 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"
 DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
-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"
+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"
 DEFAULT_CI_BASE_METADATA_VERSION="3"
 # ROCm CI forces REMOTE_VLLM=0, so content identity covers only the selected
 # local-source stages rather than unreachable remote-fetch alternatives.
@@ -1684,12 +1684,9 @@ ci_base_metadata_pairs() {
     metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
     metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
     metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
-    metadata_pair "vllm.rocm.mooncake_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "MOONCAKE_REPO")"
-    metadata_pair "vllm.rocm.mooncake_commit" "${MOONCAKE_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "MOONCAKE_BRANCH")}"
     metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
     metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
     metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
-    metadata_pair "vllm.rocm.mooncake_cache_key" "${MOONCAKE_CACHE_KEY:-}"
 }
 
 write_ci_base_metadata_annotations() {
@@ -2271,7 +2268,7 @@ extract_dependency_pins() {
         return 0
     fi
 
-    for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH MOONCAKE_BRANCH; do
+    for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
         if [[ -n "${!var:-}" ]]; then
             echo "Using provided ${var}: ${!var}"
             continue
@@ -2295,11 +2292,9 @@ compute_dependency_cache_keys() {
     local ucx_branch=""
     local rocshmem_branch=""
     local deepep_branch=""
-    local mooncake_branch=""
     local nixl_material=""
     local rocshmem_material=""
     local deepep_material=""
-    local mooncake_material=""
 
     bake_dir=$(dirname "${VLLM_BAKE_FILE}")
     dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
@@ -2307,7 +2302,6 @@ compute_dependency_cache_keys() {
     ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
     rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
     deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
-    mooncake_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "MOONCAKE_BRANCH")
 
     if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
         nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
@@ -2341,19 +2335,6 @@ compute_dependency_cache_keys() {
         export DEEPEP_CACHE_KEY
         echo "DeepEP dependency cache key: ${DEEPEP_CACHE_KEY}"
     fi
-
-    if [[ -n "${mooncake_branch}" ]]; then
-        mooncake_material=$(compose_stage_cache_material \
-            "${dockerfile_rocm}" \
-            "base mooncake_source build_mooncake package_mooncake export_mooncake")
-        MOONCAKE_CACHE_KEY=$(
-            compose_dependency_cache_key \
-                "${mooncake_branch}" \
-                "${mooncake_material}"
-        )
-        export MOONCAKE_CACHE_KEY
-        echo "Mooncake dependency cache key: ${MOONCAKE_CACHE_KEY}"
-    fi
 }
 
 compose_stage_cache_material() {
@@ -2402,13 +2383,6 @@ dependency_cache_ref_for_target() {
                 printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH:-}"
             fi
             ;;
-        mooncake-rocm-ci)
-            if [[ -n "${MOONCAKE_CACHE_KEY:-}" ]]; then
-                printf '%s\n' "${cache_repo}:mooncake-rocm-${MOONCAKE_CACHE_KEY}"
-            elif [[ -n "${MOONCAKE_BRANCH:-}" ]]; then
-                printf '%s\n' "${cache_repo}:mooncake-rocm-${MOONCAKE_BRANCH}"
-            fi
-            ;;
     esac
 }
 
@@ -2426,14 +2400,13 @@ resolve_ci_base_dependency_targets() {
     local nixl_ref=""
     local rocshmem_ref=""
     local deepep_ref=""
-    local mooncake_ref=""
 
     [[ "${TARGET}" == "ci-base-rocm-ci-with-deps" ]] || return 0
 
     case "${mode}" in
         always)
             echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
-            for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci mooncake-rocm-ci; do
+            for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
                 if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
                     add_dependency_cache_target "${target}"
                 fi
@@ -2483,16 +2456,6 @@ resolve_ci_base_dependency_targets() {
         fi
     fi
 
-    if [[ "${mode}" != "always" && -n "${MOONCAKE_CACHE_KEY:-}" ]]; then
-        mooncake_ref=$(dependency_cache_ref_for_target "mooncake-rocm-ci")
-        if dependency_cache_ref_exists "${mooncake_ref}"; then
-            echo "Mooncake dependency cache exists: ${mooncake_ref}"
-        else
-            echo "Mooncake dependency cache missing; will seed: ${mooncake_ref}"
-            add_dependency_cache_target "mooncake-rocm-ci"
-        fi
-    fi
-
     # DeepEP inherits from ROCShmem. If ROCShmem is being seeded, seed DeepEP too
     # so the pair stays consistent for future ci_base rebuilds.
     if printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "rocshmem-rocm-ci" \

docker/Dockerfile.rocm

@@ -7,10 +7,6 @@ ARG BASE_IMAGE=rocm/vllm-dev:base
 ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base
 ARG ROCM_TRITON_KERNELS_COMMIT=0f380657dbf3ee86eb57558ff71df24f03b5d4e7
 ARG INSTALL_LMCACHE=false
-ARG MOONCAKE_BRANCH=df508641d76be20803349bf715d3bfc5205a7b58
-ARG MOONCAKE_REPO=https://github.com/kvcache-ai/Mooncake.git
-ARG MOONCAKE_GO_VERSION=1.25.10
-ARG MOONCAKE_GO_SHA256=42d4f7a32316aa66591eca7e89867256057a4264451aca10570a715b3637ba70
 # NIC backend for MoRI RDMA support.
 # By default (all), drivers and userspace libraries for all supported NIC types
 # (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
@@ -647,82 +643,6 @@ RUN /bin/bash -lc 'set -euo pipefail; \
    *)     echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \
  esac'
 
-FROM base AS mooncake_source
-ARG MOONCAKE_GO_VERSION
-ARG MOONCAKE_GO_SHA256
-ENV PATH=/usr/local/go/bin:${PATH}
-ENV GOTOOLCHAIN=local
-RUN test "$(uname -m)" = "x86_64" \
-    && test ! -e /usr/local/go \
-    && curl --fail --location --retry 5 \
-        --output /tmp/go.tar.gz \
-        "https://go.dev/dl/go${MOONCAKE_GO_VERSION}.linux-amd64.tar.gz" \
-    && echo "${MOONCAKE_GO_SHA256}  /tmp/go.tar.gz" | sha256sum --check - \
-    && tar -C /usr/local -xzf /tmp/go.tar.gz \
-    && rm /tmp/go.tar.gz \
-    && test "$(go env GOVERSION)" = "go${MOONCAKE_GO_VERSION}"
-
-# Build a pinned ROCm Mooncake wheel outside the long-lived ROCm base image.
-ARG MOONCAKE_BRANCH
-ARG MOONCAKE_REPO
-RUN git init mooncake \
-    && cd mooncake \
-    && git remote add origin "${MOONCAKE_REPO}" \
-    && git fetch --depth=1 origin "${MOONCAKE_BRANCH}" \
-    && git checkout --detach FETCH_HEAD \
-    && test "$(git rev-parse HEAD)" = "${MOONCAKE_BRANCH}" \
-    && grep -qx 'GOVER=1.25.9' dependencies.sh \
-    && sed -i "s/^GOVER=.*/GOVER=${MOONCAKE_GO_VERSION}/" dependencies.sh \
-    && bash dependencies.sh -y
-
-FROM mooncake_source AS build_mooncake
-ARG USE_SCCACHE
-ENV HIPCC_COMPILE_FLAGS_APPEND=--gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/11
-RUN --mount=type=cache,id=vllm-rocm-mooncake-go-mod,target=/root/go/pkg/mod \
-    --mount=type=cache,id=vllm-rocm-mooncake-go-build,target=/root/.cache/go-build \
-    cd mooncake \
-    && cmake -G Ninja -B build \
-        -DUSE_HIP=ON -DUSE_CUDA=OFF -DWITH_EP=OFF \
-        -DUSE_HTTP=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON \
-        -DBUILD_UNIT_TESTS=OFF -DBUILD_BENCHMARK=OFF \
-        -DENABLE_DEBUG_SYMBOLS=OFF \
-        -DENABLE_SCCACHE="${USE_SCCACHE:-OFF}" \
-        -DCMAKE_BUILD_TYPE=Release \
-        -DPython3_EXECUTABLE="$(command -v python3)" \
-    && cmake --build build -j"${MAX_JOBS:-$(nproc)}" \
-    && cmake --install build
-
-FROM build_mooncake AS package_mooncake
-# Pin both the outer packaging tools and the isolated PEP 517 build environment.
-# patchelf 0.14.3 from dependencies.sh corrupts auditwheel RPATHs.
-ENV PIP_CONSTRAINT=/tmp/mooncake-build-constraints.txt
-ENV PIP_BUILD_CONSTRAINT=/tmp/mooncake-build-constraints.txt
-RUN --mount=type=cache,id=vllm-rocm-mooncake-pip,target=/root/.cache/pip \
-    printf '%s\n' \
-        'pip==26.2.1' \
-        'build==1.5.0' \
-        'setuptools==79.0.1' \
-        'wheel==0.48.0' \
-        'auditwheel==6.8.1' \
-        'packaging==26.3' \
-        'pyproject-hooks==1.2.0' \
-        'pyelftools==0.33' \
-        'patchelf==0.19.1.0' \
-        > "${PIP_CONSTRAINT}" \
-    && python3 -m pip install --upgrade \
-        pip build setuptools wheel auditwheel patchelf \
-    && cd mooncake \
-    && export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib" \
-    && HIP_BUILD=1 OUTPUT_DIR=dist PLATFORM_TAG=manylinux_2_35_x86_64 \
-        ./scripts/build_wheel.sh \
-    && test "$(find mooncake-wheel/dist -maxdepth 1 -type f \
-        -name 'mooncake_transfer_engine_rocm-*.whl' | wc -l)" -eq 1 \
-    && mkdir -p /app/install \
-    && cp mooncake-wheel/dist/mooncake_transfer_engine_rocm-*.whl /app/install/
-
-FROM scratch AS export_mooncake
-COPY --from=package_mooncake /app/install/ /
-
 # -----------------------
 # vLLM wheel release build stage (for building distributable wheels)
 # 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
         /lmcache_install/*.whl \
     && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
 
-RUN --mount=type=bind,from=export_mooncake,src=/,target=/mooncake_install,ro \
-    --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system --no-deps \
-        /mooncake_install/mooncake_transfer_engine_rocm-*.whl \
-    && python3 -c "import aiohttp, msgpack, requests" \
-    && python3 -c "from mooncake.engine import SUPPORT_HIP, TransferEngine; assert SUPPORT_HIP" \
-    && python3 -c "from mooncake.store import MooncakeDistributedStore" \
-    && mooncake_master --version
-
 # -----------------------
 # Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base.
 FROM ${CI_BASE_IMAGE} AS test
@@ -1041,21 +952,10 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
 RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
     uv pip install --system /nixl_install/*.whl
 
-RUN --mount=type=bind,from=export_mooncake,src=/,target=/mooncake_install,ro \
-    --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system --no-deps \
-        /mooncake_install/mooncake_transfer_engine_rocm-*.whl \
-    && python3 -c "import aiohttp, msgpack, requests" \
-    && python3 -c "from mooncake.engine import SUPPORT_HIP, TransferEngine; assert SUPPORT_HIP" \
-    && python3 -c "from mooncake.store import MooncakeDistributedStore" \
-    && mooncake_master --version
-
 ARG COMMON_WORKDIR
 ARG BASE_IMAGE
 ARG NIC_BACKEND
 ARG AINIC_VERSION
-ARG MOONCAKE_BRANCH
-ARG MOONCAKE_REPO
 
 # Copy over the benchmark scripts as well
 COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
@@ -1084,9 +984,7 @@ RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
 ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
 RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
     && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
-    && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt \
-    && echo "MOONCAKE_BRANCH=${MOONCAKE_BRANCH}" >> ${COMMON_WORKDIR}/versions.txt \
-    && echo "MOONCAKE_REPO=${MOONCAKE_REPO}" >> ${COMMON_WORKDIR}/versions.txt
+    && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
 
 CMD ["/bin/bash"]
 

docker/ci-rocm.hcl

@@ -91,10 +91,6 @@ variable "DEEPEP_BRANCH" {
   default = ""
 }
 
-variable "MOONCAKE_BRANCH" {
-  default = ""
-}
-
 variable "NIXL_CACHE_KEY" {
   default = ""
 }
@@ -107,10 +103,6 @@ variable "DEEPEP_CACHE_KEY" {
   default = ""
 }
 
-variable "MOONCAKE_CACHE_KEY" {
-  default = ""
-}
-
 # Docker Hub registry cache for AMD builds.
 #
 # A separate repo (rocm/vllm-ci-cache) is used for BuildKit layer cache.
@@ -255,7 +247,6 @@ function "get_cache_from_rocm_deps" {
     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}" : ""),
     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}" : ""),
     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}" : ""),
-    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}" : ""),
   ])
 }
 
@@ -280,13 +271,6 @@ function "get_cache_to_rocm_deepep" {
   ])
 }
 
-function "get_cache_to_rocm_mooncake" {
-  params = []
-  result = compact([
-    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" : ""),
-  ])
-}
-
 # CI targets
 
 target "_ci-rocm" {
@@ -417,14 +401,6 @@ target "deepep-rocm-ci" {
   output     = ["type=cacheonly"]
 }
 
-target "mooncake-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "export_mooncake"
-  cache-from = get_cache_from_rocm_deps()
-  cache-to   = get_cache_to_rocm_mooncake()
-  output     = ["type=cacheonly"]
-}
-
 # Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
 # Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
 # files drifts from the remote image label. Per-PR builds then pull the result
@@ -455,5 +431,5 @@ target "ci-base-rocm-ci" {
 # Group for ci_base builds -- exports dependency stage caches alongside the
 # ci_base image so future rebuilds can reuse them independently.
 group "ci-base-rocm-ci-with-deps" {
-  targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "mooncake-rocm-ci", "ci-base-rocm-ci"]
+  targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
 }

requirements/rocm.txt

@@ -27,3 +27,5 @@ tilelang==0.1.10
 apache-tvm-ffi==0.1.10 
 # Required for faster safetensors model loading
 fastsafetensors >= 0.3.3
+# Mooncake transfer engine for KV cache offloading/transfer (ROCm build, published on PyPI)
+mooncake-transfer-engine-rocm >= 0.3.13