File size: 3,908 Bytes
eb8a8e9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | #!/bin/bash
# Script to build and/or install DeepGEMM from source
# Default: build and install immediately
# Optional: build wheels to a directory for later installation (useful in multi-stage builds)
set -e
# Default values
DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
DEEPGEMM_GIT_REF="0f5f2662027f0db05d4e3f6a94e56e2d8fc45c51"
WHEEL_DIR=""
# 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
;;
--wheel-dir)
if [[ -z "$2" || "$2" =~ ^- ]]; then
echo "Error: --wheel-dir requires a directory path." >&2
exit 1
fi
WHEEL_DIR="$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 " --wheel-dir PATH If set, build wheel into PATH but do not install"
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 build/installation (requires CUDA 12.8+ but got ${CUDA_VERSION})"
exit 0
fi
echo "Preparing DeepGEMM build..."
echo "Repository: $DEEPGEMM_GIT_REPO"
echo "Reference: $DEEPGEMM_GIT_REF"
# Create a temporary directory for the build
# INSTALL_DIR=$(mktemp -d)
INSTALL_DIR="."
trap 'rm -rf "$INSTALL_DIR"' EXIT
# Clone the repository
git clone --recursive --shallow-submodules "$DEEPGEMM_GIT_REPO" "$INSTALL_DIR/deepgemm"
pushd "$INSTALL_DIR/deepgemm"
# Checkout the specific reference
git checkout "$DEEPGEMM_GIT_REF"
# Clean previous build artifacts
# (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh)
rm -rf build dist *.egg-info
# Build wheel
echo "🏗️ Building DeepGEMM wheel..."
python3 setup.py bdist_wheel
# If --wheel-dir was specified, copy wheels there and exit
if [ -n "$WHEEL_DIR" ]; then
mkdir -p "$WHEEL_DIR"
cp dist/*.whl "$WHEEL_DIR"/
echo "✅ Wheel built and copied to $WHEEL_DIR"
popd
exit 0
fi
# Default behaviour: install built wheel
if command -v uv >/dev/null 2>&1; then
echo "Installing DeepGEMM wheel using uv..."
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"
|