[ROCm][AMD][Installation] Add mooncake build to rocm base image

.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 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 mooncake_source build_mooncake package_mooncake export_mooncake 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,9 +1684,12 @@ 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() {
@@ -2268,7 +2271,7 @@ extract_dependency_pins() {
         return 0
     fi
 
-    for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
+    for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH MOONCAKE_BRANCH; do
         if [[ -n "${!var:-}" ]]; then
             echo "Using provided ${var}: ${!var}"
             continue
@@ -2292,16 +2295,19 @@ 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"
     nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
     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")
@@ -2335,6 +2341,19 @@ 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() {
@@ -2383,6 +2402,13 @@ 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
 }
 
@@ -2400,13 +2426,14 @@ 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; do
+            for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci mooncake-rocm-ci; do
                 if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
                     add_dependency_cache_target "${target}"
                 fi
@@ -2456,6 +2483,16 @@ 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,6 +7,10 @@ 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.
@@ -643,6 +647,82 @@ 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
@@ -873,6 +953,15 @@ 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
@@ -952,10 +1041,21 @@ 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
@@ -984,7 +1084,9 @@ 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 "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
 
 CMD ["/bin/bash"]
 

docker/ci-rocm.hcl

@@ -91,6 +91,10 @@ variable "DEEPEP_BRANCH" {
   default = ""
 }
 
+variable "MOONCAKE_BRANCH" {
+  default = ""
+}
+
 variable "NIXL_CACHE_KEY" {
   default = ""
 }
@@ -103,6 +107,10 @@ 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.
@@ -247,6 +255,7 @@ 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}" : ""),
   ])
 }
 
@@ -271,6 +280,13 @@ 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" {
@@ -401,6 +417,14 @@ 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
@@ -431,5 +455,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", "ci-base-rocm-ci"]
+  targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "mooncake-rocm-ci", "ci-base-rocm-ci"]
 }