Harshith Reddy commited on
Commit
804b6ae
·
1 Parent(s): fde793b

Cap CUDA-extension build parallelism (fix HF Spaces OOMKill, exit 137)

Browse files

Set MAX_JOBS=1 and NVCC_THREADS=1 so mamba_ssm and selective_scan_cuda_oflex
compile sequentially. The HF Spaces build VM has ~16 GB RAM; parallel nvcc
workers (3-4 GB each) consistently OOMKill on the bench Space build.
Serial build adds ~10 minutes but completes reliably.

Files changed (1) hide show
  1. Dockerfile +7 -0
Dockerfile CHANGED
@@ -12,6 +12,13 @@ ENV TORCH_CUDA_ARCH_LIST=8.9
12
  ENV CUDA_HOME=/usr/local/cuda
13
  ENV MPLCONFIGDIR=/tmp/matplotlib
14
 
 
 
 
 
 
 
 
15
  RUN apt-get update && apt-get install -y \
16
  python3.10 \
17
  python3.10-venv \
 
12
  ENV CUDA_HOME=/usr/local/cuda
13
  ENV MPLCONFIGDIR=/tmp/matplotlib
14
 
15
+ # Cap CUDA-extension build parallelism so the HF Spaces build VM (~16 GB
16
+ # RAM) does not OOMKill during mamba_ssm / selective_scan_cuda_oflex
17
+ # compilation. Each nvcc worker can use 3-4 GB; serial build is slower
18
+ # (~25 min total) but reliable.
19
+ ENV MAX_JOBS=1
20
+ ENV NVCC_THREADS=1
21
+
22
  RUN apt-get update && apt-get install -y \
23
  python3.10 \
24
  python3.10-venv \