chore(gb200): update to CUDA 12.9 and improve build process (#8772)

This commit is contained in:
ishandhanani
2025-08-08 13:42:47 -07:00
committed by GitHub
parent 36bfddecb9
commit 4e7f025219
7 changed files with 86 additions and 50 deletions

View File

@@ -79,14 +79,17 @@ RUN python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel html5li
python3 -m pip install --no-cache-dir https://github.com/sgl-project/whl/releases/download/v0.3.3/sgl_kernel-0.3.3-cp39-abi3-manylinux2014_x86_64.whl --force-reinstall --no-deps ; \
fi
# Build and install NVSHMEM + DeepEP
RUN wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.3.9/source/nvshmem_src_cuda12-all-all-3.3.9.tar.gz \
&& git clone https://github.com/deepseek-ai/DeepEP.git \
&& cd DeepEP && git checkout ${DEEPEP_COMMIT} && cd .. \
&& tar -xf nvshmem_src_cuda12-all-all-3.3.9.tar.gz && mv nvshmem_src nvshmem \
&& cd nvshmem \
&& rm -f /sgl-workspace/nvshmem_src_cuda12-all-all-3.3.9.tar.gz \
&& NVSHMEM_SHMEM_SUPPORT=0 \
# Download source files
RUN wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.3.9/source/nvshmem_src_cuda12-all-all-3.3.9.tar.gz && \
git clone https://github.com/deepseek-ai/DeepEP.git && \
cd DeepEP && git checkout ${DEEPEP_COMMIT} && cd .. && \
tar -xf nvshmem_src_cuda12-all-all-3.3.9.tar.gz && \
mv nvshmem_src nvshmem && \
rm -f /sgl-workspace/nvshmem_src_cuda12-all-all-3.3.9.tar.gz
# Build and install NVSHMEM
RUN cd /sgl-workspace/nvshmem && \
NVSHMEM_SHMEM_SUPPORT=0 \
NVSHMEM_UCX_SUPPORT=0 \
NVSHMEM_USE_NCCL=0 \
NVSHMEM_MPI_SUPPORT=0 \
@@ -94,10 +97,12 @@ RUN wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.3.9/sour
NVSHMEM_PMIX_SUPPORT=0 \
NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \
NVSHMEM_USE_GDRCOPY=1 \
cmake -S . -B build/ -DCMAKE_INSTALL_PREFIX=${NVSHMEM_DIR} -DCMAKE_CUDA_ARCHITECTURES=90 \
&& cmake --build build --target install -j${CMAKE_BUILD_PARALLEL_LEVEL} \
&& cd /sgl-workspace/DeepEP \
&& NVSHMEM_DIR=${NVSHMEM_DIR} pip install .
cmake -S . -B build/ -DCMAKE_INSTALL_PREFIX=${NVSHMEM_DIR} -DCMAKE_CUDA_ARCHITECTURES="100;120" && \
cmake --build build --target install -j${CMAKE_BUILD_PARALLEL_LEVEL}
# Install DeepEP
RUN cd /sgl-workspace/DeepEP && \
NVSHMEM_DIR=${NVSHMEM_DIR} pip install .
# Python tools
RUN python3 -m pip install --no-cache-dir \
@@ -110,7 +115,8 @@ RUN python3 -m pip install --no-cache-dir \
icdiff \
uv \
wheel \
scikit-build-core
scikit-build-core \
nixl
# Install development tools and utilities
RUN apt-get update && apt-get install -y \