diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index 95c905fc080..fb2d9bb832e 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -206,6 +206,36 @@ requires_deep_ep_v2 = pytest.mark.skipif( # should run even when FlashInfer NVLink backends are not installed. +@pytest.mark.parametrize("supports_output", [False, True]) +def test_one_sided_combine_into_compatibility(supports_output): + from vllm.distributed.device_communicators.all2all import ( + FlashInferNVLinkOneSidedManager, + ) + + class FakeMoeAlltoAll: + def combine( + self, + payload, + runtime_max_tokens_per_rank, + output=None, + ): + result = payload + runtime_max_tokens_per_rank + if output is None: + return result + output.copy_(result) + return output + + manager = FlashInferNVLinkOneSidedManager.__new__(FlashInferNVLinkOneSidedManager) + manager.moe_alltoall = FakeMoeAlltoAll() + manager._combine_supports_output = supports_output + payload = torch.arange(4, dtype=torch.float32) + output = torch.empty_like(payload) + + manager.combine_into(payload, runtime_max_tokens_per_rank=2, output=output) + + torch.testing.assert_close(output, payload + 2) + + # --------------------------------------------------------------------------- # Test 1: Two-sided manager lifecycle (init, cleanup, reinit, ensure_init) # --------------------------------------------------------------------------- diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 81006da401d..7f540bc4b1f 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -16,6 +16,7 @@ from vllm.utils.flashinfer import ( has_flashinfer_nvlink_one_sided, has_flashinfer_nvlink_two_sided, ) +from vllm.utils.func_utils import supports_kw from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2, has_mori from .base_device_communicator import All2AllManagerBase, Cache @@ -690,6 +691,7 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): self.max_num_tokens = 0 self.top_k = 0 self.num_experts = 0 + self._combine_supports_output = False def initialize( self, @@ -790,6 +792,12 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): workspace_size_per_rank=self.workspace_size, mnnvl_config=ep_config, ) + try: + self._combine_supports_output = supports_kw( + self.moe_alltoall.combine, "output", allow_var_kwargs=False + ) + except (TypeError, ValueError): + self._combine_supports_output = False self.gpus_per_node = gpus_per_node self.initialized = True @@ -804,6 +812,27 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): # different shape sequences, so a world-level barrier would deadlock. dist.barrier(group=self.cpu_group) + def combine_into( + self, + payload: torch.Tensor, + runtime_max_tokens_per_rank: int, + output: torch.Tensor, + ) -> None: + """Combine into ``output``, with a fallback for older FlashInfer.""" + assert self.moe_alltoall is not None + if self._combine_supports_output: + self.moe_alltoall.combine( + payload=payload, + runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, + output=output, + ) + else: + combined_output = self.moe_alltoall.combine( + payload=payload, + runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, + ) + output.copy_(combined_output) + def get_handle(self, kwargs): return self diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py index e49d8b2624a..74341e7681f 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py @@ -161,8 +161,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo ep_size, self.runtime_max_tokens_per_rank, hidden_size ) - combined_output = self.all2all_manager.moe_alltoall.combine( # type: ignore[attr-defined] + self.all2all_manager.combine_into( # type: ignore[attr-defined] payload=fused_expert_output, runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank, + output=output, ) - output.copy_(combined_output)