[ROCm][DeepEP] Stabilize high-throughput DBO for DP+EP (#46990)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

vllm/config/vllm.py

@@ -931,16 +931,28 @@ class VllmConfig:
                     model_type,
                 )
 
+        from vllm.platforms import current_platform
         from vllm.v1.executor.abstract import Executor
 
         executor_backend = self.parallel_config.distributed_executor_backend
         executor_class = Executor.get_class(self)
         executor_supports_async_sched = executor_class.supports_async_scheduling()
+        uses_rocm_deepep_ht_dbo = (
+            current_platform.is_rocm()
+            and self.parallel_config.enable_dbo
+            and self.parallel_config.all2all_backend == "deepep_high_throughput"
+        )
 
         if self.scheduler_config.async_scheduling:
             # Async scheduling explicitly enabled, hard fail any incompatibilities.
             # Currently, async scheduling only support eagle speculative
             # decoding.
+            if uses_rocm_deepep_ht_dbo:
+                raise ValueError(
+                    "Async scheduling is not compatible with ROCm DeepEP "
+                    "high-throughput DBO. Please use --no-async-scheduling or "
+                    "select a different all2all backend."
+                )
             if self.speculative_config is not None:
                 if (
                     self.speculative_config.method not in get_args(EagleModelTypes)
@@ -1000,6 +1012,13 @@ class VllmConfig:
                     executor_backend,
                 )
                 self.scheduler_config.async_scheduling = False
+            elif uses_rocm_deepep_ht_dbo:
+                logger.warning_once(
+                    "Async scheduling is disabled for ROCm DeepEP "
+                    "high-throughput DBO because that combination can corrupt "
+                    "DP+EP generation accuracy."
+                )
+                self.scheduler_config.async_scheduling = False
             else:
                 self.scheduler_config.async_scheduling = True
 
@@ -1044,8 +1063,6 @@ class VllmConfig:
                 "VLLM_WORKER_MULTIPROC_METHOD set to spawn"
             )
 
-        from vllm.platforms import current_platform
-
         if (
             self.model_config is not None
             and self.scheduler_config.enable_chunked_prefill

vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py

@@ -12,6 +12,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceDelegate,
 )
 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
+from vllm.platforms import current_platform
 from vllm.utils.math_utils import round_up
 from vllm.v1.worker.ubatching import (
     dbo_current_ubatch_id,
@@ -59,6 +60,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         self.dp_size = dp_size
         self.rank_expert_offset = rank_expert_offset
         self.async_prepare = True
+        self.sync_dbo_comm = current_platform.is_rocm()
 
         # The dispatch function returns a handle that the combine function
         # requires. Under DBO microbatching we must track one handle per
@@ -68,6 +70,13 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         # From https://github.com/deepseek-ai/DeepEP/blob/9fe9021f29c9083cd1808ab36b740208524d9f63/deep_ep/buffer.py#L164
         self.available_rank_configs = [2, 4, 8, 16, 24, 32, 64, 128, 144, 160]
 
+    def _sync_dbo_comm_if_needed(self) -> None:
+        if self.sync_dbo_comm and dbo_enabled():
+            # ROCm DeepEP HT dispatch/combine reuse Buffer-owned communication
+            # workspace. Do not let the next DBO ubatch reuse that workspace
+            # before this ubatch's HT kernel has completed.
+            torch.cuda.current_stream().synchronize()
+
     def num_dispatchers(self) -> int:
         return self.num_dispatchers_
 
@@ -161,6 +170,8 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
             allocate_on_comm_stream=False,
         )
 
+        self._sync_dbo_comm_if_needed()
+
         # record the handle for this ubatch
         a2a_idx = dbo_current_ubatch_id()
         self.handles[a2a_idx] = handle
@@ -375,6 +386,8 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
             allocate_on_comm_stream=False,
         )
 
+        self._sync_dbo_comm_if_needed()
+
         dbo_switch_to_compute()
 
         if do_async:

vllm/v1/worker/gpu_ubatch_wrapper.py

@@ -154,6 +154,16 @@ class UBatchWrapper:
     @staticmethod
     def _create_sm_control_context(vllm_config: VllmConfig):
         comm_sms: int = envs.VLLM_DBO_COMM_SMS
+        rocm_deepep_ht_dbo = (
+            current_platform.is_rocm()
+            and vllm_config.parallel_config.enable_dbo
+            and vllm_config.parallel_config.all2all_backend == "deepep_high_throughput"
+        )
+        if rocm_deepep_ht_dbo:
+            # On ROCm, reserving CUs for DeepEP HT communication under DBO
+            # corrupts DP+EP generation accuracy. Keep the backend active, but
+            # leave all CUs visible to the compute and communication kernels.
+            comm_sms = 0
 
         set_comm_sms = lambda sms: None
         if vllm_config.parallel_config.enable_expert_parallel: