Add MoE output contract for MoE tail fusion (#51407)

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

vllm/model_executor/layers/fused_moe/config.py

@@ -1311,6 +1311,13 @@ class FusedMoEConfig:
     # Only honored on the non-reduced (late-AR) TP path. Default False.
     skip_final_all_reduce: bool = False
 
+    # When True, experts that can stop after GEMM2 are allowed to hand back an
+    # UnfinalizedMoEOutput instead of finalized states, leaving the top-k
+    # reduction to fuse into the consumer. Set by layers that have such a
+    # consumer; read through `use_deferred_moe_finalize`, which applies the
+    # guards. Kernels without the capability ignore it. Default False.
+    defer_moe_finalize: bool = False
+
     # SwiGLU clamp limit. When set, backends that do not implement the clamp
     # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle
     # cannot silently select one and drop the clamp.
@@ -1428,6 +1435,18 @@ class FusedMoEConfig:
     def use_ep(self):
         return self.moe_parallel_config.use_ep
 
+    @property
+    def use_deferred_moe_finalize(self) -> bool:
+        """Whether experts may return an unfinalized output on this deployment.
+
+        Evaluated on read rather than in ``__post_init__`` because
+        ``defer_moe_finalize`` is set after construction, like
+        ``skip_final_all_reduce``.
+        """
+        # Without TP there is no all-reduce for the top-k reduction to fuse
+        # into, so the deferred form buys nothing.
+        return self.defer_moe_finalize and self.tp_size > 1
+
     @property
     def use_deepep_ht_kernels(self):
         return self.moe_parallel_config.use_deepep_ht_kernels

vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py

@@ -11,6 +11,9 @@ from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
     RoutingMethodType,
 )
+from vllm.model_executor.layers.fused_moe.moe_output import (
+    UnfinalizedMoEOutput,
+)
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
@@ -266,13 +269,18 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol
         e_score_correction_bias: torch.Tensor | None = None,
         routed_scaling_factor: float | None = None,
         topk_group: int | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | UnfinalizedMoEOutput:
         import flashinfer
 
         assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
 
+        num_tokens = hidden_states.shape[0]
+        # The runner divides by the token count on the host, so an idle rank's
+        # dummy 0-token forward has to keep the finalized (empty) form.
+        defer = self.moe_config.use_deferred_moe_finalize and num_tokens > 0
+
         routing_replay_out = self._maybe_make_routing_replay_buffer(
-            num_tokens=hidden_states.shape[0],
+            num_tokens=num_tokens,
             device=hidden_states.device,
         )
         out = flashinfer.fused_moe.trtllm_bf16_moe(
@@ -293,8 +301,19 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol
             activation_type=activation_to_flashinfer_int(activation),
             tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
             routing_replay_out=routing_replay_out,
+            do_finalize=not defer,
         )
-        self._maybe_dispatch_routing_replay(
-            routing_replay_out, num_tokens=hidden_states.shape[0]
-        )
-        return out
+        self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens)
+        if defer:
+            # flashinfer returns a flat permute map; the protocol wants
+            # [num_tokens, top_k] so consumers can read top_k from its shape.
+            return UnfinalizedMoEOutput(
+                gemm2_permuted=out[0],
+                expert_weights=out[1],
+                expanded_idx_to_permuted_idx=out[2]
+                .to(torch.int32)
+                .view(num_tokens, self.topk),
+            )
+        # do_finalize=True yields the finalized states (a bare tensor on some
+        # FlashInfer versions, a single-element list on others).
+        return out[0] if isinstance(out, (list, tuple)) else out

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

@@ -13,6 +13,9 @@ from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
     RoutingMethodType,
 )
+from vllm.model_executor.layers.fused_moe.moe_output import (
+    UnfinalizedMoEOutput,
+)
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
@@ -518,7 +521,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
         e_score_correction_bias: torch.Tensor | None = None,
         routed_scaling_factor: float | None = None,
         topk_group: int | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | UnfinalizedMoEOutput:
         import flashinfer
 
         assert self._supports_activation(activation)
