[ROCm][Perf] Enable DPA+ETP prefill compute/comm dual-batch-overlap for DeepSeek-V4 (#57773)

Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Signed-off-by: shen-shanshan <467638484@qq.com>

tests/kernels/moe/test_naive_dp_ep_dbo.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_moe.prepare_finalize import naive_dp_ep
+from vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep import (
+    MoEPrepareAndFinalizeNaiveDPEPModular,
+    make_moe_prepare_and_finalize_naive_dp_ep,
+)
+
+
+class _RegionProbe(MoEPrepareAndFinalizeNaiveDPEPModular):
+    def __init__(self) -> None:
+        super().__init__()
+        self.body_ran = False
+
+    def run_comm(self) -> None:
+        with self._comm_region():
+            self.body_ran = True
+
+
+@pytest.fixture
+def dbo_mocks():
+    yield_switch = MagicMock()
+    switch_back = MagicMock()
+    with (
+        patch.object(
+            naive_dp_ep, "dbo_yield_and_switch_from_compute_to_comm", yield_switch
+        ),
+        patch.object(naive_dp_ep, "dbo_switch_to_compute_sync", switch_back),
+    ):
+        yield {"yield_switch": yield_switch, "switch_back": switch_back}
+
+
+def _dbo_state(dbo_active: bool, capturing: bool):
+    return (
+        patch.object(naive_dp_ep, "dbo_enabled", return_value=dbo_active),
+        patch.object(torch.cuda, "is_current_stream_capturing", return_value=capturing),
+    )
+
+
+def test_comm_region_noop_without_dbo(dbo_mocks):
+    probe = _RegionProbe()
+    dbo, capturing = _dbo_state(dbo_active=False, capturing=False)
+    with dbo, capturing:
+        probe.run_comm()
+    assert probe.body_ran
+    dbo_mocks["yield_switch"].assert_not_called()
+    dbo_mocks["switch_back"].assert_not_called()
+
+
+def test_comm_region_handoff_when_dbo_active(dbo_mocks):
+    probe = _RegionProbe()
+    dbo, capturing = _dbo_state(dbo_active=True, capturing=False)
+    with dbo, capturing:
+        probe.run_comm()
+    assert probe.body_ran
+    dbo_mocks["yield_switch"].assert_called_once_with()
+    dbo_mocks["switch_back"].assert_called_once_with()
+
+
+def test_comm_region_skipped_during_graph_capture(dbo_mocks):
+    probe = _RegionProbe()
+    dbo, capturing = _dbo_state(dbo_active=True, capturing=True)
+    with dbo, capturing:
+        probe.run_comm()
+    assert probe.body_ran
+    dbo_mocks["yield_switch"].assert_not_called()
+    dbo_mocks["switch_back"].assert_not_called()
+
+
+def test_comm_region_restores_compute_stream_on_error(dbo_mocks):
+    probe = _RegionProbe()
+    dbo, capturing = _dbo_state(dbo_active=True, capturing=False)
+    with (
+        dbo,
+        capturing,
+        pytest.raises(RuntimeError, match="collective failed"),
+        probe._comm_region(),
+    ):
+        raise RuntimeError("collective failed")
+    dbo_mocks["yield_switch"].assert_called_once_with()
+    dbo_mocks["switch_back"].assert_called_once_with()
+
+
+@pytest.mark.parametrize("is_rocm", [True, False])
+def test_supports_dbo_only_on_rocm(is_rocm):
+    inst = MoEPrepareAndFinalizeNaiveDPEPModular()
+    platform = SimpleNamespace(is_rocm=lambda: is_rocm)
+    with patch.object(naive_dp_ep, "current_platform", platform):
+        assert inst.supports_dbo() is is_rocm
+        # supports_async() also drives MK-internal shared-expert overlap, so
+        # it must not change with DBO support.
+        assert not inst.supports_async()
+
+
+@pytest.mark.parametrize("is_rocm", [True, False])
+def test_factory_returns_base_classes(is_rocm):
+    platform = SimpleNamespace(is_rocm=lambda: is_rocm)
+    with patch.object(naive_dp_ep, "current_platform", platform):
+        mono = make_moe_prepare_and_finalize_naive_dp_ep(use_monolithic=True)
+        modular = make_moe_prepare_and_finalize_naive_dp_ep(use_monolithic=False)
+    assert type(mono) is naive_dp_ep.MoEPrepareAndFinalizeNaiveDPEPMonolithic
+    assert type(modular) is MoEPrepareAndFinalizeNaiveDPEPModular
+
+
+def test_default_supports_dbo_follows_supports_async():
+    # Classes that do not override supports_dbo() keep the old contract:
+    # DBO only with async prepare/finalize.
+    mono = naive_dp_ep.MoEPrepareAndFinalizeNaiveDPEPMonolithic()
+    assert not mono.supports_async()
+    assert not mono.supports_dbo()
+    with patch.object(type(mono), "supports_async", return_value=True):
+        assert mono.supports_dbo()

tests/v1/worker/test_ubatch_offsets.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import threading
+from unittest.mock import patch
+
+import numpy as np
+
+from vllm.v1.worker import ubatching
+from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
+from vllm.v1.worker.ubatching import (
+    dbo_current_ubatch_token_offset,
+    make_ubatch_contexts,
+)
+
+
+def test_token_offset_without_dbo_context():
+    assert dbo_current_ubatch_token_offset() == 0
+
+
+def test_token_offset_defaults_to_zero():
+    ready_barrier = threading.Barrier(3)
+    ctxs = make_ubatch_contexts(
+        num_micro_batches=2,
+        compute_stream=object(),
+        comm_stream=object(),
+        forward_contexts=[object(), object()],
+        ready_barrier=ready_barrier,
+    )
+    assert [ctx.token_offset for ctx in ctxs] == [0, 0]
+
+
+def test_token_offset_per_ubatch_thread():
+    compute_stream = object()
+    ready_barrier = threading.Barrier(3)
+    ctxs = make_ubatch_contexts(
+        num_micro_batches=2,
+        compute_stream=compute_stream,
+        comm_stream=object(),
+        forward_contexts=[object(), object()],
+        ready_barrier=ready_barrier,
+        token_offsets=[0, 7],
+    )
+    assert [ctx.token_offset for ctx in ctxs] == [0, 7]
+
+    offsets_seen: dict[int, int] = {}
+
+    def run(ctx: ubatching.UBatchContext) -> None:
+        with ctx:
+            offsets_seen[ctx.id] = dbo_current_ubatch_token_offset()
+
+    # Keep torch.cuda.set_stream out of the way. Patch once here rather than
+    # per thread: a per-thread patch exit would restore the real function
+    # while the other thread is still inside its context.
+    with patch.object(ubatching, "current_stream", return_value=compute_stream):
+        threads = [threading.Thread(target=run, args=(ctx,)) for ctx in ctxs]
+        for thread in threads:
+            thread.start()
+        ready_barrier.wait()
+        ctxs[0].cpu_wait_event.set()
+        for thread in threads:
+            thread.join()
+
+    assert offsets_seen == {0: 0, 1: 7}
+
+
+def test_ubatch_slice_bounds_are_python_ints():
+    # The last bound comes from np.cumsum; numpy scalars leaking into the
+    # per-ubatch metadata break Triton argument specialization.
+    ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
+        should_ubatch=True,
+        num_scheduled_tokens=np.array([5, 9, 6], dtype=np.int32),
+        num_tokens_padded=20,
+        num_reqs_padded=3,
+        num_ubatches=2,
+    )
+    for slices in (ubatch_slices, ubatch_slices_padded):
+        assert slices is not None
+        for s in slices:
+            for bound in (s.token_slice.start, s.token_slice.stop):
+                assert type(bound) is int

