From dd7cc85f1720ebcb65b1822cbc5054af66ae292d Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Tue, 11 Aug 2026 22:23:59 +0800 Subject: [PATCH] Add MoE output contract for MoE tail fusion (#51407) Signed-off-by: Jee Jee Li --- .../model_executor/layers/fused_moe/config.py | 19 ++++++ .../fused_moe/experts/trtllm_bf16_moe.py | 31 ++++++++-- .../fused_moe/experts/trtllm_nvfp4_moe.py | 30 +++++++--- .../layers/fused_moe/moe_output.py | 60 +++++++++++++++++++ 4 files changed, 127 insertions(+), 13 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/moe_output.py diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 054240271f..456d6c5793 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 559180bd72..ae6dc11beb 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index a5555dc132..29ccd5770e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/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] diff --git a/vllm/model_executor/layers/fused_moe/moe_output.py b/vllm/model_executor/layers/fused_moe/moe_output.py new file mode 100644 index 0000000000..d3fc1a75fb --- /dev/null +++ b/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