From 0618ad6dd559ae5f2a637d00fc1e6c40643fc318 Mon Sep 17 00:00:00 2001
From: gongwei-130 <56567052+gongwei-130@users.noreply.github.com>
Date: Thu, 2 Oct 2025 13:24:23 -0700
Subject: [PATCH] fix: shoudn't include CUDA_ARCH 100 and 120 for cuda12.6.1
 (#11176)

---
 docker/Dockerfile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index a6229288d..4aef6cf13 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -109,7 +109,7 @@ RUN wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.3.9/sour
 
 # Build and install NVSHMEM
 RUN cd /sgl-workspace/nvshmem && \
-    export CUDA_ARCH="90;100;120" && \
+    if [ "$BUILD_TYPE" = "blackwell" ] || [ "$BUILD_TYPE" = "blackwell_aarch" ]; then CUDA_ARCH="90;100;120"; else CUDA_ARCH="90"; fi && \
     NVSHMEM_SHMEM_SUPPORT=0 \
     NVSHMEM_UCX_SUPPORT=0 \
     NVSHMEM_USE_NCCL=0 \