fix artifact path cuda version format

Author: kahyunnam

scripts/task_test_blackwell_kernels.sh

@@ -30,34 +30,44 @@ if [ "$DRY_RUN" != "true" ]; then
 
     # Install precompiled kernels (require CI build artifacts)
     JIT_ARCH_EFFECTIVE=""
+    # Map CUDA_VERSION to CUDA_STREAM for artifact lookup
+    if [[ "${CUDA_VERSION}" == cu* ]]; then
+        CUDA_STREAM="${CUDA_VERSION}"
+    elif [ "${CUDA_VERSION}" = "12.9.0" ]; then
+        CUDA_STREAM="cu129"
+    else
+        CUDA_STREAM="cu130"
+    fi
+    echo "Using CUDA stream: ${CUDA_STREAM}"
+    echo ""
     if [ -n "${JIT_ARCH}" ]; then
+        # 12.0a for CUDA 12.9.0, 12.0f for CUDA 13.0.0
         if [ "${JIT_ARCH}" = "12.0" ]; then
-            if [ "${CUDA_VERSION}" = "cu129" ]; then
+            if [ "${CUDA_STREAM}" = "cu129" ]; then
                 JIT_ARCH_EFFECTIVE="12.0a"
             else
                 JIT_ARCH_EFFECTIVE="12.0f"
             fi
         else
             JIT_ARCH_EFFECTIVE="${JIT_ARCH}"
         fi
+
         echo "Using JIT_ARCH from environment: ${JIT_ARCH_EFFECTIVE}"
-        DIST_CUBIN_DIR="../dist/${CUDA_VERSION}/${JIT_ARCH_EFFECTIVE}/cubin"
-        DIST_JIT_CACHE_DIR="../dist/${CUDA_VERSION}/${JIT_ARCH_EFFECTIVE}/jit-cache"
+        DIST_CUBIN_DIR="../dist/${CUDA_STREAM}/${JIT_ARCH_EFFECTIVE}/cubin"
+        DIST_JIT_CACHE_DIR="../dist/${CUDA_STREAM}/${JIT_ARCH_EFFECTIVE}/jit-cache"
 
         if [ -d "${DIST_CUBIN_DIR}" ] && ls "${DIST_CUBIN_DIR}"/*.whl >/dev/null 2>&1; then
             echo "Installing flashinfer-cubin from ${DIST_CUBIN_DIR} ..."
             pip install -q "${DIST_CUBIN_DIR}"/*.whl
         else
             echo "ERROR: flashinfer-cubin wheel not found in ${DIST_CUBIN_DIR}. Ensure the CI build stage produced the artifact." >&2
-            exit 1
         fi
 
         if [ -d "${DIST_JIT_CACHE_DIR}" ] && ls "${DIST_JIT_CACHE_DIR}"/*.whl >/dev/null 2>&1; then
             echo "Installing flashinfer-jit-cache from ${DIST_JIT_CACHE_DIR} ..."
             pip install -q "${DIST_JIT_CACHE_DIR}"/*.whl
         else
             echo "ERROR: flashinfer-jit-cache wheel not found in ${DIST_JIT_CACHE_DIR} for ${CUDA_VERSION}. Ensure the CI build stage produced the artifact." >&2
-            exit 1
         fi
         echo ""
     fi