[Docker][KVConnector] Install mooncake from official wheels instead of a custom build (#51067)

Signed-off-by: Zhewen Li <[email protected]>
Co-authored-by: Zhewen Li <[email protected]>
Co-authored-by: Claude Opus 5 <[email protected]>
This commit is contained in:
Zhewen Li
2026-08-05 11:58:44 -07:00
committed by GitHub
co-authored by Zhewen Li Claude Opus 5
parent 613411a90c
commit 14e57ad47d
3 changed files with 26 additions and 47 deletions
-23
View File
@@ -16,13 +16,6 @@ env:
# please use CUDA 13 wheels or compile yourself on these new devices
CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
# pre-built mooncake wheels
# the manylinux_2_35 wheel has compatibility issue on Ubuntu 24.04
# so we use different wheels for the time being
MOONCAKE_WHEEL_AARCH64_2_35: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_aarch64.whl"
MOONCAKE_WHEEL_AARCH64_2_39: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_39_aarch64.whl"
MOONCAKE_WHEEL_X86_64: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_x86_64.whl"
steps:
- input: "Provide Release version here"
@@ -219,8 +212,6 @@ steps:
--build-arg CUDA_VERSION=13.0.2 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 \
--target vllm-openai \
--progress plain \
@@ -247,8 +238,6 @@ steps:
--build-arg CUDA_VERSION=13.0.2 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \
--target vllm-openai \
--progress plain \
@@ -273,8 +262,6 @@ steps:
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
@@ -301,8 +288,6 @@ steps:
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
@@ -327,8 +312,6 @@ steps:
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 \
--target vllm-openai \
--progress plain \
@@ -356,8 +339,6 @@ steps:
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \
--target vllm-openai \
--progress plain \
@@ -384,8 +365,6 @@ steps:
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
@@ -413,8 +392,6 @@ steps:
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
+17 -23
View File
@@ -975,32 +975,26 @@ RUN --mount=type=cache,target=/opt/uv/cache \
# clean up -dev packages, keep runtime libraries
rm -rf /var/lib/apt/lists/* \
); \
# Force-reinstall the matching CUDA wheel so the correct nixl_ep_cpp.so is installed.
uv pip install --system --force-reinstall --no-deps nixl-cu${CUDA_MAJOR}; \
fi
# Optional override: install mooncake-transfer-engine from a URL instead of the
# PyPI release pulled in above. Use this for wheels built with non-default CMake
# flags (e.g. `STORE_USE_ETCD=ON` for master HA). The URL's manylinux glibc
# floor must be <= the FINAL_BASE_IMAGE's glibc.
ARG MOONCAKE_WHEEL_AARCH64
ARG MOONCAKE_WHEEL_X86_64
RUN if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
WHEEL="${MOONCAKE_WHEEL_AARCH64}"; \
else \
WHEEL="${MOONCAKE_WHEEL_X86_64}"; \
fi && \
if [ -n "${WHEEL}" ]; then \
uv pip install --system "${WHEEL}" && \
CUDA_MAJOR="${CUDA_VERSION%%.*}" && \
if [ ! -f /usr/local/cuda/lib64/libcudart.so ] && \
[ -f "/usr/local/cuda/lib64/libcudart.so.${CUDA_MAJOR}" ]; then \
ln -s "libcudart.so.${CUDA_MAJOR}" /usr/local/cuda/lib64/libcudart.so; \
fi; \
# Force-reinstall the matching CUDA wheel so the correct nixl_ep_cpp.so is
# installed. Chained with && through the mooncake swap below: this RUN has no
# `set -e`, so its status is that of the last command, and a trailing no-op
# would otherwise mask a failure here.
uv pip install --system --force-reinstall --no-deps nixl-cu${CUDA_MAJOR} && \
# The default mooncake wheel links libcudart.so.12, so on CUDA 13 swap in the
# cuda13 variant. Uninstall first; both ship the `mooncake` package.
# Mirrors .buildkite/scripts/install-kv-connectors.sh.
MOONCAKE_VERSION=$(python3 -c "import importlib.metadata as m; print(m.version('mooncake-transfer-engine'))" 2>/dev/null || true) && \
if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then \
uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true; \
uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \
fi; \
fi
# Mooncake registers GPU memory for RDMA via nvidia-peermem by default. Hosts
# without that module loaded (e.g. GB200) must pass WITH_NVIDIA_PEERMEM=0 at run
# time; see docs/features/mooncake_connector_usage.md. Deliberately not an ENV
# here so the image does not override mooncake's own default.
ENV VLLM_USAGE_SOURCE production-docker-image
ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \
VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \
+9 -1
View File
@@ -10,7 +10,9 @@ For more details about Mooncake, please refer to [Mooncake project](https://gith
### Installation
Install mooncake through pip: `uv pip install mooncake-transfer-engine`.
Install mooncake through pip: `uv pip install mooncake-transfer-engine-cuda13`.
vLLM defaults to CUDA 13. On a CUDA 12 environment install `mooncake-transfer-engine` instead — the two are the same release built against different CUDA majors, and the wrong one fails to import with `libcudart.so.<major>: cannot open shared object file`.
Refer to [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for more installation instructions
@@ -44,6 +46,12 @@ Now you can send requests to the proxy server through port 8000.
- For headless instances, must be the same as the master instance
- Each instance needs a unique port on its host; using the same port number across different hosts is fine
- `WITH_NVIDIA_PEERMEM`: Selects how mooncake registers GPU memory for RDMA. Read by mooncake, not vLLM.
- Default: 1, which uses `ibv_reg_mr()` and requires the `nvidia-peermem` kernel module to be loaded
- Set to 0 to use the DMA-BUF path, which does not need that module. Required on hosts where `nvidia-peermem` is not loaded, such as GB200
- With the container image, pass it at run time: `docker run -e WITH_NVIDIA_PEERMEM=0 ...`
- Symptom when left unset on such a host: `Failed to register memory <addr>: Bad address [14]` from `rdma_context.cpp`, and KV transfers fail
- `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing the prefillers KV cache for a particular request. (Optional)
- Default: 480
- If a request is aborted and the decoder has not yet notified the prefiller, the prefill instance will release its KV-cache blocks after this timeout to avoid holding them indefinitely.