File size: 1,058 Bytes
84bcfb1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 | # Source this on a GPU machine before running anything: . /nvme/wsp-lab/code/scripts/env.sh
# It activates the lab's Python and keeps every cache inside /nvme/wsp-lab.
export LAB=/nvme/wsp-lab
. "$LAB/venv/bin/activate"
export HF_HOME=$LAB/hf NLTK_DATA=$LAB/nltk_data VLLM_CACHE_ROOT=$LAB/vllm-cache \
FLASHINFER_WORKSPACE_BASE=$LAB TRITON_CACHE_DIR=$LAB/triton-cache TORCHINDUCTOR_CACHE_DIR=$LAB/inductor-cache \
XDG_CACHE_HOME=$LAB/xdg-cache UV_CACHE_DIR=$LAB/uv-cache MPLCONFIGDIR=$LAB/mpl TOKENIZERS_PARALLELISM=false
# FlashInfer compiles its sampling kernel on first use and needs the cuRAND headers. Some machines' CUDA lacks
# them; the venv's nvidia-curand wheel has a copy, which we expose through a folder that holds only those files.
if [ ! -e "${CUDA_HOME:-/usr/local/cuda}/include/curand.h" ]; then
_inc=$LAB/cuda-extra-include
if [ ! -e "$_inc/curand.h" ]; then
mkdir -p "$_inc" && ln -sf "$LAB"/venv/lib/python3.12/site-packages/nvidia/cu13/include/curand*.h "$_inc"/
fi
export FLASHINFER_EXTRA_CUDAFLAGS="-isystem $_inc"
fi
|