| 1 | #!/bin/bash |
| 2 | set -eux |
| 3 | |
| 4 | python_executable=python$1 |
| 5 | cuda_home=/usr/local/cuda-$2 |
| 6 | |
| 7 | # Update paths |
| 8 | PATH=${cuda_home}/bin:$PATH |
| 9 | LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH |
| 10 | |
| 11 | # Install requirements |
| 12 | if [ "$(echo "$2" | cut -d. -f1)" = "12" ]; then |
| 13 | sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt |
| 14 | fi |
| 15 | $python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt |
| 16 | |
| 17 | # Limit the number of parallel jobs to avoid OOM |
| 18 | export MAX_JOBS=1 |
| 19 | # Make sure release wheels are built for the following architectures |
| 20 | # Do not add +PTX here: vLLM filters torch's top-level PTX flag when it |
| 21 | # converts global gencode flags into per-kernel arch lists. If a specific |
| 22 | # kernel needs PTX, add +PTX to that kernel's CMake arch list instead. |
| 23 | export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0" |
| 24 | |
| 25 | bash tools/check_repo.sh |
| 26 | |
| 27 | # Build |
| 28 | $python_executable setup.py bdist_wheel --dist-dir=dist |