From aff584fa541e494cd8815888a479d03a387dd2a5 Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Tue, 6 May 2025 19:10:57 -0400 Subject: [PATCH] Fix sgl-kernel build on aarch64 platforms (#6062) --- sgl-kernel/build.sh | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/sgl-kernel/build.sh b/sgl-kernel/build.sh index 65d4edc53..7fe089770 100755 --- a/sgl-kernel/build.sh +++ b/sgl-kernel/build.sh @@ -7,12 +7,23 @@ PYTHON_ROOT_PATH=/opt/python/cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.} ARCH=$(uname -i) 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 - 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//.}" 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//.}" fi @@ -39,7 +50,7 @@ docker run --rm \ export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0+PTX' && \ export CUDA_VERSION=${CUDA_VERSION} && \ 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 && \ ls -la ${PYTHON_ROOT_PATH}/lib/python${PYTHON_VERSION}/site-packages/wheel/ && \