Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 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
3 changes: 2 additions & 1 deletion scripts/build_flashinfer_jit_cache_whl.sh
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@ echo "=========================================="
# MAX_JOBS = min(nproc, max(1, MemAvailable_GB/4))
MEM_AVAILABLE_GB=$(free -g | awk '/^Mem:/ {print $7}')
NPROC=$(nproc)
MAX_JOBS=$(( MEM_AVAILABLE_GB / $([ "$(uname -m)" = "aarch64" ] && echo 8 || echo 4) ))
# MAX_JOBS=$(( MEM_AVAILABLE_GB / $([ "$(uname -m)" = "aarch64" ] && echo 8 || echo 4) ))
MAX_JOBS=$(( MEM_AVAILABLE_GB / 8 ))
if (( MAX_JOBS < 1 )); then
MAX_JOBS=1
elif (( NPROC < MAX_JOBS )); then
Expand Down
44 changes: 44 additions & 0 deletions scripts/task_test_blackwell_kernels.sh
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,51 @@ 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)
JIT_ARCH_EFFECTIVE=""
if [ -n "${JIT_ARCH}" ]; then
if [ "${JIT_ARCH}" = "12.0" ]; then
if [ "${CUDA_VERSION}" = "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"

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

# 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