From 392de007cb6f378c9b75546e81f7403b778635e6 Mon Sep 17 00:00:00 2001
From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
Date: Fri, 15 Aug 2025 08:41:49 +0800
Subject: [PATCH] Minor fix docker container DeepEP on multi platforms (#9205)

---
 docker/Dockerfile | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index aac130d8a..2801b0737 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -102,7 +102,18 @@ RUN cd /sgl-workspace/nvshmem && \
 
 # Install DeepEP
 RUN cd /sgl-workspace/DeepEP && \
-    TORCH_CUDA_ARCH_LIST='9.0;10.0' NVSHMEM_DIR=${NVSHMEM_DIR} pip install .
+    case "$CUDA_VERSION" in \
+      12.6.1) \
+        CHOSEN_TORCH_CUDA_ARCH_LIST='9.0' \
+        ;; \
+      12.8.1|12.9.1) \
+        CHOSEN_TORCH_CUDA_ARCH_LIST='9.0;10.0' \
+        ;; \
+      *) \
+        echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 \
+        ;; \
+    esac && \
+    NVSHMEM_DIR=${NVSHMEM_DIR} TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" pip install .
 
 # Python tools
 RUN python3 -m pip install --no-cache-dir \