diff --git a/tests/kernels/moe/test_naive_dp_ep_dbo.py b/tests/kernels/moe/test_naive_dp_ep_dbo.py new file mode 100644 index 0000000000..c83328aef9 --- /dev/null +++ b/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() diff --git a/tests/v1/worker/test_ubatch_offsets.py b/tests/v1/worker/test_ubatch_offsets.py new file mode 100644 index 0000000000..39de6c5537 --- /dev/null +++ b/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 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 06614f8326..87d2187459 100644 --- a/vllm/config/vllm.py +++ b/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: diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 158e144e72..785369a344 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/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, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py index 7087b9016d..86168f683b 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py +++ b/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): diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 76db8d132f..2facab7230 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/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, diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index c6b380d63e..816df53776 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/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] = [] diff --git a/vllm/v1/worker/ubatch_utils.py b/vllm/v1/worker/ubatch_utils.py index 329c261a0a..85e83489c9 100644 --- a/vllm/v1/worker/ubatch_utils.py +++ b/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, ) diff --git a/vllm/v1/worker/ubatching.py b/vllm/v1/worker/ubatching.py index a773ee715d..a39f85d388 100644 --- a/vllm/v1/worker/ubatching.py +++ b/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)