Fix sgl-kernel build on aarch64 platforms (#6062)
This commit is contained in:
@@ -7,12 +7,23 @@ PYTHON_ROOT_PATH=/opt/python/cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}
|
|||||||
|
|
||||||
ARCH=$(uname -i)
|
ARCH=$(uname -i)
|
||||||
echo "ARCH: $ARCH"
|
echo "ARCH: $ARCH"
|
||||||
|
if [ ${ARCH} = "aarch64" ]; then
|
||||||
|
LIBCUDA_ARCH="sbsa"
|
||||||
|
BUILDER_NAME="pytorch/manylinuxaarch64-builder"
|
||||||
|
else
|
||||||
|
LIBCUDA_ARCH=${ARCH}
|
||||||
|
if [ ${CUDA_VERSION} = "12.8" ]; then
|
||||||
|
BUILDER_NAME="pytorch/manylinux2_28-builder"
|
||||||
|
else
|
||||||
|
BUILDER_NAME="pytorch/manylinux-builder"
|
||||||
|
fi
|
||||||
|
fi
|
||||||
|
|
||||||
if [ ${CUDA_VERSION} = "12.8" ]; then
|
if [ ${CUDA_VERSION} = "12.8" ]; then
|
||||||
DOCKER_IMAGE="pytorch/manylinux2_28-builder:cuda${CUDA_VERSION}"
|
DOCKER_IMAGE="${BUILDER_NAME}:cuda${CUDA_VERSION}"
|
||||||
TORCH_INSTALL="pip install --no-cache-dir --pre torch --index-url https://download.pytorch.org/whl/nightly/cu${CUDA_VERSION//.}"
|
TORCH_INSTALL="pip install --no-cache-dir --pre torch --index-url https://download.pytorch.org/whl/nightly/cu${CUDA_VERSION//.}"
|
||||||
else
|
else
|
||||||
DOCKER_IMAGE="pytorch/manylinux-builder:cuda${CUDA_VERSION}"
|
DOCKER_IMAGE="${BUILDER_NAME}:cuda${CUDA_VERSION}"
|
||||||
TORCH_INSTALL="pip install --no-cache-dir torch==2.6.0 --index-url https://download.pytorch.org/whl/cu${CUDA_VERSION//.}"
|
TORCH_INSTALL="pip install --no-cache-dir torch==2.6.0 --index-url https://download.pytorch.org/whl/cu${CUDA_VERSION//.}"
|
||||||
fi
|
fi
|
||||||
|
|
||||||
@@ -39,7 +50,7 @@ docker run --rm \
|
|||||||
export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0+PTX' && \
|
export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0+PTX' && \
|
||||||
export CUDA_VERSION=${CUDA_VERSION} && \
|
export CUDA_VERSION=${CUDA_VERSION} && \
|
||||||
mkdir -p /usr/lib/${ARCH}-linux-gnu/ && \
|
mkdir -p /usr/lib/${ARCH}-linux-gnu/ && \
|
||||||
ln -s /usr/local/cuda-${CUDA_VERSION}/targets/x86_64-linux/lib/stubs/libcuda.so /usr/lib/${ARCH}-linux-gnu/libcuda.so && \
|
ln -s /usr/local/cuda-${CUDA_VERSION}/targets/${LIBCUDA_ARCH}-linux/lib/stubs/libcuda.so /usr/lib/${ARCH}-linux-gnu/libcuda.so && \
|
||||||
|
|
||||||
cd /sgl-kernel && \
|
cd /sgl-kernel && \
|
||||||
ls -la ${PYTHON_ROOT_PATH}/lib/python${PYTHON_VERSION}/site-packages/wheel/ && \
|
ls -la ${PYTHON_ROOT_PATH}/lib/python${PYTHON_VERSION}/site-packages/wheel/ && \
|
||||||
|
|||||||
Reference in New Issue
Block a user