[Perf][MoE] Write FlashInfer combine into final output (#47156)

Signed-off-by: snordmann <snordmann@nvidia.com>
Co-authored-by: Codex <codex@openai.com>

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)
 # ---------------------------------------------------------------------------

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
 

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)