| # 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 | |