@@ -543,8 +546,13 @@ class TrtLlmNvFp4ExpertsMonolithic(
 
         output1_scale_gate_scalar = self.quant_config.g1_alphas
 
+        num_tokens = hidden_states.shape[0]
+        # The runner divides by the token count on the host, so an idle rank's
+        # dummy 0-token forward has to keep the finalized (empty) form.
+        defer = self.moe_config.use_deferred_moe_finalize and num_tokens > 0
+
         routing_replay_out = self._maybe_make_routing_replay_buffer(
-            num_tokens=hidden_states.shape[0],
+            num_tokens=num_tokens,
             device=hidden_states.device,
         )
         # Invoke kernel.
@@ -578,13 +586,21 @@ class TrtLlmNvFp4ExpertsMonolithic(
             local_num_experts=self.local_num_experts,
             routed_scaling_factor=routed_scaling_factor,
             routing_method_type=self.routing_method_type,
-            do_finalize=True,
+            do_finalize=not defer,
             activation_type=activation_to_flashinfer_int(activation),
             per_token_scale=per_token_scale,
             tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
             routing_replay_out=routing_replay_out,
-        )[0]
-        self._maybe_dispatch_routing_replay(
-            routing_replay_out, num_tokens=hidden_states.shape[0]
         )
-        return result
+        self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens)
+        if defer:
+            # flashinfer returns a flat permute map; the protocol wants
+            # [num_tokens, top_k] so consumers can read top_k from its shape.
+            return UnfinalizedMoEOutput(
+                gemm2_permuted=result[0],
+                expert_weights=result[1],
+                expanded_idx_to_permuted_idx=result[2]
+                .to(torch.int32)
+                .view(num_tokens, self.topk),
+            )
+        return result[0]

vllm/model_executor/layers/fused_moe/moe_output.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Output contract between a MoE layer and a consumer that fuses its tail."""
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class UnfinalizedMoEOutput:
+    """Unfinalized output of a monolithic MoE kernel.
+
+    Kernels that can stop after GEMM2 (the TRTLLM-Gen ``do_finalize=False`` path)
+    hand back their permuted, unweighted output plus the routing weights and the
+    permute map, so that the top-k reduction can be fused with whatever follows
+    -- the shared-expert add and the tensor-parallel all-reduce -- instead of
+    running as its own kernel.
+
+    The buffers are consumed as-is by the fused kernels, which index
+    ``gemm2_permuted`` by row: it must be densely packed at ``hidden_dim``, and
+    its row count (an autotuner-dependent padded value) is never referenced.
+    """
+
+    # [num_permuted_rows, hidden_dim] permuted, unweighted GEMM2 output.
+    gemm2_permuted: torch.Tensor
+    # [num_tokens, top_k] routing weights, in the activation dtype -- consumers
+    # index them as such, and a buffer typed wider than its contents reads as
+    # garbage. These already carry routed_scaling_factor for routing methods that
+    # fold it in.
+    expert_weights: torch.Tensor
+    # [num_tokens, top_k] int32 permute map; -1 = expert not local to this rank.
+    expanded_idx_to_permuted_idx: torch.Tensor
+
+
+@dataclass
+class MoEOutput:
+    """A MoE layer's output with its final reduction still open.
+
+    Returned by layers whose MoE runs un-reduced (``reduce_results=False``) so
+    that the consumer -- typically the next layer's RMSNorm -- can fuse the
+    tensor-parallel all-reduce into itself instead of paying for a standalone
+    one. Keeping the shared-expert output and the routed scale separate leaves
+    that reduction the consumer's to schedule; when the routed output is still
+    unfinalized, the top-k reduction is open too and can fold into the same
+    kernel.
+
+    Producers only leave the routed output unfinalized when a fused consumer can
+    actually take that form -- the token ceiling and topology support are theirs
+    to check -- so an ``UnfinalizedMoEOutput`` here means the fused path applies,
+    and a consumer need not re-derive that.
+    """
+
+    # Un-reduced routed output, either finalized or not.
+    routed: torch.Tensor | UnfinalizedMoEOutput
+    # Un-reduced shared-expert output, folded in by whoever finalizes.
+    shared_output: torch.Tensor | None = None
+    # Applied to the routed sum before the shared add. Already folded into the
+    # routing weights for methods that do so, in which case this is 1.0.
+    routed_scaling_factor: float = 1.0