diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 7a93823cd2d..f416c813611 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -200,9 +200,9 @@ RUN cd /opt/rixl && \ # DeepEP build stage FROM base AS build_deep -ARG ROCSHMEM_BRANCH="ba0bf0f3" +ARG ROCSHMEM_BRANCH="f0acb0c6" ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" -ARG DEEPEP_BRANCH="5d90af8b" +ARG DEEPEP_BRANCH="a9ea9774" ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" ARG DEEPEP_NIC="cx7" ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" @@ -213,18 +213,15 @@ RUN git clone ${ROCSHMEM_REPO} \ && git checkout ${ROCSHMEM_BRANCH} \ && mkdir -p projects/rocshmem/build \ && cd projects/rocshmem/build \ - && bash ../scripts/build_configs/all_backends \ - -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \ - -DROCM_PATH=/opt/rocm \ - -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ - -DUSE_EXTERNAL_MPI=OFF + && INSTALL_PREFIX=${ROCSHMEM_DIR} \ + ../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF # Build DeepEP wheel. # DeepEP looks for rocshmem at ROCSHMEM_DIR. RUN git clone ${DEEPEP_REPO} \ && cd DeepEP \ && git checkout ${DEEPEP_BRANCH} \ - && python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install + && python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install # MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do # 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 \ # above are not available once that RUN step completes. COPY --from=export_vllm /*.whl /opt/vllm-wheels/ +# Update rdma-core to support latest rocshmem +ARG DEEPEP_NIC +RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ + git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ + cd /tmp/rdma-core && \ + mkdir -p build && cd build && \ + cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ + ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ +fi + # Install RIXL wheel RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ uv pip install --system /rixl_install/*.whl diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 6a15d3f6168..340b6ff1cf2 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -10,7 +10,6 @@ import vllm.envs as envs from vllm.distributed import get_dp_group, get_ep_group from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.platforms import current_platform from vllm.utils.flashinfer import ( has_flashinfer_nvlink_one_sided, has_flashinfer_nvlink_two_sided, @@ -225,11 +224,8 @@ class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase): num_rdma_bytes=num_rdma_bytes, low_latency_mode=False, num_qps_per_rank=num_qps_per_rank, + explicitly_destroy=True, ) - if not current_platform.is_rocm(): - kwargs.update( - explicitly_destroy=True, - ) return kwargs def get_handle(self, kwargs): @@ -303,13 +299,10 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase): num_rdma_bytes=num_rdma_bytes, low_latency_mode=True, num_qps_per_rank=num_qps_per_rank, + allow_nvlink_for_low_latency_mode=True, + allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL, + explicitly_destroy=True, ) - if not current_platform.is_rocm(): - kwargs.update( - allow_nvlink_for_low_latency_mode=True, - allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL, - explicitly_destroy=True, - ) return kwargs def get_handle(self, kwargs): diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py index 058d09d23bf..e5d2b601a76 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py @@ -16,7 +16,6 @@ from vllm.model_executor.layers.fused_moe.utils import ( moe_kernel_quantize_input, normalize_batched_scales_shape, ) -from vllm.platforms import current_platform from vllm.v1.worker.ubatching import ( dbo_current_ubatch_id, dbo_enabled, @@ -290,46 +289,29 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): # Dispatch dispatch_topk_ids = self._map_global_to_physical_ids(topk_ids) - if current_platform.is_rocm(): - ( - expert_x, - expert_num_tokens, - handle, - _, - hook, - ) = self.buffer.low_latency_dispatch( - a1, - dispatch_topk_ids, - self.max_tokens_per_rank, - num_experts, - use_fp8=self.use_fp8_dispatch, - async_finish=False, - return_recv_hook=True, - ) - else: - ( - expert_x, - expert_num_tokens, - handle, - _, - hook, - ) = self.buffer.low_latency_dispatch( - a1, - dispatch_topk_ids, - self.max_tokens_per_rank, - num_experts, - use_fp8=self.use_fp8_dispatch, - round_scale=self.use_ue8m0_dispatch, - use_ue8m0=self.use_ue8m0_dispatch, - **(dict(use_nvfp4=True) if use_nvfp4 else dict()), - **( - dict(x_global_scale=qc_a1_gscale_or_scale) - if qc_a1_gscale_or_scale is not None and nvfp4_dispatch - else dict() - ), - async_finish=False, - return_recv_hook=True, - ) + ( + expert_x, + expert_num_tokens, + handle, + _, + hook, + ) = self.buffer.low_latency_dispatch( + a1, + dispatch_topk_ids, + self.max_tokens_per_rank, + num_experts, + use_fp8=self.use_fp8_dispatch, + round_scale=self.use_ue8m0_dispatch, + use_ue8m0=self.use_ue8m0_dispatch, + **(dict(use_nvfp4=True) if use_nvfp4 else dict()), + **( + dict(x_global_scale=qc_a1_gscale_or_scale) + if qc_a1_gscale_or_scale is not None and nvfp4_dispatch + else dict() + ), + async_finish=False, + return_recv_hook=True, + ) self.handles[a2a_idx] = handle return (