| 1 | #!/bin/bash |
| 2 | # Script to build and/or install DeepGEMM from source |
| 3 | # Default: build and install immediately |
| 4 | # Optional: build wheels to a directory for later installation (useful in multi-stage builds) |
| 5 | set -e |
| 6 | |
| 7 | # Default values |
| 8 | # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake |
| 9 | DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" |
| 10 | # NOTE: This targets the vLLM fork's dev branch, which carries the sm120 |
| 11 | # and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE |
| 12 | # activations, plus the CUDA 12.x layout header fix from vllm-project/DeepGEMM#12. |
| 13 | DEEPGEMM_GIT_REF="e1f418c2a4f20818221f6b0e578b4c2f634d4c3f" |
| 14 | WHEEL_DIR="" |
| 15 | |
| 16 | # Parse command line arguments |
| 17 | while [[ $# -gt 0 ]]; do |
| 18 | case $1 in |
| 19 | --ref) |
| 20 | if [[ -z "$2" || "$2" =~ ^- ]]; then |
| 21 | echo "Error: --ref requires an argument." >&2 |
| 22 | exit 1 |
| 23 | fi |
| 24 | DEEPGEMM_GIT_REF="$2" |
| 25 | shift 2 |
| 26 | ;; |
| 27 | --cuda-version) |
| 28 | if [[ -z "$2" || "$2" =~ ^- ]]; then |
| 29 | echo "Error: --cuda-version requires an argument." >&2 |
| 30 | exit 1 |
| 31 | fi |
| 32 | CUDA_VERSION="$2" |
| 33 | shift 2 |
| 34 | ;; |
| 35 | --wheel-dir) |
| 36 | if [[ -z "$2" || "$2" =~ ^- ]]; then |
| 37 | echo "Error: --wheel-dir requires a directory path." >&2 |
| 38 | exit 1 |
| 39 | fi |
| 40 | WHEEL_DIR="$2" |
| 41 | shift 2 |
| 42 | ;; |
| 43 | -h|--help) |
| 44 | echo "Usage: $0 [OPTIONS]" |
| 45 | echo "Options:" |
| 46 | echo " --ref REF Git reference to checkout (default: $DEEPGEMM_GIT_REF)" |
| 47 | echo " --cuda-version VER CUDA version (auto-detected if not provided)" |
| 48 | echo " --wheel-dir PATH If set, build wheel into PATH but do not install" |
| 49 | echo " -h, --help Show this help message" |
| 50 | exit 0 |
| 51 | ;; |
| 52 | *) |
| 53 | echo "Unknown option: $1" >&2 |
| 54 | exit 1 |
| 55 | ;; |
| 56 | esac |
| 57 | done |
| 58 | |
| 59 | # Auto-detect CUDA version if not provided |
| 60 | if [ -z "$CUDA_VERSION" ]; then |
| 61 | if command -v nvcc >/dev/null 2>&1; then |
| 62 | CUDA_VERSION=$(nvcc --version | grep "release" | sed -n 's/.*release \([0-9]\+\.[0-9]\+\).*/\1/p') |
| 63 | echo "Auto-detected CUDA version: $CUDA_VERSION" |
| 64 | else |
| 65 | echo "Warning: Could not auto-detect CUDA version. Please specify with --cuda-version" |
| 66 | exit 1 |
| 67 | fi |
| 68 | fi |
| 69 | |
| 70 | # Extract major and minor version numbers |
| 71 | CUDA_MAJOR="${CUDA_VERSION%%.*}" |
| 72 | CUDA_MINOR="${CUDA_VERSION#"${CUDA_MAJOR}".}" |
| 73 | CUDA_MINOR="${CUDA_MINOR%%.*}" |
| 74 | echo "CUDA version: $CUDA_VERSION (major: $CUDA_MAJOR, minor: $CUDA_MINOR)" |
| 75 | |
| 76 | # Check CUDA version requirement |
| 77 | if [ "$CUDA_MAJOR" -lt 12 ] || { [ "$CUDA_MAJOR" -eq 12 ] && [ "$CUDA_MINOR" -lt 8 ]; }; then |
| 78 | echo "Skipping DeepGEMM build/installation (requires CUDA 12.8+ but got ${CUDA_VERSION})" |
| 79 | exit 0 |
| 80 | fi |
| 81 | |
| 82 | echo "Preparing DeepGEMM build..." |
| 83 | echo "Repository: $DEEPGEMM_GIT_REPO" |
| 84 | echo "Reference: $DEEPGEMM_GIT_REF" |
| 85 | |
| 86 | # Create a temporary directory for the build |
| 87 | INSTALL_DIR=$(mktemp -d) |
| 88 | trap 'rm -rf "$INSTALL_DIR"' EXIT |
| 89 | |
| 90 | # Clone the repository |
| 91 | git clone --recursive --shallow-submodules "$DEEPGEMM_GIT_REPO" "$INSTALL_DIR/deepgemm" |
| 92 | pushd "$INSTALL_DIR/deepgemm" |
| 93 | |
| 94 | # Checkout the specific reference |
| 95 | git checkout "$DEEPGEMM_GIT_REF" |
| 96 | |
| 97 | # Clean previous build artifacts |
| 98 | # (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh) |
| 99 | rm -rf -- build dist *.egg-info 2>/dev/null || true |
| 100 | |
| 101 | # Build wheel |
| 102 | echo "🏗️ Building DeepGEMM wheel..." |
| 103 | python3 setup.py bdist_wheel |
| 104 | |
| 105 | # If --wheel-dir was specified, copy wheels there and exit |
| 106 | if [ -n "$WHEEL_DIR" ]; then |
| 107 | mkdir -p "$WHEEL_DIR" |
| 108 | cp dist/*.whl "$WHEEL_DIR"/ |
| 109 | echo "✅ Wheel built and copied to $WHEEL_DIR" |
| 110 | popd |
| 111 | exit 0 |
| 112 | fi |
| 113 | |
| 114 | # Default behaviour: install built wheel |
| 115 | if command -v uv >/dev/null 2>&1; then |
| 116 | echo "Installing DeepGEMM wheel using uv..." |
| 117 | if [ -n "$VLLM_DOCKER_BUILD_CONTEXT" ]; then |
| 118 | uv pip install --system dist/*.whl |
| 119 | else |
| 120 | uv pip install dist/*.whl |
| 121 | fi |
| 122 | else |
| 123 | echo "Installing DeepGEMM wheel using pip..." |
| 124 | python3 -m pip install dist/*.whl |
| 125 | fi |
| 126 | |
| 127 | popd |
| 128 | echo "✅ DeepGEMM installation completed successfully" |