[Perf][MoE] Eliminate staging from NCCL symmetric reduce-scatter (#49194)

Signed-off-by: samnordmann <snordmann@nvidia.com>
Co-authored-by: Misha Goin <mgoin64@gmail.com>

tests/distributed/test_mnnvl_alltoall.py

@@ -150,7 +150,9 @@ def _init_dp_environment(world_size, rank, port, dp_size, dp_port):
 
 
 @contextmanager
-def _make_forward_context(rank, world_size, num_tokens_per_rank):
+def _make_forward_context(
+    rank, world_size, num_tokens_per_rank, num_tokens_across_dp=None
+):
     """Keep the vLLM config and mock-DP forward context active together.
 
     Returns a context manager suitable for ``with`` statements.
@@ -175,15 +177,15 @@ def _make_forward_context(rank, world_size, num_tokens_per_rank):
         is_moe_model=True,
         data_parallel_rank=rank,
     )
+    if num_tokens_across_dp is None:
+        num_tokens_across_dp = [num_tokens_per_rank] * world_size
     with (
         set_current_vllm_config(vllm_config),
         set_forward_context(
             _AttnMeta(),
             vllm_config,
             num_tokens=num_tokens_per_rank,
-            num_tokens_across_dp=torch.tensor(
-                [num_tokens_per_rank] * world_size, dtype=torch.int
-            ),
+            num_tokens_across_dp=torch.tensor(num_tokens_across_dp, dtype=torch.int),
         ),
     ):
         yield
@@ -516,6 +518,121 @@ def test_one_sided_manager_workspace_grow(world_size):
 # ---------------------------------------------------------------------------
 
 
+class _FakeCombineGroup:
+    def __init__(self, combine_input):
+        self.combine_input = combine_input
+        self.allocation_calls = 0
+        self.combined_input = None
+
+    def allocate_combine_input(self, shape, dtype, device, **kwargs):
+        self.allocation_calls += 1
+        assert self.combine_input is not None
+        assert self.combine_input.shape == shape
+        assert self.combine_input.dtype == dtype
+        assert self.combine_input.device == device
+        return self.combine_input
+
+    def combine_into_output(self, hidden_states, output, **kwargs):
+        self.combined_input = hidden_states
+        output.copy_(hidden_states)
+        return output
+
+
+def test_args_cpu_workspace_allocation(monkeypatch):
+    from types import SimpleNamespace
+
+    import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+    from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+
+    kernel = object.__new__(mk.FusedMoEKernelModularImpl)
+    kernel.fused_experts = SimpleNamespace(
+        workspace_dtype=lambda dtype: dtype,
+        workspace_shapes=lambda *args: ((0,), (0,), (2, 3)),
+        finalize_weight_and_reduce_impl=lambda: None,
+    )
+    kernel.prepare_finalize = SimpleNamespace(
+        allocate_fused_expert_output=lambda *args: None,
+    )
+    monkeypatch.setattr(mk.current_platform, "is_cpu", lambda: True)
+
+    workspace13, workspace2, output, specialized = kernel._allocate_buffers(
+        torch.float32,
+        torch.device("cpu"),
+        2,
+        2,
+        3,
+        3,
+        1,
+        1,
+        1,
+        None,
+        MoEActivation.SILU,
+    )
+    assert workspace13.numel() == workspace2.numel() == 0
+    assert output.shape == (2, 3)
+    assert output.device.type == "cpu"
+    assert not specialized
+
+
+def test_args_finalize_noop_reuses_fused_expert_output(monkeypatch):
+    import vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep as pf
+    from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+        TopKWeightAndReduceNoOP,
+    )
+
+    fused_expert_output = torch.arange(6, dtype=torch.float32).view(2, 3)
+    output = torch.empty_like(fused_expert_output)
+    group = _FakeCombineGroup(combine_input=None)
+    monkeypatch.setattr(pf, "get_ep_group", lambda: group)
+
+    impl = pf.MoEPrepareAndFinalizeNaiveDPEPModular()
+    impl.finalize(
+        output,
+        fused_expert_output,
+        torch.ones((2, 1)),
+        torch.zeros((2, 1), dtype=torch.long),
+        False,
+        TopKWeightAndReduceNoOP(),
+    )
+
+    assert group.allocation_calls == 0
+    assert group.combined_input is fused_expert_output
+    torch.testing.assert_close(output, fused_expert_output)
+
+
+def test_args_finalize_contiguous_writes_combine_input(monkeypatch):
+    import vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep as pf
+    import vllm.model_executor.layers.fused_moe.topk_weight_and_reduce as wr
+
+    fused_expert_output = torch.arange(12, dtype=torch.float32).view(2, 2, 3)
+    original = fused_expert_output.clone()
+    topk_weights = torch.tensor([[0.25, 0.75], [0.4, 0.6]])
+    combine_input = torch.empty((2, 3), dtype=torch.float32)
+    output = torch.empty_like(combine_input)
+    group = _FakeCombineGroup(combine_input)
+    monkeypatch.setattr(pf, "get_ep_group", lambda: group)
+    monkeypatch.setattr(
+        wr.ops,
+        "moe_sum",
+        lambda hidden_states, out: out.copy_(hidden_states.sum(dim=1)),
+    )
+
+    impl = pf.MoEPrepareAndFinalizeNaiveDPEPModular()
+    impl.finalize(
+        output,
+        fused_expert_output,
+        topk_weights,
+        torch.zeros((2, 2), dtype=torch.long),
+        False,
+        wr.TopKWeightAndReduceContiguous(),
+    )
+
+    expected = (original * topk_weights.unsqueeze(-1)).sum(dim=1)
+    assert group.allocation_calls == 1
+    assert group.combined_input is combine_input
+    torch.testing.assert_close(output, expected)
+
+
 def _args_dispatch_combine_worker(rank, world_size):
     from vllm.distributed.device_communicators.all2all import AgRsAll2AllManager
     from vllm.forward_context import get_forward_context
@@ -604,14 +721,39 @@ def _args_dispatch_combine_worker(rank, world_size):
             # -- combine (reduce-scatter) --
             # Each token i has value i in all columns; after reduce-scatter
             # each rank gets its slice, summed across ranks.
-            expert_out = (
+            expert_values = (
                 torch.arange(total_tokens, device=device, dtype=torch.float32)
                 .unsqueeze(1)
                 .expand(total_tokens, hidden_size)
                 .contiguous()
             )
+            expert_out = manager.allocate_combine_input(
+                expert_values.shape,
+                expert_values.dtype,
+                expert_values.device,
+                is_sequence_parallel=True,
+            )
+            if expert_out is None:
+                expert_out = expert_values
+            else:
+                from vllm.distributed.device_communicators.pynccl_allocator import (
+                    is_symmetric_memory_tensor,
+                )
 
-            combined = manager.combine(expert_out, is_sequence_parallel=True)
+                assert is_symmetric_memory_tensor(expert_out)
+                expert_out.copy_(expert_values)
+
+            combine_output = torch.empty(
+                (tokens_per_rank, hidden_size),
+                device=device,
+                dtype=expert_out.dtype,
+            )
+            combined = manager.combine_into_output(
+                expert_out,
+                combine_output,
+                is_sequence_parallel=True,
+            )
+            assert combined.data_ptr() == combine_output.data_ptr()
             assert combined.shape == (tokens_per_rank, hidden_size)
 
             for i in range(tokens_per_rank):
@@ -623,6 +765,61 @@ def _args_dispatch_combine_worker(rank, world_size):
 
             torch.distributed.barrier()
 
+    uneven_sizes = [tokens_per_rank + r for r in range(world_size)]
+    with _make_forward_context(
+        rank,
+        world_size,
+        uneven_sizes[rank],
+        num_tokens_across_dp=uneven_sizes,
+    ):
+        dp_metadata = get_forward_context().dp_metadata
+        assert dp_metadata is not None
+        with dp_metadata.sp_local_sizes(sequence_parallel_size=1):
+            assert (
+                manager.allocate_combine_input(
+                    (sum(uneven_sizes), hidden_size),
+                    torch.float32,
+                    device,
+                    is_sequence_parallel=True,
+                )
+                is None
+            )
+
+            total_uneven_tokens = sum(uneven_sizes)
+            uneven_input = (
+                torch.arange(
+                    total_uneven_tokens,
+                    dtype=torch.float32,
+                    device=device,
+                )
+                .unsqueeze(1)
+                .expand(total_uneven_tokens, hidden_size)
+                .contiguous()
+            )
+            uneven_input.add_(rank)
+            uneven_output = torch.empty(
+                (uneven_sizes[rank], hidden_size),
+                dtype=torch.float32,
+                device=device,
+            )
+            combined = get_ep_group().reduce_scatterv_into_output(
+                uneven_input,
+                uneven_output,
+                dim=0,
+                sizes=uneven_sizes,
+            )
+            assert combined.data_ptr() == uneven_output.data_ptr()
+
+            start = sum(uneven_sizes[:rank])
+            expected = torch.arange(
+                start,
+                start + uneven_sizes[rank],
+                dtype=torch.float32,
+                device=device,
+            ) * world_size + sum(range(world_size))
+            expected = expected.unsqueeze(1).expand_as(combined)
+            torch.testing.assert_close(combined, expected)
+
 
 @requires_multi_gpu
 @pytest.mark.parametrize("world_size", [2])

tests/distributed/test_nccl_symm_mem.py

@@ -3,6 +3,7 @@
 
 import random
 import typing
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -12,11 +13,15 @@ import torch.multiprocessing as mp
 import vllm.envs as envs
 from tests.utils import ensure_current_vllm_config
 from vllm.distributed import cleanup_dist_env_and_memory
-from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
+from vllm.distributed.device_communicators.cuda_communicator import (
+    NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION,
+    CudaCommunicator,
+)
 from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
 from vllm.distributed.device_communicators.pynccl_allocator import (
     get_nccl_mem_pool,
     is_symmetric_memory_enabled,
+    is_symmetric_memory_tensor,
 )
 from vllm.distributed.parallel_state import (
     get_tp_group,
@@ -32,6 +37,53 @@ random.seed(44)
 test_size_elements = 4 * 1024 * 1024
 
 
+def test_disabled_pynccl_does_not_allocate_symmetric_buffer():
+    communicator = object.__new__(CudaCommunicator)
+    communicator.pynccl_comm = SimpleNamespace(disabled=True)
+
+    assert (
+        communicator.get_symmetric_memory_buffer(
+            "test", (1,), torch.float32, torch.device("cuda")
+        )
+        is None
+    )
+
+
+@pytest.mark.parametrize("use_aiter", [False, True])
+def test_reduce_scatterv_preserves_caller_output(monkeypatch, use_aiter):
+    import vllm.distributed.device_communicators.cuda_communicator as cc
+
+    communicator = object.__new__(CudaCommunicator)
+    communicator.world_size = 2
+    communicator.rank_in_group = 0
+    input_ = torch.arange(12, dtype=torch.float32).view(4, 3)
+    output = torch.full((2, 3), float("nan"))
+
+    def reduce_scatter(out, inp):
+        assert out is output
+        out.copy_(inp.view(2, 2, 3).sum(dim=0))
+
+    communicator.pynccl_comm = SimpleNamespace(
+        disabled=False, nccl_version=22902, reduce_scatter=reduce_scatter
+    )
+    communicator._can_use_aiter_ag_rs = lambda sizes: use_aiter and sizes is None
+    communicator.aiter_ar_comm = SimpleNamespace(
+        should_custom_rs=lambda inp, dim: True,
+        custom_reduce_scatter=lambda inp, out, dim: reduce_scatter(out, inp),
+    )
+    monkeypatch.setattr(cc, "should_nccl_symm_mem_ag_rs", lambda: True)
+    monkeypatch.setattr(cc, "is_symmetric_memory_tensor", lambda tensor: False)
+    monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0")
+
+    # Explicit uniform sizes must not stage an ordinary input into symmetric
+    # scratch, and AITER must write the supplied output after normalization.
+    result = communicator.reduce_scatterv_into_output(
+        input_, output, dim=0, sizes=[2, 2]
+    )
+    assert result.data_ptr() == output.data_ptr()
+    torch.testing.assert_close(output, input_.view(2, 2, 3).sum(dim=0))
+
+
 def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
     monkeypatch = pytest.MonkeyPatch()
     with monkeypatch.context() as m:
@@ -193,15 +245,107 @@ def nccl_symm_mem_reduce_scatter_worker(local_rank: int, world_size: int):
             pytest.skip("NCCL symmetric memory is disabled.")
 
         per_rank_size = test_size_elements // world_size
-        input_tensor = torch.randint(
+        ordinary_input = torch.randint(
             1, 23, (test_size_elements,), dtype=dtype, device=device
         )
-        input_clone = input_tensor.clone()
-        output = cuda_communicator.reduce_scatter(input_tensor, dim=0)
+        ordinary_clone = ordinary_input.clone()
+        ordinary_output = cuda_communicator.reduce_scatter(ordinary_input, dim=0)
 
         group = get_tp_group().device_group
-        expected = torch.empty(per_rank_size, dtype=dtype, device=device)
+        ordinary_expected = torch.empty(per_rank_size, dtype=dtype, device=device)
+        dist.reduce_scatter_tensor(ordinary_expected, ordinary_clone, group=group)
+        torch.testing.assert_close(
+            ordinary_output, ordinary_expected, atol=2.5, rtol=0.1
+        )
+        assert is_symmetric_memory_tensor(ordinary_output)
+
+        # Both calls reuse rs_out; validate the first result before overwriting it.
+        for sizes in (None, [per_rank_size] * world_size):
+            ordinary_v_output = cuda_communicator.reduce_scatterv(
+                ordinary_input, dim=0, sizes=sizes
+            )
+            torch.testing.assert_close(
+                ordinary_v_output, ordinary_expected, atol=2.5, rtol=0.1
+            )
+            assert is_symmetric_memory_tensor(ordinary_v_output)
+
+        pynccl_comm = cuda_communicator.pynccl_comm
+        assert pynccl_comm is not None
+        if pynccl_comm.nccl_version < NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION:
+            return
+
+        from vllm.v1.worker import ubatching
+
+        m.setattr(ubatching, "dbo_current_ubatch_id", lambda: 0)
+        ubatch0 = cuda_communicator._get_symm_scratch("ag_out", (128,), dtype, device)
+        ubatch0.fill_(1)
+        m.setattr(ubatching, "dbo_current_ubatch_id", lambda: 1)
+        ubatch1 = cuda_communicator._get_symm_scratch("ag_out", (128,), dtype, device)
+        ubatch1.fill_(2)
+        assert ubatch0.data_ptr() != ubatch1.data_ptr()
+        assert (ubatch0 == 1).all()
+
+        m.setattr(ubatching, "dbo_current_ubatch_id", lambda: 0)
+        smaller = cuda_communicator._get_symm_scratch("ag_out", (64,), dtype, device)
+        assert smaller.shape == (64,)
+        assert smaller.data_ptr() == ubatch0.data_ptr()
+
+        grown = cuda_communicator._get_symm_scratch("ag_out", (129,), dtype, device)
+        assert grown.shape == (129,)
+        assert grown.data_ptr() != ubatch0.data_ptr()
+        reused = cuda_communicator._get_symm_scratch("ag_out", (200,), dtype, device)
+        assert reused.data_ptr() == grown.data_ptr()
+
+        input_tensor = cuda_communicator._get_symm_scratch(
+            "test_graph_rs_input", (test_size_elements,), dtype, device
+        )
+        input_tensor.random_(1, 23)
+        input_clone = input_tensor.clone()
+        output = torch.empty(per_rank_size, dtype=dtype, device=device)
+        assert is_symmetric_memory_tensor(input_tensor)
+        assert not is_symmetric_memory_tensor(output)
+
+        result = cuda_communicator.reduce_scatterv_into_output(
+            input_tensor,
+            output,
+            dim=0,
+            sizes=[per_rank_size] * world_size,
+        )
+        assert result.data_ptr() == output.data_ptr()
+
+        expected = torch.empty_like(output)
         dist.reduce_scatter_tensor(expected, input_clone, group=group)
+        torch.testing.assert_close(result, expected, atol=2.5, rtol=0.1)
+
+        dist.barrier(group=get_tp_group().cpu_group)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            graph_result = cuda_communicator.reduce_scatterv_into_output(
+                input_tensor,
+                output,
+                dim=0,
+                sizes=[per_rank_size] * world_size,
+            )
+        assert graph_result.data_ptr() == output.data_ptr()
+
+        input_tensor.random_(24, 47)
+        input_clone.copy_(input_tensor)
+        dist.reduce_scatter_tensor(expected, input_clone, group=group)
+        captured_input_ptr = input_tensor.data_ptr()
+        del input_tensor
+        grown_input = cuda_communicator._get_symm_scratch(
+            "test_graph_rs_input", (test_size_elements + 1,), dtype, device
+        )
+        assert grown_input.data_ptr() != captured_input_ptr
+        retired = cuda_communicator.__dict__["_retired_symm_scratch_bufs"]
+        assert any(
+            buf.data_ptr() == captured_input_ptr
+            for buffers in retired.values()
+            for buf in buffers
+        )
+        output.zero_()
+        graph.replay()
+        torch.accelerator.synchronize()
         torch.testing.assert_close(output, expected, atol=2.5, rtol=0.1)
 
 

vllm/distributed/device_communicators/all2all.py

@@ -171,6 +171,42 @@ class AgRsAll2AllManager(All2AllManagerBase):
         hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes)
         return hidden_states
 
+    def allocate_combine_input(
+        self,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+        is_sequence_parallel: bool = False,
+    ) -> torch.Tensor | None:
+        dist_group = self._get_comm_group(is_sequence_parallel)
+        sizes = self._get_sizes(shape[0] // dist_group.world_size, dist_group)
+        if sum(sizes) != shape[0] or any(size != sizes[0] for size in sizes):
+            return None
+        device_communicator = dist_group.device_communicator
+        if device_communicator is None:
+            return None
+        return device_communicator.get_symmetric_memory_buffer(
+            "moe_ag_rs_combine", shape, dtype, device
+        )
+
+    def combine_into_output(
+        self,
+        hidden_states: torch.Tensor,
+        output: torch.Tensor,
+        is_sequence_parallel: bool = False,
+    ) -> torch.Tensor:
+        dist_group = self._get_comm_group(is_sequence_parallel)
+        sizes = self._get_sizes(
+            hidden_states.shape[0] // dist_group.world_size,
+            dist_group,
+        )
+        return dist_group.reduce_scatterv_into_output(
+            hidden_states,
+            output,
+            dim=0,
+            sizes=sizes,
+        )
+
     def destroy(self):
         pass
 

vllm/distributed/device_communicators/base_device_communicator.py

@@ -157,6 +157,24 @@ class All2AllManagerBase:
     def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False):
         raise NotImplementedError
 
+    def allocate_combine_input(
+        self,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+        is_sequence_parallel: bool = False,
+    ) -> torch.Tensor | None:
+        return None
+
+    def combine_into_output(
+        self,
+        hidden_states: torch.Tensor,
+        output: torch.Tensor,
+        is_sequence_parallel: bool = False,
+    ) -> torch.Tensor:
+        output.copy_(self.combine(hidden_states, is_sequence_parallel))
+        return output
+
     def destroy(self):
         pass
 
@@ -315,6 +333,25 @@ class DeviceCommunicatorBase:
     ) -> torch.Tensor:
         raise NotImplementedError
 
+    def reduce_scatterv_into_output(
+        self,
+        input_: torch.Tensor,
+        output: torch.Tensor,
+        dim: int = -1,
+        sizes: list[int] | None = None,
+    ) -> torch.Tensor:
+        output.copy_(self.reduce_scatterv(input_, dim, sizes))
+        return output
+
+    def get_symmetric_memory_buffer(
+        self,
+        role: str,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        return None
+
     def gather(
         self, input_: torch.Tensor, dst: int = 0, dim: int = -1
     ) -> torch.Tensor | None:
@@ -417,5 +454,23 @@ class DeviceCommunicatorBase:
         """
         return hidden_states
 
+    def allocate_combine_input(
+        self,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+        is_sequence_parallel: bool = False,
+    ) -> torch.Tensor | None:
+        return None
+
+    def combine_into_output(
+        self,
+        hidden_states: torch.Tensor,
+        output: torch.Tensor,
+        is_sequence_parallel: bool = False,
+    ) -> torch.Tensor:
+        output.copy_(self.combine(hidden_states, is_sequence_parallel))
+        return output
+
     def batch_isend_irecv(self, p2p_ops: list):
         raise NotImplementedError

vllm/distributed/device_communicators/cuda_communicator.py

@@ -16,6 +16,7 @@ from vllm.distributed.device_communicators.all_reduce_utils import (
 from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
 from vllm.distributed.device_communicators.pynccl_allocator import (
     is_symmetric_memory_enabled,
+    is_symmetric_memory_tensor,
 )
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
@@ -26,6 +27,10 @@ from .base_device_communicator import DeviceCommunicatorBase
 
 logger = init_logger(__name__)
 
+# Since NCCL 2.29.2, ncclSymkPickKernel accepts registered input with
+# non-registered output for ReduceScatter (ncclSymSendRegRecvNonreg).
+NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION = 22902
+
 
 class CudaCommunicator(DeviceCommunicatorBase):
     def __init__(
@@ -505,13 +510,33 @@ class CudaCommunicator(DeviceCommunicatorBase):
     def reduce_scatterv(
         self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None
     ):
+        return self._reduce_scatterv(input_, dim, sizes)
+
+    def reduce_scatterv_into_output(
+        self,
+        input_: torch.Tensor,
+        output: torch.Tensor,
+        dim: int = -1,
+        sizes: list[int] | None = None,
+    ) -> torch.Tensor:
+        return self._reduce_scatterv(input_, dim, sizes, output)
+
+    def _reduce_scatterv(
+        self,
+        input_: torch.Tensor,
+        dim: int,
+        sizes: list[int] | None,
+        output: torch.Tensor | None = None,
+    ) -> torch.Tensor:
         world_size = self.world_size
         pynccl_comm = self.pynccl_comm
         assert pynccl_comm is not None
+        assert not pynccl_comm.disabled, "reduce_scatterv requires PyNccl"
         if dim < 0:
             # Convert negative dim to positive.
             dim += input_.dim()
 
+        sizes_are_explicit = sizes is not None
         # 'sizes' is not needed if all inputs in the same group have the same
         # shape
         if sizes is not None and all(s == sizes[0] for s in sizes):
@@ -529,35 +554,63 @@ class CudaCommunicator(DeviceCommunicatorBase):
             assert input_tensor.shape[0] % world_size == 0
             chunk_size = input_tensor.shape[0] // world_size
         output_shape = (chunk_size,) + input_tensor.shape[1:]
+        if output is not None:
+            assert dim == 0, "preallocated reduce-scatter output requires dim=0"
+            assert output.shape == output_shape
+            assert output.dtype == input_tensor.dtype
+            assert output.device == input_tensor.device
+            assert output.is_contiguous()
 
         if self._can_use_aiter_ag_rs(sizes):
             aiter_comm = self.aiter_ar_comm
             assert aiter_comm is not None
             if aiter_comm.should_custom_rs(input_tensor, dim=0):
-                output = torch.empty(
-                    output_shape, dtype=input_tensor.dtype, device=input_tensor.device
-                )
+                if output is None:
+                    output = torch.empty(
+                        output_shape,
+                        dtype=input_tensor.dtype,
+                        device=input_tensor.device,
+                    )
                 aiter_comm.custom_reduce_scatter(input_tensor, output, dim=0)
                 return output.movedim(0, dim).contiguous()
 
         # Symmetric memory is only used when all ranks have uniform sizes.
         # ncclCommWindowRegister is collective: asymmetric pool allocations
         # from variable per-rank sizes cause deadlocks.
-        use_symm_mem = sizes is None and should_nccl_symm_mem_ag_rs()
-        if use_symm_mem:
-            output = self._reduce_scatter_symm_mem(input_tensor)
-        else:
-            output = torch.empty(
-                output_shape, dtype=input_tensor.dtype, device=input_tensor.device
+        uniform_sizes = sizes is None or all(size == sizes[0] for size in sizes)
+        direct_output_supported = (
+            output is None
+            or is_symmetric_memory_tensor(output)
+            or pynccl_comm.nccl_version >= NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION
+        )
+        use_symm_mem = (
+            uniform_sizes
+            and direct_output_supported
+            and (
+                output is None
+                or not sizes_are_explicit
+                or is_symmetric_memory_tensor(input_tensor)
             )
+            and should_nccl_symm_mem_ag_rs()
+        )
+        if use_symm_mem:
+            output = self._reduce_scatter_symm_mem(input_tensor, output)
+        else:
+            if output is None:
+                output = torch.empty(
+                    output_shape, dtype=input_tensor.dtype, device=input_tensor.device
+                )
             use_deterministic_rs = envs.VLLM_BATCH_INVARIANT and world_size > 2
             if use_deterministic_rs:
                 # Reduce to a fixed root (0) for determinism
                 reduced = torch.empty_like(input_tensor)
-                sizes = sizes if sizes else [chunk_size] * world_size
+                scatter_sizes = sizes
+                if scatter_sizes is None:
+                    scatter_sizes = [chunk_size] * world_size
                 pynccl_comm.reduce(reduced, input_tensor, root=0)
-                pynccl_comm.scatter(output, reduced, sizes, root=0)
-            elif sizes is not None and sizes.count(sizes[0]) != len(sizes):
+                pynccl_comm.scatter(output, reduced, scatter_sizes, root=0)
+            elif not uniform_sizes:
+                assert sizes is not None
                 pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes)
             else:
                 pynccl_comm.reduce_scatter(output, input_tensor)
@@ -576,54 +629,64 @@ class CudaCommunicator(DeviceCommunicatorBase):
 
         Allocating a fresh symm tensor per collective pays the
         ``nccl_symm_mem_context`` snapshot + window-registration scan on every
-        call (~0.5 ms/RS+AG pair, dwarfing the NVLS transfer itself). Instead we
-        allocate once per ``(role, shape, dtype)``, register once, and reuse.
+        call (~0.5 ms/RS+AG pair, dwarfing the NVLS transfer itself). Instead,
+        each ``(role, shape[1:], dtype, device)`` keeps a geometric high-water
+        allocation and returns a leading slice. Superseded allocations remain
+        referenced because a captured CUDA graph may still use their pointers;
+        geometric growth bounds their total capacity to less than twice the
+        current allocation.
 
-        Safe for serial (eager) sequence parallelism: each collective's result
-        is consumed on the same stream before the next same-role collective
-        reuses the buffer. Distinct roles (e.g. ``rs_in`` vs ``ag_out``, both
-        full-size) get distinct buffers so a reduce-scatter input copy never
-        clobbers a still-live all-gather output.
+        Safe across serial MoE layers and eager sequence parallelism: the
+        producer and collective are ordered on the same stream before the next
+        same-role operation reuses the buffer. DBO microbatches use distinct
+        cache entries. Any future cross-layer communication overlap must also
+        use distinct roles.
         """
         from vllm.distributed.device_communicators.pynccl_allocator import (
             nccl_symm_mem_context,
         )
+        from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
         pynccl_comm = self.pynccl_comm
         assert pynccl_comm is not None
+        assert shape, "symmetric scratch buffers require at least one dimension"
         cache = self.__dict__.setdefault("_symm_scratch_bufs", {})
-        key = (role, tuple(shape), dtype)
+        key = (role, dbo_current_ubatch_id(), tuple(shape[1:]), dtype, device)
         buf = cache.get(key)
-        if buf is None:
+        requested_rows = shape[0]
+        if buf is None or buf.shape[0] < requested_rows:
+            capacity = (
+                requested_rows if buf is None else max(requested_rows, 2 * buf.shape[0])
+            )
             with nccl_symm_mem_context(pynccl_comm):
-                buf = torch.empty(shape, dtype=dtype, device=device)
+                new_buf = torch.empty(
+                    (capacity, *shape[1:]), dtype=dtype, device=device
+                )
+            if buf is not None:
+                retired = self.__dict__.setdefault("_retired_symm_scratch_bufs", {})
+                retired.setdefault(key, []).append(buf)
+            buf = new_buf
             cache[key] = buf
-        return buf
+        return buf[:requested_rows]
 
     def _reduce_scatter_symm_mem(
         self,
         input_tensor: torch.Tensor,
+        output: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """ReduceScatter using NCCL symmetric memory (NVLS).
 
-        Only called for uniform-size reduce_scatter (variable sizes are
-        guarded out by the caller to avoid asymmetric ncclCommWindowRegister).
-        Uses persistent pre-registered scratch (see _get_symm_scratch).
+        The MoE reduce_scatterv path passes explicit uniform sizes and calls
+        this only with an already-registered input, so that path never stages.
+        Uniform calls without a caller-owned output retain their existing
+        opt-in behavior and stage ordinary input into registered scratch.
+        The output may be ordinary memory.
         """
-        from vllm.distributed.device_communicators.pynccl_allocator import (
-            is_symmetric_memory_tensor,
-        )
-
         pynccl_comm = self.pynccl_comm
         assert pynccl_comm is not None
-
         chunk = input_tensor.shape[0] // self.world_size
         output_shape = (chunk,) + tuple(input_tensor.shape[1:])
 
-        symm_output = self._get_symm_scratch(
-            "rs_out", output_shape, input_tensor.dtype, input_tensor.device
-        )
-        # NVLS reduce-scatter (LDMC) requires the input in symmetric memory.
         if is_symmetric_memory_tensor(input_tensor):
             symm_input = input_tensor
         else:
@@ -635,8 +698,31 @@ class CudaCommunicator(DeviceCommunicatorBase):
             )
             symm_input.copy_(input_tensor)
 
-        pynccl_comm.reduce_scatter(symm_output, symm_input)
-        return symm_output
+        if output is None:
+            output = self._get_symm_scratch(
+                "rs_out", output_shape, input_tensor.dtype, input_tensor.device
+            )
+        pynccl_comm.reduce_scatter(output, symm_input)
+        return output
+
+    def get_symmetric_memory_buffer(
+        self,
+        role: str,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        pynccl_comm = self.pynccl_comm
+        if (
+            pynccl_comm is None
+            or pynccl_comm.disabled
+            or pynccl_comm.world_size == 1
+            or pynccl_comm.nccl_version < NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION
+            or not should_nccl_symm_mem_ag_rs()
+        ):
+            return None
+        output = self._get_symm_scratch(role, shape, dtype, device)
+        return output if is_symmetric_memory_tensor(output) else None
 
     def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
         """Sends a tensor to the destination rank in a blocking way."""
@@ -925,6 +1011,29 @@ class CudaCommunicator(DeviceCommunicatorBase):
             is_sequence_parallel,
         )
 
+    def allocate_combine_input(
+        self,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+        is_sequence_parallel: bool = False,
+    ) -> torch.Tensor | None:
+        assert self.all2all_manager is not None
+        return self.all2all_manager.allocate_combine_input(
+            shape, dtype, device, is_sequence_parallel
+        )
+
+    def combine_into_output(
+        self,
+        hidden_states: torch.Tensor,
+        output: torch.Tensor,
+        is_sequence_parallel: bool = False,
+    ) -> torch.Tensor:
+        assert self.all2all_manager is not None
+        return self.all2all_manager.combine_into_output(
+            hidden_states, output, is_sequence_parallel
+        )
+
     def batch_isend_irecv(self, p2p_ops: list):
         pynccl_comm = self.pynccl_comm
         if pynccl_comm is not None and not pynccl_comm.disabled:

vllm/distributed/parallel_state.py

@@ -801,6 +801,19 @@ class GroupCoordinator:
             raise ValueError("No device communicator found")
         return self.device_communicator.reduce_scatterv(input_, dim, sizes)
 
+    def reduce_scatterv_into_output(
+        self,
+        input_: torch.Tensor,
+        output: torch.Tensor,
+        dim: int = -1,
+        sizes: list[int] | None = None,
+    ) -> torch.Tensor:
+        if self.device_communicator is None:
+            raise ValueError("No device communicator found")
+        return self.device_communicator.reduce_scatterv_into_output(
+            input_, output, dim, sizes
+        )
+
     def _reduce_scatter_out_place(self, input_: torch.Tensor, dim: int) -> torch.Tensor:
         if self.device_communicator is None:
             raise ValueError("No device communicator found")
@@ -1442,6 +1455,32 @@ class GroupCoordinator:
         else:
             return hidden_states
 
+    def allocate_combine_input(
+        self,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+        is_sequence_parallel: bool = False,
+    ) -> torch.Tensor | None:
+        if self.device_communicator is None:
+            return None
+        return self.device_communicator.allocate_combine_input(
+            shape, dtype, device, is_sequence_parallel
+        )
+
+    def combine_into_output(
+        self,
+        hidden_states: torch.Tensor,
+        output: torch.Tensor,
+        is_sequence_parallel: bool = False,
+    ) -> torch.Tensor:
+        if self.device_communicator is None:
+            output.copy_(hidden_states)
+            return output
+        return self.device_communicator.combine_into_output(
+            hidden_states, output, is_sequence_parallel
+        )
+
 
 _WORLD: GroupCoordinator | None = None
 _INNER_DP_WORLD: GroupCoordinator | None = None

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -197,6 +197,20 @@ class FusedMoEPrepareAndFinalize(ABC):
         """
         return
 
+    def allocate_fused_expert_output(
+        self,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+        weight_and_reduce_impl: TopKWeightAndReduce,
+    ) -> torch.Tensor | None:
+        """Optionally provide the exact output buffer for the expert kernel.
+
+        Implementations returning a tensor must preserve its identity through
+        finalize; replacing it may silently reintroduce a write-back copy.
+        """
+        return None
+
     @property
     @abstractmethod
     def activation_format(self) -> FusedMoEActivationFormat:
@@ -1054,13 +1068,14 @@ class FusedMoEKernelModularImpl:
         local_num_experts: int,
         expert_tokens_meta: ExpertTokensMetadata | None,
         activation: MoEActivation,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]:
         """Allocate temporary and output buffers for the fused experts op.
         Inputs:
         - out_dtype: output type of workspace and output tensors.
         - device: the device of the workspace and output tensors.
         See `workspace_shapes` for a description of the remainder of arguments.
-        Returns a tuple of (workspace13, workspace2, output) tensors.
+        Returns the workspace tensors, output tensor, and whether the output
+        uses a prepare/finalize-provided allocation.
         """
         assert M_full > 0 and M_chunk > 0
 
@@ -1090,9 +1105,20 @@ class FusedMoEKernelModularImpl:
             activation,
         )
 
+        fused_out_alias = self.prepare_finalize.allocate_fused_expert_output(
+            fused_out_shape,
+            workspace_dtype,
+            device,
+            self.fused_experts.finalize_weight_and_reduce_impl(),
+        )
+
         # We can reuse the memory between cache1 and cache3 because by the
         # time we need cache3, we're done with cache1.
         # Reuse workspace13 for the output since there is only one chunk.
+        # Keep the original fused-output reservation even when the output is
+        # externally aliased. The workspace manager profiles and locks a
+        # process-wide high-water mark; later graph shapes may need this
+        # storage for workspace13.
         max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape))
 
         if current_platform.is_cpu():