vllm/config/vllm.py

@@ -2116,17 +2116,28 @@ class VllmConfig:
 
         if self.parallel_config.use_ubatching:
             a2a_backend = self.parallel_config.all2all_backend
-            assert a2a_backend in [
-                "deepep_low_latency",
-                "deepep_high_throughput",
-                "nixl_ep",
-            ], (
-                "Microbatching currently only supports the deepep_low_latency, "
-                "deepep_high_throughput, and nixl_ep all2all backends. "
-                f"{a2a_backend} is not supported. To fix use "
-                "--all2all-backend=deepep_low_latency, "
-                "--all2all-backend=deepep_high_throughput, or "
-                "--all2all-backend=nixl_ep and install the matching kernels."
+            if current_platform.is_rocm():
+                supported_backends = [
+                    "deepep_low_latency",
+                    "deepep_high_throughput",
+                    "nixl_ep",
+                    # The naive DP/EP prepare/finalize runs its all-gather /
+                    # reduce-scatter on the comm stream under DBO.
+                    "allgather_reducescatter",
+                ]
+            else:
+                # CUDA and other platforms.
+                supported_backends = [
+                    "deepep_low_latency",
+                    "deepep_high_throughput",
+                    "nixl_ep",
+                ]
+            assert a2a_backend in supported_backends, (
+                "Microbatching currently only supports the "
+                f"{', '.join(supported_backends)} all2all backends on this "
+                f"platform. {a2a_backend} is not supported. To fix use one of "
+                f"{', '.join('--all2all-backend=' + b for b in supported_backends)}"
+                " and install the matching kernels."
             )
 
         if not self.instance_id:

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -258,6 +258,15 @@ class FusedMoEPrepareAndFinalize(ABC):
         """
         return False
 
+    def supports_dbo(self) -> bool:
+        """Indicates whether or not this class can run under DBO.
+
+        Defaults to supports_async(). A class whose synchronous prepare and
+        finalize do their own DBO compute/comm handoff can return True here
+        while keeping supports_async() False.
+        """
+        return self.supports_async()
+
     def supports_deferred_moe_finalize(self) -> bool:
         """Whether ``finalize`` can be skipped for a deferring consumer.
 
@@ -1206,7 +1215,7 @@ class FusedMoEKernelModularImpl:
             # We shouldn't be running an a2a kernel that doesn't
             # support async prepare/finalize
             # TODO(lucas): enable in follow-up
-            assert not dbo_enabled()
+            assert not dbo_enabled() or self.prepare_finalize.supports_dbo()
 
             (
                 a1q,
@@ -1394,7 +1403,7 @@ class FusedMoEKernelModularImpl:
 
         """
         if not self.prepare_finalize.supports_async():
-            assert not dbo_enabled()
+            assert not dbo_enabled() or self.prepare_finalize.supports_dbo()
 
             self.prepare_finalize.finalize(
                 output,

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

@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterator
+from contextlib import contextmanager
+
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -11,7 +14,13 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
+from vllm.platforms import current_platform
 from vllm.utils.flashinfer import nvfp4_block_scale_interleave
+from vllm.v1.worker.ubatching import (
+    dbo_enabled,
+    dbo_switch_to_compute_sync,
+    dbo_yield_and_switch_from_compute_to_comm,
+)
 
 
 def _quantize_and_setup_dispatch(
@@ -133,6 +142,34 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
             return self._allocate_combine_input(shape, dtype, device)
         return None
 
+    def supports_dbo(self) -> bool:
+        # prepare()/finalize() hand the collective to the comm stream
+        # themselves (see _comm_region), so the synchronous path is DBO-safe.
+        # supports_async() stays False: it also selects MK-internal
+        # shared-expert overlap. Only validated on ROCm so far.
+        return current_platform.is_rocm()
+
+    @contextmanager
+    def _comm_region(self) -> Iterator[None]:
+        """Scope the dispatch/combine collective runs in.
+
+        Without DBO the collective runs on the calling stream. Under DBO it
+        runs on the comm stream so the peer ubatch can compute while it is in
+        flight, using the same compute/comm handoff as DeepEP-HT. Graph
+        capture keeps every collective on the capture stream, so overlap is
+        limited to eager steps (prefill).
+        """
+        if not dbo_enabled() or torch.cuda.is_current_stream_capturing():
+            yield
+            return
+        # Release the CPU to the peer ubatch so it can enqueue compute while
+        # this ubatch's collective runs; the compute stream waits on exit.
+        dbo_yield_and_switch_from_compute_to_comm()
+        try:
+            yield
+        finally:
+            dbo_switch_to_compute_sync()
+
     def prepare(
         self,
         a1: torch.Tensor,
@@ -178,13 +215,14 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
                 extra_tensors = []
             extra_tensors.append(local_token_lora_mapping)
 
-        res = get_ep_group().dispatch(
-            a1q,
-            topk_weights,
-            topk_ids,
-            is_sequence_parallel=self.is_sequence_parallel,
-            extra_tensors=extra_tensors,
-        )
+        with self._comm_region():
+            res = get_ep_group().dispatch(
+                a1q,
+                topk_weights,
+                topk_ids,
+                is_sequence_parallel=self.is_sequence_parallel,
+                extra_tensors=extra_tensors,
+            )
 
         if extra_tensors is None:
             assert len(res) == 3
@@ -238,11 +276,12 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
             apply_router_weight_on_input=apply_router_weight_on_input,
         )
 
-        get_ep_group().combine_into_output(
-            out,
-            output,
-            is_sequence_parallel=self.is_sequence_parallel,
-        )
+        with self._comm_region():
+            get_ep_group().combine_into_output(
+                out,
+                output,
+                is_sequence_parallel=self.is_sequence_parallel,
+            )
 
 
 class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMonolithic):

vllm/models/deepseek_v4/amd/model.py

@@ -82,6 +82,7 @@ from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
 from vllm.platforms import current_platform
 from vllm.platforms.rocm import on_gfx950
 from vllm.sequence import IntermediateTensors
+from vllm.v1.worker.ubatching import dbo_current_ubatch_token_offset
 
 from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
 
@@ -1228,7 +1229,13 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
 
         if self._mtp_hidden_buffer is not None:
             num_tokens = hidden_states.shape[0]
-            self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
+            # Each DBO ubatch owns a different slice of the batch. Writing
+            # from row 0 in both threads would let one ubatch clobber the
+            # other's rows in the buffer the MTP drafter reads after the step.
+            offset = dbo_current_ubatch_token_offset()
+            self._mtp_hidden_buffer[offset : offset + num_tokens].copy_(
+                hidden_states.flatten(1)
+            )
 
         hidden_states = self.hc_head_op(
             hidden_states,

vllm/v1/worker/gpu_ubatch_wrapper.py

@@ -213,6 +213,7 @@ class UBatchWrapper:
     def _run_ubatches(self, ubatch_metadata, model) -> torch.Tensor:
         @torch.inference_mode()
         def _ubatch_thread(results, model, ubatch_metadata):
+            torch.accelerator.set_device_index(self.device)
             with ubatch_metadata.context:
                 model_output = model(
                     input_ids=ubatch_metadata.input_ids,
@@ -291,6 +292,7 @@ class UBatchWrapper:
             compute_stream=compute_stream,
             forward_contexts=forward_contexts,
             ready_barrier=self.ready_barrier,
+            token_offsets=[s.token_slice.start for s in ubatch_slices],
         )
 
         ubatch_metadata: list[UbatchMetadata] = []

vllm/v1/worker/ubatch_utils.py

@@ -176,7 +176,12 @@ def maybe_create_ubatch_slices(
     if split_point is None:
         split_point = int(num_tokens_padded) // num_ubatches
 
-    token_split_points = [split_point * i for i in range(1, num_ubatches)]
+    if isinstance(split_point, int):
+        token_split_points = [split_point * i for i in range(1, num_ubatches)]
+    else:
+        # Explicit split points, one per ubatch boundary.
+        assert len(split_point) == num_ubatches - 1
+        token_split_points = [int(p) for p in split_point]
 
     # TODO(lucas): Refactor the gpu_model_runner.py so we can pass
     # in cu_num_tokens directly (i.e. query_start_loc)
@@ -187,7 +192,7 @@ def maybe_create_ubatch_slices(
     start_token = 0
 
     # Add the end point to the split points to make iteration easier
-    all_points = token_split_points + [cu_num_tokens[-1]]
+    all_points = token_split_points + [int(cu_num_tokens[-1])]
 
     for end_token in all_points:
         token_slice = slice(start_token, end_token)
@@ -316,6 +321,10 @@ def _make_metadata_with_slice(
     block_table_tensor = attn_metadata.block_table_tensor[request_slice]
     slot_mapping = attn_metadata.slot_mapping[token_slice]
 
+    positions = attn_metadata.positions
+    if positions is not None:
+        positions = positions[token_slice]
+
     return CommonAttentionMetadata(
         query_start_loc=query_start_loc,
         query_start_loc_cpu=query_start_loc_cpu,
@@ -326,6 +335,7 @@ def _make_metadata_with_slice(
         max_seq_len=max_seq_len,
         block_table_tensor=block_table_tensor,
         slot_mapping=slot_mapping,
+        positions=positions,
         seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
     )
 

vllm/v1/worker/ubatching.py

@@ -32,6 +32,7 @@ class UBatchContext:
         gpu_comm_done_event: torch.Event,
         gpu_compute_done_event: torch.Event,
         schedule: str = "default",
+        token_offset: int = 0,
     ):
         self.id = id
         self.comm_stream = comm_stream
@@ -44,6 +45,7 @@ class UBatchContext:
         self.gpu_comm_done_event = gpu_comm_done_event
         self.gpu_compute_done_event = gpu_compute_done_event
         self.schedule = schedule
+        self.token_offset = token_offset
         self.recv_hook = None
 
     def __enter__(self):
@@ -155,6 +157,17 @@ def dbo_current_ubatch_id() -> int:
     return _THREAD_ID_TO_CONTEXT[threading.get_ident()]
 
 
+def dbo_current_ubatch_token_offset() -> int:
+    """Row offset of the calling ubatch inside the parent batch.
+
+    Returns 0 when no DBO context is active on this thread.
+    """
+    if len(_THREAD_ID_TO_CONTEXT) == 0:
+        return 0
+    ctx = _CURRENT_CONTEXTS[_THREAD_ID_TO_CONTEXT[threading.get_ident()]]
+    return ctx.token_offset
+
+
 def _register_ubatch_function(func):
     def wrapper(*args, **kwargs):
         if len(_THREAD_ID_TO_CONTEXT) > 0:
@@ -204,6 +217,7 @@ def make_ubatch_contexts(
     forward_contexts: list[ForwardContext],
     ready_barrier: threading.Barrier,
     schedule: str = "default",
+    token_offsets: list[int] | None = None,
 ) -> list[UBatchContext]:
     global _NUM_UBATCHES, _CURRENT_CONTEXTS
     assert num_micro_batches > 1, "num_micro_batches must be greater than 1"
@@ -233,6 +247,7 @@ def make_ubatch_contexts(
             gpu_comm_done_event=gpu_comm_done_events[i],
             gpu_compute_done_event=gpu_compute_done_events[i],
             schedule=schedule,
+            token_offset=token_offsets[i] if token_offsets else 0,
         )
         ctxs.append(ctx)