A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[UX] Move Dockerfile DeepGEMM install to tools/install_deepgemm.sh (#23360)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 22, 2025 at 22:52 UTC
f6818a92cbad18b7bac0b5cb424549ba941e11b9
2 files changed
+112
-24
docker/Dockerfile
+4
-24
@@ -432,31 +432,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \
432
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
433
434
# Install DeepGEMM from source
435
-ARG DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
435
ARG DEEPGEMM_GIT_REF="7b6b5563b9d4c1ae07ffbce7f78ad3ac9204827c"
437
-RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
438
- . /etc/environment
439
- CUDA_MAJOR="${CUDA_VERSION%%.*}"
440
- CUDA_MINOR="${CUDA_VERSION#${CUDA_MAJOR}.}"
441
- CUDA_MINOR="${CUDA_MINOR%%.*}"
442
- if [ "$CUDA_MAJOR" -ge 12 ] && [ "$CUDA_MINOR" -ge 8 ]; then
443
- git clone --recursive --shallow-submodules \
444
- ${DEEPGEMM_GIT_REPO} deepgemm
445
- echo "🏗️ Building DeepGEMM"
446
- pushd deepgemm
447
- git checkout ${DEEPGEMM_GIT_REF}
448
- # Build DeepGEMM
449
- # (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh)
450
- rm -rf build dist
451
- rm -rf *.egg-info
452
- python3 setup.py bdist_wheel
453
- uv pip install --system dist/*.whl
454
- popd
455
- rm -rf deepgemm
456
- else
457
- echo "Skipping DeepGEMM installation (requires CUDA 12.8+ but got ${CUDA_VERSION})"
458
- fi
459
-BASH
436
+COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
437
+RUN --mount=type=cache,target=/root/.cache/uv \
438
+ VLLM_DOCKER_BUILD_CONTEXT=1 /tmp/install_deepgemm.sh --cuda-version "${CUDA_VERSION}" --ref "${DEEPGEMM_GIT_REF}" \
439
+ && rm /tmp/install_deepgemm.sh
440
441
# Install EP kernels(pplx-kernels and DeepEP), NixL
442
COPY tools/ep_kernels/install_python_libraries.sh install_python_libraries.sh
tools/install_deepgemm.sh
new
+108
@@ -0,0 +1,108 @@
1
+#!/bin/bash
2
+# Script to install DeepGEMM from source
3
+# This script can be used both in Docker builds and by users locally
4
+
5
+set -e
6
+
7
+# Default values
8
+DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
9
+DEEPGEMM_GIT_REF="7b6b5563b9d4c1ae07ffbce7f78ad3ac9204827c"
10
+
11
+# Parse command line arguments
12
+while [[ $# -gt 0 ]]; do
13
+ case $1 in
14
+ --ref)
15
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
16
+ echo "Error: --ref requires an argument." >&2
17
+ exit 1
18
+ fi
19
+ DEEPGEMM_GIT_REF="$2"
20
+ shift 2
21
+ ;;
22
+ --cuda-version)
23
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
24
+ echo "Error: --cuda-version requires an argument." >&2
25
+ exit 1
26
+ fi
27
+ CUDA_VERSION="$2"
28
+ shift 2
29
+ ;;
30
+ -h|--help)
31
+ echo "Usage: $0 [OPTIONS]"
32
+ echo "Options:"
33
+ echo " --ref REF Git reference to checkout (default: $DEEPGEMM_GIT_REF)"
34
+ echo " --cuda-version VER CUDA version (auto-detected if not provided)"
35
+ echo " -h, --help Show this help message"
36
+ exit 0
37
+ ;;
38
+ *)
39
+ echo "Unknown option: $1" >&2
40
+ exit 1
41
+ ;;
42
+ esac
43
+done
44
+
45
+# Auto-detect CUDA version if not provided
46
+if [ -z "$CUDA_VERSION" ]; then
47
+ if command -v nvcc >/dev/null 2>&1; then
48
+ CUDA_VERSION=$(nvcc --version | grep "release" | sed -n 's/.*release \([0-9]\+\.[0-9]\+\).*/\1/p')
49
+ echo "Auto-detected CUDA version: $CUDA_VERSION"
50
+ else
51
+ echo "Warning: Could not auto-detect CUDA version. Please specify with --cuda-version"
52
+ exit 1
53
+ fi
54
+fi
55
+
56
+# Extract major and minor version numbers
57
+CUDA_MAJOR="${CUDA_VERSION%%.*}"
58
+CUDA_MINOR="${CUDA_VERSION#${CUDA_MAJOR}.}"
59
+CUDA_MINOR="${CUDA_MINOR%%.*}"
60
+
61
+echo "CUDA version: $CUDA_VERSION (major: $CUDA_MAJOR, minor: $CUDA_MINOR)"
62
+
63
+# Check CUDA version requirement
64
+if [ "$CUDA_MAJOR" -lt 12 ] || { [ "$CUDA_MAJOR" -eq 12 ] && [ "$CUDA_MINOR" -lt 8 ]; }; then
65
+ echo "Skipping DeepGEMM installation (requires CUDA 12.8+ but got ${CUDA_VERSION})"
66
+ exit 0
67
+fi
68
+
69
+echo "Installing DeepGEMM from source..."
70
+echo "Repository: $DEEPGEMM_GIT_REPO"
71
+echo "Reference: $DEEPGEMM_GIT_REF"
72
+
73
+# Create a temporary directory for the build
74
+INSTALL_DIR=$(mktemp -d)
75
+trap 'rm -rf "$INSTALL_DIR"' EXIT
76
+
77
+# Clone the repository
78
+git clone --recursive --shallow-submodules "$DEEPGEMM_GIT_REPO" "$INSTALL_DIR/deepgemm"
79
+
80
+echo "🏗️ Building DeepGEMM"
81
+pushd "$INSTALL_DIR/deepgemm"
82
+
83
+# Checkout the specific reference
84
+git checkout "$DEEPGEMM_GIT_REF"
85
+
86
+# Build DeepGEMM
87
+# (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh)
88
+rm -rf build dist
89
+rm -rf *.egg-info
90
+python3 setup.py bdist_wheel
91
+
92
+# Install the wheel
93
+if command -v uv >/dev/null 2>&1; then
94
+ echo "Installing DeepGEMM wheel using uv..."
95
+ # Use --system in Docker contexts, respect user's environment otherwise
96
+ if [ -n "$VLLM_DOCKER_BUILD_CONTEXT" ]; then
97
+ uv pip install --system dist/*.whl
98
+ else
99
+ uv pip install dist/*.whl
100
+ fi
101
+else
102
+ echo "Installing DeepGEMM wheel using pip..."
103
+ python3 -m pip install dist/*.whl
104
+fi
105
+
106
+popd
107
+
108
+echo "✅ DeepGEMM installation completed successfully"
\ No newline at end of file