Support Blackwell DeepEP docker images (#6868)

This commit is contained in:
fzyzcjy
2025-06-05 15:07:53 +08:00
committed by GitHub
parent bcf66ef3e1
commit 0166403c20
5 changed files with 23 additions and 121 deletions

View File

@@ -1,4 +1,5 @@
FROM lmsysorg/sglang:latest
ARG BASE_IMAGE
FROM ${BASE_IMAGE}
# CMake
RUN apt-get update \
@@ -55,6 +56,9 @@ RUN tar -xf nvshmem_src_3.2.5-1.txz \
WORKDIR /sgl-workspace/nvshmem
RUN git apply /sgl-workspace/DeepEP/third-party/nvshmem.patch
RUN sed -i '1i#include <unistd.h>' /sgl-workspace/nvshmem/examples/moe_shuffle.cu && \
cat /sgl-workspace/nvshmem/examples/moe_shuffle.cu
WORKDIR /sgl-workspace/nvshmem
ENV CUDA_HOME=/usr/local/cuda
RUN NVSHMEM_SHMEM_SUPPORT=0 \
@@ -71,7 +75,7 @@ RUN NVSHMEM_SHMEM_SUPPORT=0 \
WORKDIR /sgl-workspace/DeepEP
ENV NVSHMEM_DIR=/sgl-workspace/nvshmem/install
RUN NVSHMEM_DIR=/sgl-workspace/nvshmem/install pip install .
RUN NVSHMEM_DIR=/sgl-workspace/nvshmem/install pip install --break-system-packages .
# Set workspace
WORKDIR /sgl-workspace