@@ -1110,16 +1136,23 @@ class FusedMoEKernelModularImpl:
             )
             workspace13 = _resize_cache(common_workspace, workspace13_shape)
             fused_out = _resize_cache(common_workspace, fused_out_shape)
-            return workspace13, workspace2, fused_out
+            return workspace13, workspace2, fused_out, False
 
         common_workspace, workspace2 = current_workspace_manager().get_simultaneous(
             ((max_shape_size,), workspace_dtype),
             (workspace2_shape, workspace_dtype),
         )
         workspace13 = _resize_cache(common_workspace, workspace13_shape)
-        fused_out = _resize_cache(common_workspace, fused_out_shape)
+        if fused_out_alias is None:
+            fused_out = _resize_cache(common_workspace, fused_out_shape)
+        else:
+            assert fused_out_alias.shape == fused_out_shape
+            assert fused_out_alias.dtype == workspace_dtype
+            assert fused_out_alias.device == device
+            assert fused_out_alias.is_contiguous()
+            fused_out = fused_out_alias
 
-        return workspace13, workspace2, fused_out
+        return workspace13, workspace2, fused_out, fused_out_alias is not None
 
     def _maybe_apply_shared_experts(
         self,
@@ -1251,22 +1284,25 @@ class FusedMoEKernelModularImpl:
         if M_full == 0:
             return torch.empty_like(a1q, dtype=in_dtype)
 
-        workspace13, workspace2, fused_out = self._allocate_buffers(
-            in_dtype,
-            a1q.device,
-            M_full,
-            M_full,
-            N,
-            K,
-            top_k,
-            global_num_experts,
-            local_num_experts,
-            expert_tokens_meta,
-            activation,
+        workspace13, workspace2, fused_out, fused_out_is_specialized = (
+            self._allocate_buffers(
+                in_dtype,
+                a1q.device,
+                M_full,
+                M_full,
+                N,
+                K,
+                top_k,
+                global_num_experts,
+                local_num_experts,
+                expert_tokens_meta,
+                activation,
+            )
         )
 
         use_output_alias = (
-            output_alias is not None
+            not fused_out_is_specialized
+            and output_alias is not None
             and output_alias.shape == fused_out.shape
             and output_alias.dtype == fused_out.dtype
             and output_alias.device == fused_out.device

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

@@ -8,6 +8,7 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceContiguous,
     TopKWeightAndReduceDelegate,
+    TopKWeightAndReduceNoOP,
 )
 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
 from vllm.utils.flashinfer import nvfp4_block_scale_interleave
@@ -108,6 +109,30 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
     def output_is_reduced(self) -> bool:
         return False
 
+    def _allocate_combine_input(
+        self,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        return get_ep_group().allocate_combine_input(
+            shape,
+            dtype,
+            device,
+            is_sequence_parallel=self.is_sequence_parallel,
+        )
+
+    def allocate_fused_expert_output(
+        self,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+    ) -> torch.Tensor | None:
+        if isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP):
+            return self._allocate_combine_input(shape, dtype, device)
+        return None
+
     def prepare(
         self,
         a1: torch.Tensor,
@@ -194,16 +219,29 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
         if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
             weight_and_reduce_impl = TopKWeightAndReduceContiguous()
 
+        if isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP):
+            # allocate_fused_expert_output already selected the combine input.
+            # Reuse it directly; the symmetric allocator may return a new view
+            # object for the same high-water backing allocation.
+            combine_input = fused_expert_output
+        else:
+            combine_input = self._allocate_combine_input(
+                (topk_ids.shape[0], fused_expert_output.shape[-1]),
+                fused_expert_output.dtype,
+                fused_expert_output.device,
+            )
         out = weight_and_reduce_impl.apply(
-            output=None,
+            output=combine_input,
             fused_expert_output=fused_expert_output,
             topk_weights=topk_weights,
             topk_ids=topk_ids,
             apply_router_weight_on_input=apply_router_weight_on_input,
         )
 
-        output.copy_(
-            get_ep_group().combine(out, is_sequence_parallel=self.is_sequence_parallel)
+        get_ep_group().combine_into_output(
+            out,
+            output,
+            is_sequence_parallel=self.is_sequence_parallel,
         )