[Docker][KVConnector] Install mooncake from official wheels instead of a custom build (#51067)
Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Zhewen Li committed
Aug 5, 2026 at 11:58 UTC
14e57ad47dec65059f4f04700439a8465ea83a87
3 files changed
+26
-47
.buildkite/release-pipeline.yaml
-23
@@ -16,13 +16,6 @@ env:
16
# please use CUDA 13 wheels or compile yourself on these new devices
17
CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
18
CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
19
-
20
- # pre-built mooncake wheels
21
- # the manylinux_2_35 wheel has compatibility issue on Ubuntu 24.04
22
- # so we use different wheels for the time being
23
- 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"
24
- 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"
25
- 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"
19
20
steps:
21
- input: "Provide Release version here"
@@ -219,8 +212,6 @@ steps:
212
--build-arg CUDA_VERSION=13.0.2 \
213
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
214
--build-arg INSTALL_KV_CONNECTORS=true \
222
- --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
223
- --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
215
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 \
216
--target vllm-openai \
217
--progress plain \
@@ -247,8 +238,6 @@ steps:
238
--build-arg CUDA_VERSION=13.0.2 \
239
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
240
--build-arg INSTALL_KV_CONNECTORS=true \
250
- --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
251
- --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
241
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \
242
--target vllm-openai \
243
--progress plain \
@@ -273,8 +262,6 @@ steps:
262
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 \
263
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
264
--build-arg INSTALL_KV_CONNECTORS=true \
276
- --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
277
- --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
265
--target vllm-openai \
266
--progress plain \
267
-f docker/Dockerfile .
@@ -301,8 +288,6 @@ steps:
288
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 \
289
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
290
--build-arg INSTALL_KV_CONNECTORS=true \
304
- --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
305
- --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
291
--target vllm-openai \
292
--progress plain \
293
-f docker/Dockerfile .
@@ -327,8 +312,6 @@ steps:
312
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
313
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
314
--build-arg INSTALL_KV_CONNECTORS=true \
330
- --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
331
- --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
315
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 \
316
--target vllm-openai \
317
--progress plain \
@@ -356,8 +339,6 @@ steps:
339
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
340
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
341
--build-arg INSTALL_KV_CONNECTORS=true \
359
- --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
360
- --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
342
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \
343
--target vllm-openai \
344
--progress plain \
@@ -384,8 +365,6 @@ steps:
365
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
366
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
367
--build-arg INSTALL_KV_CONNECTORS=true \
387
- --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
388
- --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
368
--target vllm-openai \
369
--progress plain \
370
-f docker/Dockerfile .
@@ -413,8 +392,6 @@ steps:
392
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
393
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
394
--build-arg INSTALL_KV_CONNECTORS=true \
416
- --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
417
- --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
395
--target vllm-openai \
396
--progress plain \
397
-f docker/Dockerfile .
docker/Dockerfile
+17
-23
@@ -975,32 +975,26 @@ RUN --mount=type=cache,target=/opt/uv/cache \
975
# clean up -dev packages, keep runtime libraries
976
rm -rf /var/lib/apt/lists/* \
977
); \
978
- # Force-reinstall the matching CUDA wheel so the correct nixl_ep_cpp.so is installed.
979
- uv pip install --system --force-reinstall --no-deps nixl-cu${CUDA_MAJOR}; \
980
- fi
981
-
982
-# Optional override: install mooncake-transfer-engine from a URL instead of the
983
-# PyPI release pulled in above. Use this for wheels built with non-default CMake
984
-# flags (e.g. `STORE_USE_ETCD=ON` for master HA). The URL's manylinux glibc
985
-# floor must be <= the FINAL_BASE_IMAGE's glibc.
986
-ARG MOONCAKE_WHEEL_AARCH64
987
-ARG MOONCAKE_WHEEL_X86_64
988
-RUN if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
989
- if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
990
- WHEEL="${MOONCAKE_WHEEL_AARCH64}"; \
991
- else \
992
- WHEEL="${MOONCAKE_WHEEL_X86_64}"; \
993
- fi && \
994
- if [ -n "${WHEEL}" ]; then \
995
- uv pip install --system "${WHEEL}" && \
996
- CUDA_MAJOR="${CUDA_VERSION%%.*}" && \
997
- if [ ! -f /usr/local/cuda/lib64/libcudart.so ] && \
998
- [ -f "/usr/local/cuda/lib64/libcudart.so.${CUDA_MAJOR}" ]; then \
999
- ln -s "libcudart.so.${CUDA_MAJOR}" /usr/local/cuda/lib64/libcudart.so; \
1000
- fi; \
978
+ # Force-reinstall the matching CUDA wheel so the correct nixl_ep_cpp.so is
979
+ # installed. Chained with && through the mooncake swap below: this RUN has no
980
+ # `set -e`, so its status is that of the last command, and a trailing no-op
981
+ # would otherwise mask a failure here.
982
+ uv pip install --system --force-reinstall --no-deps nixl-cu${CUDA_MAJOR} && \
983
+ # The default mooncake wheel links libcudart.so.12, so on CUDA 13 swap in the
984
+ # cuda13 variant. Uninstall first; both ship the `mooncake` package.
985
+ # Mirrors .buildkite/scripts/install-kv-connectors.sh.
986
+ MOONCAKE_VERSION=$(python3 -c "import importlib.metadata as m; print(m.version('mooncake-transfer-engine'))" 2>/dev/null || true) && \
987
+ if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then \
988
+ uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true; \
989
+ uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \
990
fi; \
991
fi
992
993
+# Mooncake registers GPU memory for RDMA via nvidia-peermem by default. Hosts
994
+# without that module loaded (e.g. GB200) must pass WITH_NVIDIA_PEERMEM=0 at run
995
+# time; see docs/features/mooncake_connector_usage.md. Deliberately not an ENV
996
+# here so the image does not override mooncake's own default.
997
+
998
ENV VLLM_USAGE_SOURCE production-docker-image
999
ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \
1000
VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \
docs/features/mooncake_connector_usage.md
+9
-1
@@ -10,7 +10,9 @@ For more details about Mooncake, please refer to [Mooncake project](https://gith
10
11
### Installation
12
13
-Install mooncake through pip: `uv pip install mooncake-transfer-engine`.
13
+Install mooncake through pip: `uv pip install mooncake-transfer-engine-cuda13`.
14
+
15
+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`.
16
17
Refer to [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for more installation instructions
18
@@ -44,6 +46,12 @@ Now you can send requests to the proxy server through port 8000.
46
- For headless instances, must be the same as the master instance
47
- Each instance needs a unique port on its host; using the same port number across different hosts is fine
48
49
+- `WITH_NVIDIA_PEERMEM`: Selects how mooncake registers GPU memory for RDMA. Read by mooncake, not vLLM.
50
+ - Default: 1, which uses `ibv_reg_mr()` and requires the `nvidia-peermem` kernel module to be loaded
51
+ - 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
52
+ - With the container image, pass it at run time: `docker run -e WITH_NVIDIA_PEERMEM=0 ...`
53
+ - Symptom when left unset on such a host: `Failed to register memory <addr>: Bad address [14]` from `rdma_context.cpp`, and KV transfers fail
54
+
55
- `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing the prefiller’s KV cache for a particular request. (Optional)
56
- Default: 480
57
- 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.