[Perf][MoE] Write FlashInfer combine into final output (#47156)
Signed-off-by: snordmann <snordmann@nvidia.com> Co-authored-by: Codex <codex@openai.com>
Samuel Nordmann committed
Jul 16, 2026 at 16:21 UTC
d08eebad162bbd1f2e99cca550313daaa81c7654
3 files changed
+61
-2
tests/distributed/test_mnnvl_alltoall.py
+30
@@ -206,6 +206,36 @@ requires_deep_ep_v2 = pytest.mark.skipif(
206
# should run even when FlashInfer NVLink backends are not installed.
207
208
209
+@pytest.mark.parametrize("supports_output", [False, True])
210
+def test_one_sided_combine_into_compatibility(supports_output):
211
+ from vllm.distributed.device_communicators.all2all import (
212
+ FlashInferNVLinkOneSidedManager,
213
+ )
214
+
215
+ class FakeMoeAlltoAll:
216
+ def combine(
217
+ self,
218
+ payload,
219
+ runtime_max_tokens_per_rank,
220
+ output=None,
221
+ ):
222
+ result = payload + runtime_max_tokens_per_rank
223
+ if output is None:
224
+ return result
225
+ output.copy_(result)
226
+ return output
227
+
228
+ manager = FlashInferNVLinkOneSidedManager.__new__(FlashInferNVLinkOneSidedManager)
229
+ manager.moe_alltoall = FakeMoeAlltoAll()
230
+ manager._combine_supports_output = supports_output
231
+ payload = torch.arange(4, dtype=torch.float32)
232
+ output = torch.empty_like(payload)
233
+
234
+ manager.combine_into(payload, runtime_max_tokens_per_rank=2, output=output)
235
+
236
+ torch.testing.assert_close(output, payload + 2)
237
+
238
+
239
# ---------------------------------------------------------------------------
240
# Test 1: Two-sided manager lifecycle (init, cleanup, reinit, ensure_init)
241
# ---------------------------------------------------------------------------
vllm/distributed/device_communicators/all2all.py
+29
@@ -16,6 +16,7 @@ from vllm.utils.flashinfer import (
16
has_flashinfer_nvlink_one_sided,
17
has_flashinfer_nvlink_two_sided,
18
)
19
+from vllm.utils.func_utils import supports_kw
20
from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2, has_mori
21
22
from .base_device_communicator import All2AllManagerBase, Cache
@@ -690,6 +691,7 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
691
self.max_num_tokens = 0
692
self.top_k = 0
693
self.num_experts = 0
694
+ self._combine_supports_output = False
695
696
def initialize(
697
self,
@@ -790,6 +792,12 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
792
workspace_size_per_rank=self.workspace_size,
793
mnnvl_config=ep_config,
794
)
795
+ try:
796
+ self._combine_supports_output = supports_kw(
797
+ self.moe_alltoall.combine, "output", allow_var_kwargs=False
798
+ )
799
+ except (TypeError, ValueError):
800
+ self._combine_supports_output = False
801
802
self.gpus_per_node = gpus_per_node
803
self.initialized = True
@@ -804,6 +812,27 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
812
# different shape sequences, so a world-level barrier would deadlock.
813
dist.barrier(group=self.cpu_group)
814
815
+ def combine_into(
816
+ self,
817
+ payload: torch.Tensor,
818
+ runtime_max_tokens_per_rank: int,
819
+ output: torch.Tensor,
820
+ ) -> None:
821
+ """Combine into ``output``, with a fallback for older FlashInfer."""
822
+ assert self.moe_alltoall is not None
823
+ if self._combine_supports_output:
824
+ self.moe_alltoall.combine(
825
+ payload=payload,
826
+ runtime_max_tokens_per_rank=runtime_max_tokens_per_rank,
827
+ output=output,
828
+ )
829
+ else:
830
+ combined_output = self.moe_alltoall.combine(
831
+ payload=payload,
832
+ runtime_max_tokens_per_rank=runtime_max_tokens_per_rank,
833
+ )
834
+ output.copy_(combined_output)
835
+
836
def get_handle(self, kwargs):
837
return self
838
vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py
+2
-2
@@ -161,8 +161,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
161
ep_size, self.runtime_max_tokens_per_rank, hidden_size
162
)
163
164
- combined_output = self.all2all_manager.moe_alltoall.combine( # type: ignore[attr-defined]
164
+ self.all2all_manager.combine_into( # type: ignore[attr-defined]
165
payload=fused_expert_output,
166
runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
167
+ output=output,
168
)
168
- output.copy_(combined_output)