From 71babdef510ef3d95ea829ec61e5f7c3f9ae7232 Mon Sep 17 00:00:00 2001
From: ZhenshengWu <49758440+ZhenshengWu@users.noreply.github.com>
Date: Sun, 1 Feb 2026 15:41:21 +0800
Subject: [PATCH] Fix CUDA 12 dependency when importing Mooncake in official
 CUDA 13.x image (#17540)

Co-authored-by: wuzhensheng01 <wuzhensheng01@baidu.com>
---
 docker/Dockerfile | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 366efa327..f76be5fae 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -19,6 +19,11 @@ ARG PIP_DEFAULT_INDEX
 ARG UBUNTU_MIRROR
 ARG GITHUB_ARTIFACTORY=github.com
 ARG INSTALL_FLASHINFER_JIT_CACHE=0
+ARG FLASHINFER_VERSION=0.6.1
+ARG MOONCAKE_VERSION=0.3.8.post1
+#if need other arg please add in MOONCAKE_COMPILE_ARG
+ARG MOONCAKE_COMPILE_ARG="-DUSE_HTTP=ON -DUSE_MNNVL=ON -DUSE_CUDA=ON -DWITH_EP=ON"
+
 ARG FLASHINFER_VERSION=0.6.2
 
 ENV DEBIAN_FRONTEND=noninteractive \
@@ -269,11 +274,27 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     fi && \
     TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" MAX_JOBS=${BUILD_AND_DOWNLOAD_PARALLEL} pip install --no-build-isolation .
 
+# Install Mooncake
+RUN --mount=type=cache,target=/root/.cache/pip \
+    CUDA_MAJOR="${CUDA_VERSION%%.*}" && \
+    if [ "$CUDA_MAJOR" -ge 13 ]; then \
+        echo "CUDA >= 13, installing mooncake-transfer-engine from source code"; \
+        git clone --branch v${MOONCAKE_VERSION} --depth 1 https://github.com/kvcache-ai/Mooncake.git && \
+        cd Mooncake && \
+        bash dependencies.sh && \
+        mkdir -p build && \
+        cd build && \
+        cmake .. ${MOONCAKE_COMPILE_ARG} && \
+        make -j$(nproc) && \
+        make install; \
+    else \
+        echo "CUDA < 13, installing mooncake-transfer-engine from pip"; \
+        python3 -m pip install mooncake-transfer-engine==${MOONCAKE_VERSION}; \
+    fi
 # Install essential Python packages
 RUN --mount=type=cache,target=/root/.cache/pip \
     python3 -m pip install \
     datamodel_code_generator \
-    mooncake-transfer-engine==0.3.8.post1 \
     pre-commit \
     pytest \
     black \