From f6818a92cbad18b7bac0b5cb424549ba941e11b9 Mon Sep 17 00:00:00 2001
From: Michael Goin <mgoin64@gmail.com>
Date: Fri, 22 Aug 2025 22:52:50 -0400
Subject: [PATCH] [UX] Move Dockerfile DeepGEMM install to
 tools/install_deepgemm.sh (#23360)

Signed-off-by: mgoin <mgoin64@gmail.com>
---
 docker/Dockerfile         |  28 ++--------
 tools/install_deepgemm.sh | 108 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 112 insertions(+), 24 deletions(-)
 create mode 100755 tools/install_deepgemm.sh

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 977dd7a3d31d8..839ac501dbaf0 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -432,31 +432,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 
 # Install DeepGEMM from source
-ARG DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
 ARG DEEPGEMM_GIT_REF="7b6b5563b9d4c1ae07ffbce7f78ad3ac9204827c"
-RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
-  . /etc/environment
-    CUDA_MAJOR="${CUDA_VERSION%%.*}"
-    CUDA_MINOR="${CUDA_VERSION#${CUDA_MAJOR}.}"
-    CUDA_MINOR="${CUDA_MINOR%%.*}"
-    if [ "$CUDA_MAJOR" -ge 12 ] && [ "$CUDA_MINOR" -ge 8 ]; then
-        git clone --recursive --shallow-submodules \
-            ${DEEPGEMM_GIT_REPO} deepgemm
-        echo "🏗️  Building DeepGEMM"
-        pushd deepgemm
-            git checkout ${DEEPGEMM_GIT_REF}
-            # Build DeepGEMM
-            # (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh)
-            rm -rf build dist
-            rm -rf *.egg-info
-            python3 setup.py bdist_wheel
-            uv pip install --system dist/*.whl
-        popd
-        rm -rf deepgemm
-    else
-        echo "Skipping DeepGEMM installation (requires CUDA 12.8+ but got ${CUDA_VERSION})"
-    fi
-BASH
+COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
+RUN --mount=type=cache,target=/root/.cache/uv \
+    VLLM_DOCKER_BUILD_CONTEXT=1 /tmp/install_deepgemm.sh --cuda-version "${CUDA_VERSION}" --ref "${DEEPGEMM_GIT_REF}" \
+    && rm /tmp/install_deepgemm.sh
 
 # Install EP kernels(pplx-kernels and DeepEP), NixL
 COPY tools/ep_kernels/install_python_libraries.sh install_python_libraries.sh
diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh
new file mode 100755
index 0000000000000..33849581d2c0e
--- /dev/null
+++ b/tools/install_deepgemm.sh
@@ -0,0 +1,108 @@
+#!/bin/bash
+# Script to install DeepGEMM from source
+# This script can be used both in Docker builds and by users locally
+
+set -e
+
+# Default values
+DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
+DEEPGEMM_GIT_REF="7b6b5563b9d4c1ae07ffbce7f78ad3ac9204827c"
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --ref)
+            if [[ -z "$2" || "$2" =~ ^- ]]; then
+                echo "Error: --ref requires an argument." >&2
+                exit 1
+            fi
+            DEEPGEMM_GIT_REF="$2"
+            shift 2
+            ;;
+        --cuda-version)
+            if [[ -z "$2" || "$2" =~ ^- ]]; then
+                echo "Error: --cuda-version requires an argument." >&2
+                exit 1
+            fi
+            CUDA_VERSION="$2"
+            shift 2
+            ;;
+        -h|--help)
+            echo "Usage: $0 [OPTIONS]"
+            echo "Options:"
+            echo "  --ref REF          Git reference to checkout (default: $DEEPGEMM_GIT_REF)"
+            echo "  --cuda-version VER CUDA version (auto-detected if not provided)"
+            echo "  -h, --help         Show this help message"
+            exit 0
+            ;;
+        *)
+            echo "Unknown option: $1" >&2
+            exit 1
+            ;;
+    esac
+done
+
+# Auto-detect CUDA version if not provided
+if [ -z "$CUDA_VERSION" ]; then
+    if command -v nvcc >/dev/null 2>&1; then
+        CUDA_VERSION=$(nvcc --version | grep "release" | sed -n 's/.*release \([0-9]\+\.[0-9]\+\).*/\1/p')
+        echo "Auto-detected CUDA version: $CUDA_VERSION"
+    else
+        echo "Warning: Could not auto-detect CUDA version. Please specify with --cuda-version"
+        exit 1
+    fi
+fi
+
+# Extract major and minor version numbers
+CUDA_MAJOR="${CUDA_VERSION%%.*}"
+CUDA_MINOR="${CUDA_VERSION#${CUDA_MAJOR}.}"
+CUDA_MINOR="${CUDA_MINOR%%.*}"
+
+echo "CUDA version: $CUDA_VERSION (major: $CUDA_MAJOR, minor: $CUDA_MINOR)"
+
+# Check CUDA version requirement
+if [ "$CUDA_MAJOR" -lt 12 ] || { [ "$CUDA_MAJOR" -eq 12 ] && [ "$CUDA_MINOR" -lt 8 ]; }; then
+    echo "Skipping DeepGEMM installation (requires CUDA 12.8+ but got ${CUDA_VERSION})"
+    exit 0
+fi
+
+echo "Installing DeepGEMM from source..."
+echo "Repository: $DEEPGEMM_GIT_REPO"
+echo "Reference: $DEEPGEMM_GIT_REF"
+
+# Create a temporary directory for the build
+INSTALL_DIR=$(mktemp -d)
+trap 'rm -rf "$INSTALL_DIR"' EXIT
+
+# Clone the repository
+git clone --recursive --shallow-submodules "$DEEPGEMM_GIT_REPO" "$INSTALL_DIR/deepgemm"
+
+echo "🏗️  Building DeepGEMM"
+pushd "$INSTALL_DIR/deepgemm"
+
+# Checkout the specific reference
+git checkout "$DEEPGEMM_GIT_REF"
+
+# Build DeepGEMM
+# (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh)
+rm -rf build dist
+rm -rf *.egg-info
+python3 setup.py bdist_wheel
+
+# Install the wheel
+if command -v uv >/dev/null 2>&1; then
+    echo "Installing DeepGEMM wheel using uv..."
+    # Use --system in Docker contexts, respect user's environment otherwise
+    if [ -n "$VLLM_DOCKER_BUILD_CONTEXT" ]; then
+        uv pip install --system dist/*.whl
+    else
+        uv pip install dist/*.whl
+    fi
+else
+    echo "Installing DeepGEMM wheel using pip..."
+    python3 -m pip install dist/*.whl
+fi
+
+popd
+
+echo "✅ DeepGEMM installation completed successfully"
\ No newline at end of file