From 141a45964453b7a3b827303ca7847d5902f974c7 Mon Sep 17 00:00:00 2001
From: Yineng Zhang <me@zhyncs.com>
Date: Fri, 2 May 2025 11:31:29 -0700
Subject: [PATCH] fix: only upgrade nccl for cu128 (#5986)

---
 docker/Dockerfile | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 125ce670b..ab1b5c664 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -29,7 +29,7 @@ RUN python3 -m pip install --upgrade pip setuptools wheel html5lib six \
          export CUINDEX=121; \
        elif [ "$CUDA_VERSION" = "12.4.1" ]; then \
          export CUINDEX=124; \
-       elif [ "$CUDA_VERSION" = "12.5.1" ]; then \
+       elif [ "$CUDA_VERSION" = "12.8.1" ]; then \
          export CUINDEX=124; \
        elif [ "$CUDA_VERSION" = "11.8.0" ]; then \
          export CUINDEX=118; \
@@ -40,6 +40,8 @@ RUN python3 -m pip install --upgrade pip setuptools wheel html5lib six \
     && python3 -m pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu${CUINDEX} \
     && cd sglang \
     && python3 -m pip --no-cache-dir install -e "python[${BUILD_TYPE}]" --find-links https://flashinfer.ai/whl/cu${CUINDEX}/torch2.6/flashinfer-python \
-    && python3 -m pip install nvidia-nccl-cu12==2.26.2.post1 --force-reinstall --no-deps
+    && if [ "$CUDA_VERSION" = "12.8.1" ]; then \
+         python3 -m pip install nvidia-nccl-cu12==2.26.2.post1 --force-reinstall --no-deps; \
+       fi
 
 ENV DEBIAN_FRONTEND=interactive