diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index b093b1788a..b36c02a48e 100644 --- a/vllm/config/vllm.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py index 45f9e815ac..f30e0bea3c 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py +++ b/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: diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 657fc82673..76fa12b412 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/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: