Skip to content
Open
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions scripts/task_test_blackwell_kernels.sh
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,38 @@ if [[ "$1" == "--dry-run" ]] || [[ "${DRY_RUN}" == "true" ]]; then
fi

if [ "$DRY_RUN" != "true" ]; then
echo "Using CUDA version: ${CUDA_VERSION}"
echo ""

# Install precompiled kernels (require CI build artifacts)
: ${DIST_CUBIN_DIR:="../dist/${CUDA_VERSION}/cubin"}
: ${DIST_JIT_CACHE_DIR:="../dist/${CUDA_VERSION}/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 ""

# Install local python sources
pip install -e . -v --no-deps
echo ""

# Verify installation
echo "Verifying installation..."
(cd /tmp && python -m flashinfer show-config)
echo ""
fi

EXIT_CODE=0
Expand Down