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 4, 11:15 UTC
[ROCm] ROCm DeepEP API updated to latest (#39721)
Signed-off-by: Tej Kiran <vpolamre@amd.com> Signed-off-by: tej <37236721+itej89@users.noreply.github.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: HAIAI <39548240+HAIAI@users.noreply.github.com>
tej committed
Apr 30, 2026 at 09:46 UTC
121dbe7a221d2bf6415caea0acb812c666a7665e
3 files changed
+42
-60
docker/Dockerfile.rocm
+15
-8
@@ -200,9 +200,9 @@ RUN cd /opt/rixl && \
200
201
# DeepEP build stage
202
FROM base AS build_deep
203
-ARG ROCSHMEM_BRANCH="ba0bf0f3"
203
+ARG ROCSHMEM_BRANCH="f0acb0c6"
204
ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
205
-ARG DEEPEP_BRANCH="5d90af8b"
205
+ARG DEEPEP_BRANCH="a9ea9774"
206
ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
207
ARG DEEPEP_NIC="cx7"
208
ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
@@ -213,18 +213,15 @@ RUN git clone ${ROCSHMEM_REPO} \
213
&& git checkout ${ROCSHMEM_BRANCH} \
214
&& mkdir -p projects/rocshmem/build \
215
&& cd projects/rocshmem/build \
216
- && bash ../scripts/build_configs/all_backends \
217
- -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
218
- -DROCM_PATH=/opt/rocm \
219
- -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
220
- -DUSE_EXTERNAL_MPI=OFF
216
+ && INSTALL_PREFIX=${ROCSHMEM_DIR} \
217
+ ../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF
218
219
# Build DeepEP wheel.
220
# DeepEP looks for rocshmem at ROCSHMEM_DIR.
221
RUN git clone ${DEEPEP_REPO} \
222
&& cd DeepEP \
223
&& git checkout ${DEEPEP_BRANCH} \
227
- && python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
224
+ && python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
225
226
# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
227
# not force users to rebuild the long-lived Dockerfile.rocm_base image.
@@ -388,6 +385,16 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
385
# above are not available once that RUN step completes.
386
COPY --from=export_vllm /*.whl /opt/vllm-wheels/
387
388
+# Update rdma-core to support latest rocshmem
389
+ARG DEEPEP_NIC
390
+RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
391
+ git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
392
+ cd /tmp/rdma-core && \
393
+ mkdir -p build && cd build && \
394
+ cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
395
+ ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
396
+fi
397
+
398
# Install RIXL wheel
399
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
400
uv pip install --system /rixl_install/*.whl
vllm/distributed/device_communicators/all2all.py
+4
-11
@@ -10,7 +10,6 @@ import vllm.envs as envs
10
from vllm.distributed import get_dp_group, get_ep_group
11
from vllm.forward_context import get_forward_context
12
from vllm.logger import init_logger
13
-from vllm.platforms import current_platform
13
from vllm.utils.flashinfer import (
14
has_flashinfer_nvlink_one_sided,
15
has_flashinfer_nvlink_two_sided,
@@ -225,11 +224,8 @@ class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):
224
num_rdma_bytes=num_rdma_bytes,
225
low_latency_mode=False,
226
num_qps_per_rank=num_qps_per_rank,
227
+ explicitly_destroy=True,
228
)
229
- if not current_platform.is_rocm():
230
- kwargs.update(
231
- explicitly_destroy=True,
232
- )
229
return kwargs
230
231
def get_handle(self, kwargs):
@@ -303,13 +299,10 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
299
num_rdma_bytes=num_rdma_bytes,
300
low_latency_mode=True,
301
num_qps_per_rank=num_qps_per_rank,
302
+ allow_nvlink_for_low_latency_mode=True,
303
+ allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
304
+ explicitly_destroy=True,
305
)
307
- if not current_platform.is_rocm():
308
- kwargs.update(
309
- allow_nvlink_for_low_latency_mode=True,
310
- allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
311
- explicitly_destroy=True,
312
- )
306
return kwargs
307
308
def get_handle(self, kwargs):
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py
+23
-41
@@ -16,7 +16,6 @@ from vllm.model_executor.layers.fused_moe.utils import (
16
moe_kernel_quantize_input,
17
normalize_batched_scales_shape,
18
)
19
-from vllm.platforms import current_platform
19
from vllm.v1.worker.ubatching import (
20
dbo_current_ubatch_id,
21
dbo_enabled,
@@ -290,46 +289,29 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
289
290
# Dispatch
291
dispatch_topk_ids = self._map_global_to_physical_ids(topk_ids)
293
- if current_platform.is_rocm():
294
- (
295
- expert_x,
296
- expert_num_tokens,
297
- handle,
298
- _,
299
- hook,
300
- ) = self.buffer.low_latency_dispatch(
301
- a1,
302
- dispatch_topk_ids,
303
- self.max_tokens_per_rank,
304
- num_experts,
305
- use_fp8=self.use_fp8_dispatch,
306
- async_finish=False,
307
- return_recv_hook=True,
308
- )
309
- else:
310
- (
311
- expert_x,
312
- expert_num_tokens,
313
- handle,
314
- _,
315
- hook,
316
- ) = self.buffer.low_latency_dispatch(
317
- a1,
318
- dispatch_topk_ids,
319
- self.max_tokens_per_rank,
320
- num_experts,
321
- use_fp8=self.use_fp8_dispatch,
322
- round_scale=self.use_ue8m0_dispatch,
323
- use_ue8m0=self.use_ue8m0_dispatch,
324
- **(dict(use_nvfp4=True) if use_nvfp4 else dict()),
325
- **(
326
- dict(x_global_scale=qc_a1_gscale_or_scale)
327
- if qc_a1_gscale_or_scale is not None and nvfp4_dispatch
328
- else dict()
329
- ),
330
- async_finish=False,
331
- return_recv_hook=True,
332
- )
292
+ (
293
+ expert_x,
294
+ expert_num_tokens,
295
+ handle,
296
+ _,
297
+ hook,
298
+ ) = self.buffer.low_latency_dispatch(
299
+ a1,
300
+ dispatch_topk_ids,
301
+ self.max_tokens_per_rank,
302
+ num_experts,
303
+ use_fp8=self.use_fp8_dispatch,
304
+ round_scale=self.use_ue8m0_dispatch,
305
+ use_ue8m0=self.use_ue8m0_dispatch,
306
+ **(dict(use_nvfp4=True) if use_nvfp4 else dict()),
307
+ **(
308
+ dict(x_global_scale=qc_a1_gscale_or_scale)
309
+ if qc_a1_gscale_or_scale is not None and nvfp4_dispatch
310
+ else dict()
311
+ ),
312
+ async_finish=False,
313
+ return_recv_hook=True,
314
+ )
315
self.handles[a2a_idx] = handle
316
317
return (