From 8fe9bbffb6fd34365343db637948f4126cab9e0c Mon Sep 17 00:00:00 2001
From: Baizhou Zhang <sobereddiezhang@gmail.com>
Date: Mon, 13 Apr 2026 23:04:23 -0700
Subject: [PATCH] [CI] Reinstall flashinfer-jit-cache on CUDA version mismatch
 (#22741)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh | 2 +-
 scripts/ci/cuda/ci_install_dependency.sh            | 6 ++++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh b/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh
index ab41dcef2..0f35e57e0 100755
--- a/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh
+++ b/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh
@@ -25,7 +25,7 @@ if [ "$FLASHINFER_JIT_CACHE_INSTALLED" = false ]; then
     FLASHINFER_CACHE_DIR="${HOME}/.cache/flashinfer-wheels"
     mkdir -p "${FLASHINFER_CACHE_DIR}"
 
-    FLASHINFER_WHEEL_PATTERN="flashinfer_jit_cache-${FLASHINFER_PYTHON_REQUIRED}*.whl"
+    FLASHINFER_WHEEL_PATTERN="flashinfer_jit_cache-${FLASHINFER_PYTHON_REQUIRED}+${CU_VERSION}*.whl"
     CACHED_WHEEL=$(find "${FLASHINFER_CACHE_DIR}" -name "${FLASHINFER_WHEEL_PATTERN}" -type f 2>/dev/null | head -n 1)
 
     if [ -n "$CACHED_WHEEL" ] && [ -f "$CACHED_WHEEL" ]; then
diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh
index 7735977ef..afd0383bb 100755
--- a/scripts/ci/cuda/ci_install_dependency.sh
+++ b/scripts/ci/cuda/ci_install_dependency.sh
@@ -187,6 +187,7 @@ FLASHINFER_PYTHON_REQUIRED=$(grep -Po -m1 '(?<=flashinfer_python==)[0-9A-Za-z\.\
 FLASHINFER_CUBIN_REQUIRED=$(grep -Po -m1 '(?<=flashinfer_cubin==)[0-9A-Za-z\.\-]+' python/pyproject.toml || echo "")
 FLASHINFER_CUBIN_INSTALLED=$(pip show flashinfer-cubin 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
 FLASHINFER_JIT_INSTALLED=$(pip show flashinfer-jit-cache 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//' || echo "")
+FLASHINFER_JIT_CU_VERSION=$(pip show flashinfer-jit-cache 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed -n 's/.*+//p' || echo "")
 
 UNINSTALL_CUBIN=true
 UNINSTALL_JIT_CACHE=true
@@ -205,6 +206,11 @@ else
     echo "flashinfer-jit-cache version mismatch (installed: ${FLASHINFER_JIT_INSTALLED:-none}, required: ${FLASHINFER_PYTHON_REQUIRED}), will reinstall"
 fi
 
+if [ "$UNINSTALL_JIT_CACHE" = false ] && [ "$FLASHINFER_JIT_CU_VERSION" != "$CU_VERSION" ]; then
+    echo "flashinfer-jit-cache CUDA version mismatch (installed: ${FLASHINFER_JIT_CU_VERSION:-none}, required: ${CU_VERSION}), will reinstall"
+    UNINSTALL_JIT_CACHE=true
+fi
+
 # Build uninstall list based on what needs updating
 FLASHINFER_UNINSTALL="flashinfer-python"
 [ "$UNINSTALL_CUBIN" = true ] && FLASHINFER_UNINSTALL="$FLASHINFER_UNINSTALL flashinfer-cubin"