[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>
Andreas Karatzas committed
Jun 29, 2026 at 16:28 UTC
53f7553f099c2cbb88d2161959fc49dd71c8205b
3 files changed
+42
-2
vllm/config/vllm.py
+19
-2
@@ -931,16 +931,28 @@ class VllmConfig:
931
model_type,
932
)
933
934
+ from vllm.platforms import current_platform
935
from vllm.v1.executor.abstract import Executor
936
937
executor_backend = self.parallel_config.distributed_executor_backend
938
executor_class = Executor.get_class(self)
939
executor_supports_async_sched = executor_class.supports_async_scheduling()
940
+ uses_rocm_deepep_ht_dbo = (
941
+ current_platform.is_rocm()
942
+ and self.parallel_config.enable_dbo
943
+ and self.parallel_config.all2all_backend == "deepep_high_throughput"
944
+ )
945
946
if self.scheduler_config.async_scheduling:
947
# Async scheduling explicitly enabled, hard fail any incompatibilities.
948
# Currently, async scheduling only support eagle speculative
949
# decoding.
950
+ if uses_rocm_deepep_ht_dbo:
951
+ raise ValueError(
952
+ "Async scheduling is not compatible with ROCm DeepEP "
953
+ "high-throughput DBO. Please use --no-async-scheduling or "
954
+ "select a different all2all backend."
955
+ )
956
if self.speculative_config is not None:
957
if (
958
self.speculative_config.method not in get_args(EagleModelTypes)
@@ -1000,6 +1012,13 @@ class VllmConfig:
1012
executor_backend,
1013
)
1014
self.scheduler_config.async_scheduling = False
1015
+ elif uses_rocm_deepep_ht_dbo:
1016
+ logger.warning_once(
1017
+ "Async scheduling is disabled for ROCm DeepEP "
1018
+ "high-throughput DBO because that combination can corrupt "
1019
+ "DP+EP generation accuracy."
1020
+ )
1021
+ self.scheduler_config.async_scheduling = False
1022
else:
1023
self.scheduler_config.async_scheduling = True
1024
@@ -1044,8 +1063,6 @@ class VllmConfig:
1063
"VLLM_WORKER_MULTIPROC_METHOD set to spawn"
1064
)
1065
1047
- from vllm.platforms import current_platform
1048
-
1066
if (
1067
self.model_config is not None
1068
and self.scheduler_config.enable_chunked_prefill
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py
+13
@@ -12,6 +12,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
12
TopKWeightAndReduceDelegate,
13
)
14
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
15
+from vllm.platforms import current_platform
16
from vllm.utils.math_utils import round_up
17
from vllm.v1.worker.ubatching import (
18
dbo_current_ubatch_id,
@@ -59,6 +60,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
60
self.dp_size = dp_size
61
self.rank_expert_offset = rank_expert_offset
62
self.async_prepare = True
63
+ self.sync_dbo_comm = current_platform.is_rocm()
64
65
# The dispatch function returns a handle that the combine function
66
# requires. Under DBO microbatching we must track one handle per
@@ -68,6 +70,13 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
70
# From https://github.com/deepseek-ai/DeepEP/blob/9fe9021f29c9083cd1808ab36b740208524d9f63/deep_ep/buffer.py#L164
71
self.available_rank_configs = [2, 4, 8, 16, 24, 32, 64, 128, 144, 160]
72
73
+ def _sync_dbo_comm_if_needed(self) -> None:
74
+ if self.sync_dbo_comm and dbo_enabled():
75
+ # ROCm DeepEP HT dispatch/combine reuse Buffer-owned communication
76
+ # workspace. Do not let the next DBO ubatch reuse that workspace
77
+ # before this ubatch's HT kernel has completed.
78
+ torch.cuda.current_stream().synchronize()
79
+
80
def num_dispatchers(self) -> int:
81
return self.num_dispatchers_
82
@@ -161,6 +170,8 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
170
allocate_on_comm_stream=False,
171
)
172
173
+ self._sync_dbo_comm_if_needed()
174
+
175
# record the handle for this ubatch
176
a2a_idx = dbo_current_ubatch_id()
177
self.handles[a2a_idx] = handle
@@ -375,6 +386,8 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
386
allocate_on_comm_stream=False,
387
)
388
389
+ self._sync_dbo_comm_if_needed()
390
+
391
dbo_switch_to_compute()
392
393
if do_async:
vllm/v1/worker/gpu_ubatch_wrapper.py
+10
@@ -154,6 +154,16 @@ class UBatchWrapper:
154
@staticmethod
155
def _create_sm_control_context(vllm_config: VllmConfig):
156
comm_sms: int = envs.VLLM_DBO_COMM_SMS
157
+ rocm_deepep_ht_dbo = (
158
+ current_platform.is_rocm()
159
+ and vllm_config.parallel_config.enable_dbo
160
+ and vllm_config.parallel_config.all2all_backend == "deepep_high_throughput"
161
+ )
162
+ if rocm_deepep_ht_dbo:
163
+ # On ROCm, reserving CUs for DeepEP HT communication under DBO
164
+ # corrupts DP+EP generation accuracy. Keep the backend active, but
165
+ # leave all CUs visible to the compute and communication kernels.
166
+ comm_sms = 0
167
168
set_comm_sms = lambda sms: None
169
if vllm_config.parallel_config.enable_expert_parallel: