feat: support sgl-kernel PyPI (#2433)

Co-authored-by: Zhangyi <1109276519@qq.com>
This commit is contained in:
Yineng Zhang
2024-12-11 06:06:19 +08:00
committed by GitHub
parent 2b340adfb1
commit 56fcd8e8a5
4 changed files with 97 additions and 10 deletions

View File

@@ -1,13 +1,16 @@
#!/bin/bash
set -ex
PYTHON_VERSION=$1
CUDA_VERSION=$2
PYTHON_ROOT_PATH=/opt/python/cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}
docker run --rm -it \
docker run --rm \
-v "$(pwd)":/sgl-kernel \
pytorch/manylinux-builder:cuda12.1 \
pytorch/manylinux-builder:cuda${CUDA_VERSION} \
bash -c "
pip install --no-cache-dir torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121 && \
${PYTHON_ROOT_PATH}/bin/pip install --no-cache-dir torch==2.4.0 --index-url https://download.pytorch.org/whl/cu${CUDA_VERSION//.} && \
export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0+PTX' && \
export CUDA_VERSION=${CUDA_VERSION} && \
cd /sgl-kernel && \
python setup.py bdist_wheel
${PYTHON_ROOT_PATH}/bin/python setup.py bdist_wheel
"