ci: align torch CUDA wheel checks (#24940)

This commit is contained in:
Hank Han
2026-05-11 16:22:20 +08:00
committed by GitHub
parent f7ba136279
commit c027ae677c
+13
View File
@@ -277,7 +277,20 @@ install_sglang_kernel() {
# TODO: Remove after torch 2.11 where cu13 is enabled by default # TODO: Remove after torch 2.11 where cu13 is enabled by default
TORCH_CUDA_VER=$(python3 -c "import torch; v=torch.version.cuda; parts=v.split('.'); print(f'cu{parts[0]}{parts[1]}')") TORCH_CUDA_VER=$(python3 -c "import torch; v=torch.version.cuda; parts=v.split('.'); print(f'cu{parts[0]}{parts[1]}')")
echo "Detected torch CUDA version: ${TORCH_CUDA_VER}" echo "Detected torch CUDA version: ${TORCH_CUDA_VER}"
TORCHAUDIO_CUDA_VER=$(pip show torchaudio 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed -n 's/.*+\(cu[0-9][0-9]*\)$/\1/p' || true)
TORCHVISION_CUDA_VER=$(pip show torchvision 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed -n 's/.*+\(cu[0-9][0-9]*\)$/\1/p' || true)
REINSTALL_TORCH=false
if [ "${TORCH_CUDA_VER}" != "${CU_VERSION}" ]; then if [ "${TORCH_CUDA_VER}" != "${CU_VERSION}" ]; then
REINSTALL_TORCH=true
else
for cuda_ver in "${TORCHAUDIO_CUDA_VER}" "${TORCHVISION_CUDA_VER}"; do
if [ -n "${cuda_ver}" ] && [ "${cuda_ver}" != "${CU_VERSION}" ]; then
REINSTALL_TORCH=true
break
fi
done
fi
if [ "${REINSTALL_TORCH}" = true ]; then
TORCH_VER=$(pip show torch 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//') TORCH_VER=$(pip show torch 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
TORCHAUDIO_VER=$(pip show torchaudio 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//') TORCHAUDIO_VER=$(pip show torchaudio 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
TORCHVISION_VER=$(pip show torchvision 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//') TORCHVISION_VER=$(pip show torchvision 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')