Fix sgl-kernel build on aarch64 platforms (#6062)

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/ && \