A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 19:15 UTC
[AMD][Build] Add DeepEP to ROCm Dockerfile (#36086)
Signed-off-by: Ryan Rock <ryan.rock@amd.com>
Ryan Rock committed
Mar 12, 2026 at 16:33 UTC
a79c1c2c806c7426931f02ad0b81d4656a07cba5
2 files changed
+49
.buildkite/test-amd.yaml
+16
@@ -2071,6 +2071,14 @@ steps:
2071
- pytest -v -s kernels/moe --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
2072
parallelism: 2
2073
2074
+- label: Kernels FP8 MoE Test
2075
+ timeout_in_minutes: 60
2076
+ mirror_hardwares: [amdexperimental, amdproduction]
2077
+ agent_pool: mi325_2
2078
+ optional: true
2079
+ commands:
2080
+ - pytest -v -s kernels/moe/test_deepep_moe.py
2081
+
2082
- label: Kernels Mamba Test # 31min
2083
timeout_in_minutes: 45
2084
mirror_hardwares: [amdexperimental, amdproduction]
@@ -3801,6 +3809,14 @@ steps:
3809
- pytest -v -s kernels/moe --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
3810
parallelism: 2
3811
3812
+- label: Kernels FP8 MoE Test
3813
+ timeout_in_minutes: 60
3814
+ mirror_hardwares: [amdexperimental, amdproduction]
3815
+ agent_pool: mi355_2
3816
+ optional: true
3817
+ commands:
3818
+ - pytest -v -s kernels/moe/test_deepep_moe.py
3819
+
3820
- label: Kernels Mamba Test # 31min
3821
timeout_in_minutes: 45
3822
mirror_hardwares: [amdexperimental, amdproduction]
docker/Dockerfile.rocm
+33
@@ -184,6 +184,34 @@ RUN cd /opt/rixl && mkdir -p /app/install && \
184
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
185
--nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
186
187
+# DeepEP build stage
188
+FROM base AS build_deep
189
+ARG ROCSHMEM_BRANCH="ba0bf0f3"
190
+ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
191
+ARG DEEPEP_BRANCH="e84464ec"
192
+ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
193
+ARG DEEPEP_NIC="cx7"
194
+ENV ROCSHMEM_DIR=/opt/rocshmem
195
+
196
+RUN git clone ${ROCSHMEM_REPO} \
197
+ && cd rocm-systems \
198
+ && git checkout ${ROCSHMEM_BRANCH} \
199
+ && mkdir -p projects/rocshmem/build \
200
+ && cd projects/rocshmem/build \
201
+ && cmake .. \
202
+ -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
203
+ -DROCM_PATH=/opt/rocm \
204
+ -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
205
+ -DUSE_EXTERNAL_MPI=OFF \
206
+ && make -j \
207
+ && make install
208
+
209
+# Build DeepEP wheel.
210
+# DeepEP looks for rocshmem at ROCSHMEM_DIR.
211
+RUN git clone ${DEEPEP_REPO} \
212
+ && cd DeepEP \
213
+ && git checkout ${DEEPEP_BRANCH} \
214
+ && python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
215
216
# -----------------------
217
# vLLM wheel release build stage (for building distributable wheels)
@@ -305,6 +333,11 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
333
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
334
uv pip install --system /rixl_install/*.whl
335
336
+# Install DeepEP wheel
337
+RUN --mount=type=bind,from=build_deep,src=/app/deep_install,target=/deep_install \
338
+ uv pip install --system /deep_install/*.whl
339
+COPY --from=build_deep /opt/rocshmem /opt/rocshmem
340
+
341
# RIXL/MoRIIO runtime dependencies (RDMA userspace libraries)
342
RUN apt-get update -q -y && apt-get install -q -y \
343
librdmacm1 \