diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index f7dac37940..bd7ff4df60 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/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]) diff --git a/tests/distributed/test_nccl_symm_mem.py b/tests/distributed/test_nccl_symm_mem.py index bd0270fc48..d3013445cf 100644 --- a/tests/distributed/test_nccl_symm_mem.py +++ b/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) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index d3c48756ca..d8790dbf42 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/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 diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index f70bf09215..588163b846 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/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 diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 99ebcd717d..d4d588bf45 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/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: diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 88bb3d7c5a..de6680e027 100644 --- a/vllm/distributed/parallel_state.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 9bde03795c..7a4cbd4237 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/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 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 97d4e0d284..7087b9016d 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 @@ -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, )