[Kimi K2.5] Fuse deferred MoE finalize into next-layer AR + residual + RMSNorm

python/sglang/srt/layers/communicator.py

@@ -567,7 +567,28 @@ def prepare_attn(
                 and hasattr(hidden_states, "_sglang_needs_allreduce_fusion")
                 and hidden_states._sglang_needs_allreduce_fusion
             ):
-                if (
+                moe_finalize_bundle = getattr(
+                    hidden_states, "_sglang_moe_finalize_bundle", None
+                )
+                if moe_finalize_bundle is not None and hasattr(
+                    self.input_layernorm, "forward_with_moe_finalize_allreduce_fusion"
+                ):
+                    # Fuse the deferred MoE finalize into the AR + residual +
+                    # RMSNorm. residual_out (kept for Eagle3 capture) is the
+                    # post-all-reduce, pre-norm residual. ``hidden_states`` here
+                    # is the un-finalized gemm2_out carried from the prev layer.
+                    hidden_states, residual = (
+                        self.input_layernorm.forward_with_moe_finalize_allreduce_fusion(
+                            gemm2_out=hidden_states,
+                            residual=residual,
+                            expanded_idx_to_permuted_idx=moe_finalize_bundle.expanded_idx_to_permuted_idx,
+                            expert_scale_factor=moe_finalize_bundle.expert_scale_factor,
+                            shared_expert_output=moe_finalize_bundle.shared_output,
+                            top_k=moe_finalize_bundle.top_k,
+                            use_attn_tp_group=False,
+                        )
+                    )
+                elif (
                     apply_aiter_all_reduce_fusion(hidden_states)
                     or apply_flashinfer_allreduce_fusion(hidden_states.shape[0])
                 ) and hasattr(self.input_layernorm, "forward_with_allreduce_fusion"):

python/sglang/srt/layers/flashinfer_comm_fusion.py

@@ -19,6 +19,7 @@
     get_tp_group,
 )
 from sglang.srt.distributed.parallel_state import in_the_same_node_as
+from sglang.srt.environ import envs
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     ceil_align,
@@ -41,6 +42,9 @@
 _flashinfer_create_workspace_supports_comm_backend = False
 _flashinfer_allreduce_supports_trigger_completion = False
 _mnnvl_non_blackwell_fallback_logged = False
+# Mirrors flashinfer.comm.trtllm_ar.MAX_COMM_SIZE (MAX_INT32 rounded down to 2MB):
+# the largest one-shot lamport comm the MoE-finalize fusion kernel accepts.
+_FLASHINFER_MAX_COMM_SIZE = 2147483647 & ~((1 << 21) - 1)
 
 
 def _mnnvl_supported(is_multi_node: bool) -> bool:
@@ -62,17 +66,28 @@ def _resolve_backend(backend: str, is_multi_node: bool = False) -> str:
     if backend == "auto":
         # Prefer mnnvl wherever it is supported (any Blackwell system, or SM90
         # single-node); fall back to trtllm otherwise.
-        return "mnnvl" if _mnnvl_supported(is_multi_node) else "trtllm"
-
-    if backend == "mnnvl" and not _mnnvl_supported(is_multi_node):
+        resolved = "mnnvl" if _mnnvl_supported(is_multi_node) else "trtllm"
+    elif backend == "mnnvl" and not _mnnvl_supported(is_multi_node):
         if not _mnnvl_non_blackwell_fallback_logged:
             logger.info(
                 "FlashInfer allreduce fusion: forcing trtllm backend "
                 "(mnnvl requires a Blackwell system, or SM90 single-node)."
             )
             _mnnvl_non_blackwell_fallback_logged = True
-        return "trtllm"
-    return backend
+        resolved = "trtllm"
+    else:
+        resolved = backend
+
+    # The deferred MoE finalize fusion (kMoEFinalizeARResidualRMSNorm, pattern 7)
+    # is only implemented by the trtllm backend; mnnvl rejects it. When deferred
+    # finalize is enabled the finalize is folded into the AR + residual + RMSNorm
+    # fusion, so the shared workspace must be trtllm (which also serves the plain
+    # kARResidualRMSNorm pattern used by the non-finalize fusion path). Note this
+    # ties the whole fusion to trtllm's intra-node lamport all-reduce, so deferred
+    # finalize fusion is single-(NVLink-)node only -- see _mnnvl_supported.
+    if resolved == "mnnvl" and envs.SGLANG_ENABLE_MOE_DEFERRED_FINALIZE.get():
+        resolved = "trtllm"
+    return resolved
 
 
 def resolve_flashinfer_allreduce_fusion_backend(server_args) -> Optional[str]:
@@ -817,26 +832,214 @@ def flashinfer_allreduce_residual_rmsnorm(
     return norm_out, residual_out
 
 
+def fake_flashinfer_moe_finalize_allreduce_residual_rmsnorm(
+    gemm2_out: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    expanded_idx_to_permuted_idx: torch.Tensor,
+    expert_scale_factor: torch.Tensor,
+    shared_expert_output: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+    max_token_num: int = 16384,
+    use_oneshot: Optional[bool] = None,
+    use_attn_tp_group: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    residual_out = torch.empty_like(residual)
+    norm_out = torch.empty_like(residual)
+    return norm_out, residual_out
+
+
+@register_custom_op(
+    mutates_args=["gemm2_out", "residual", "weight"],
+    fake_impl=fake_flashinfer_moe_finalize_allreduce_residual_rmsnorm,
+)
+def flashinfer_moe_finalize_allreduce_residual_rmsnorm(
+    gemm2_out: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    expanded_idx_to_permuted_idx: torch.Tensor,
+    expert_scale_factor: torch.Tensor,
+    shared_expert_output: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+    max_token_num: int = 2048,
+    use_oneshot: Optional[bool] = None,
+    use_attn_tp_group: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Fused MoE finalize + all-reduce + residual add + RMSNorm.
+
+    Uses FlashInfer's ``kMoEFinalizeARResidualRMSNorm`` pattern, which combines
+    the per-rank expert outputs (weighted by ``expert_scale_factor`` and scattered
+    via ``expanded_idx_to_permuted_idx``), adds the shared-expert output, runs the
+    all-reduce across the MoE TP group, adds the residual and applies RMSNorm --
+    all in a single kernel. This lets the deferred MoE finalize be folded into the
+    next layer's input-RMSNorm AR fusion, while the post-all-reduce ``residual_out``
+    stays available for Eagle3 aux capture.
+
+    Args:
+        gemm2_out: Un-finalized expert outputs (the kernel's ``allreduce_in``),
+            shape [num_permuted_tokens, hidden_dim].
+        residual: Residual stream, shape [num_tokens, hidden_dim].
+        weight: RMSNorm weight (rms_gamma).
+        expanded_idx_to_permuted_idx: Map from (token, top_k slot) to the row in
+            ``gemm2_out`` (int32).
+        expert_scale_factor: Per-(token, expert) combine weights.
+        shared_expert_output: Optional shared-expert output [num_tokens, hidden_dim].
+        eps: RMSNorm epsilon.
+        max_token_num: Maximum token number for the workspace.
+        use_oneshot: Whether to use oneshot mode.
+        use_attn_tp_group: If True, use the attention TP group; otherwise the MoE
+            TP group (the all-reduce group the MoE output lives on).
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: (norm_output, residual_output), or
+        (None, None) if the fused kernel is unavailable and the caller should fall
+        back to a separate finalize + AR + RMSNorm.
+    """
+    if not is_flashinfer_available() or _flashinfer_comm is None:
+        return None, None
+
+    moe_finalize_pattern = getattr(
+        _flashinfer_comm.AllReduceFusionPattern,
+        "kMoEFinalizeARResidualRMSNorm",
+        None,
+    )
+    if moe_finalize_pattern is None:
+        logger.debug(
+            "FlashInfer build lacks kMoEFinalizeARResidualRMSNorm, "
+            "falling back to separate finalize + allreduce fusion"
+        )
+        return None, None
+
+    if use_attn_tp_group:
+        world_size = get_attn_tensor_model_parallel_world_size()
+    else:
+        if get_moe_expert_parallel_world_size() > 1:
+            world_size = get_moe_expert_parallel_world_size()
+        else:
+            world_size = get_moe_tensor_parallel_world_size()
+
+    if world_size <= 1:
+        return None, None
+
+    # The MoE-finalize pattern all-reduces the *permuted* expert output
+    # (``gemm2_out``), whose row count is dominated by per-expert tile padding and
+    # is many times larger than the token count. FlashInfer sizes the lamport comm
+    # buffer from ``allreduce_in.numel()``, so the workspace must cover
+    # ``gemm2_out.shape[0]`` rows, not ``residual.shape[0]``.
+    comm_token_num = gemm2_out.shape[0]
+    # FlashInfer's MoE-finalize kernel derives ``top_k`` from
+    # ``expanded_idx_to_permuted_idx.size(-1)`` and indexes both the index map and
+    # ``expert_scale_factor`` as ``[token, slot]``. The deferred-finalize path hands
+    # us a *flat* ``[token_num * top_k]`` index tensor, which would make the kernel
+    # treat ``top_k == token_num * top_k`` and read far past the end of every per-
+    # token row (illegal access). Restore the ``[token_num, top_k]`` shape the kernel
+    # expects (``expert_scale_factor`` is already 2D, but normalize defensively).
+    num_token = residual.shape[0]
+    top_k = expanded_idx_to_permuted_idx.numel() // max(num_token, 1)
+    expanded_idx_to_permuted_idx = expanded_idx_to_permuted_idx.reshape(
+        num_token, top_k
+    )
+    expert_scale_factor = expert_scale_factor.reshape(num_token, top_k)
+    if (
+        not gemm2_out.is_contiguous()
+        or not residual.is_contiguous()
+        or not weight.is_contiguous()
+        or not expanded_idx_to_permuted_idx.is_contiguous()
+        or not expert_scale_factor.is_contiguous()
+        or (
+            shared_expert_output is not None
+            and not shared_expert_output.is_contiguous()
+        )
+    ):
+        logger.debug("Non-contiguous tensors, skipping FlashInfer MoE finalize fusion")
+        return None, None
+
+    # FlashInfer's MoE-finalize fusion is one-shot only and rejects comms whose
+    # lamport size exceeds MAX_COMM_SIZE. Mirror that bound here and fall back for
+    # large (e.g. prefill) shapes instead of letting the kernel raise.
+    required_comm_size = comm_token_num * residual.shape[-1] * 2 * world_size
+    if required_comm_size > _FLASHINFER_MAX_COMM_SIZE:
+        logger.debug(
+            "MoE finalize fusion comm size %d exceeds MAX_COMM_SIZE; falling back",
+            required_comm_size,
+        )
+        return None, None
+
+    # Never grow the workspace here: this op runs inside captured CUDA graphs, and
+    # reallocating the symmetric workspace would invalidate graph-captured pointers.
+    # The workspace is pre-sized for the deferred-finalize comm in
+    # pre_initialize_workspaces(); if it is too small for this shape (e.g. an
+    # un-captured large prefill), fall back to the separate finalize + AR fusion.
+    workspace_manager = _get_workspace_manager(use_attn_tp_group)
+    if (
+        not workspace_manager.initialized
+        or workspace_manager.workspace is None
+        or not workspace_manager.is_buffer_size_sufficient(
+            token_num=comm_token_num,
+            hidden_dim=residual.shape[-1],
+            dtype=residual.dtype,
+            use_oneshot=use_oneshot,
+        )
+    ):
+        logger.debug(
+            "FlashInfer workspace insufficient for MoE finalize fusion "
+            "(need %d permuted rows); falling back",
+            comm_token_num,
+        )
+        return None, None
+
+    residual_out = torch.empty_like(residual)
+    norm_out = torch.empty_like(residual)
+
+    _flashinfer_comm.allreduce_fusion(
+        input=gemm2_out,
+        workspace=workspace_manager.workspace,
+        pattern=moe_finalize_pattern,
+        launch_with_pdl=True,
+        residual_out=residual_out,
+        norm_out=norm_out,
+        residual_in=residual,
+        rms_gamma=weight,
+        rms_eps=eps,
+        use_oneshot=use_oneshot,
+        expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx,
+        expert_scale_factor=expert_scale_factor,
+        shared_expert_output=shared_expert_output,
+    )
+
+    return norm_out, residual_out
+
+
 def pre_initialize_workspaces(
     max_token_num: int,
     hidden_dim: int,
     dtype: torch.dtype,
     use_oneshot: Optional[bool] = None,
+    moe_finalize_max_token_num: Optional[int] = None,
 ):
     """Pre-initialize flashinfer workspaces before CUDA graph capture.
 
     This must be called before graph capture to avoid collective operations
     (broadcasts, barriers) inside the graph capture context, which can
     deadlock with custom_all_reduce.register_graph_buffers.
+
+    ``moe_finalize_max_token_num`` sizes the MoE-TP workspace for the deferred
+    MoE-finalize AR fusion, whose ``allreduce_in`` is the *permuted* expert
+    output (``num_local_experts`` tile padding + ``num_tokens * top_k`` rows) and
+    is far larger than the token count. It must be sized up front because the
+    fused op runs inside captured CUDA graphs and must never grow the workspace.
     """
     if _flashinfer_allreduce_unavailable or _flashinfer_comm is None:
         return
 
-    # Initialize MoE workspace
+    # Initialize MoE workspace (sized for the permuted MoE-finalize comm when the
+    # deferred-finalize fusion is enabled).
+    moe_max_token_num = max(max_token_num, moe_finalize_max_token_num or 0)
     ensure_workspace_initialized(
-        max_token_num=max_token_num,
+        max_token_num=moe_max_token_num,
         hidden_dim=hidden_dim,
         dtype=dtype,
+        token_num=moe_max_token_num,
         use_oneshot=use_oneshot,
         use_attn_tp_group=False,
     )

python/sglang/srt/layers/layernorm.py

@@ -202,6 +202,60 @@ def _forward_with_allreduce_fusion(
     return norm_module.forward(x, residual, post_residual_addition)
 
 
+def _forward_with_moe_finalize_allreduce_fusion(
+    norm_module,
+    gemm2_out: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    expanded_idx_to_permuted_idx: torch.Tensor,
+    expert_scale_factor: torch.Tensor,
+    shared_expert_output: Optional[torch.Tensor],
+    top_k: int,
+    use_attn_tp_group: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Fuse MoE finalize + all-reduce + residual + RMSNorm.
+
+    Returns (norm_out, residual_out). ``residual_out`` is the post-all-reduce,
+    pre-norm residual that Eagle3 aux capture reads. Falls back to a separate
+    finalize kernel followed by the standard AR+RMSNorm fusion when the fused
+    MoE-finalize kernel is unavailable.
+    """
+    from sglang.srt.layers.flashinfer_comm_fusion import (
+        flashinfer_moe_finalize_allreduce_residual_rmsnorm,
+    )
+
+    norm_out, residual_out = flashinfer_moe_finalize_allreduce_residual_rmsnorm(
+        gemm2_out=gemm2_out,
+        residual=residual,
+        weight=weight,
+        expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx,
+        expert_scale_factor=expert_scale_factor,
+        shared_expert_output=shared_expert_output,
+        eps=norm_module.variance_epsilon,
+        max_token_num=max(residual.shape[0], 2048),
+        use_attn_tp_group=use_attn_tp_group,
+    )
+    if norm_out is not None:
+        return norm_out, residual_out
+
+    # Fallback: run the finalize kernel separately, then the standard
+    # all-reduce + residual + RMSNorm fusion (still emits a fresh residual_out).
+    from sglang.jit_kernel.moe_finalize_fuse_shared import moe_finalize_fuse_shared
+    from sglang.jit_kernel.utils import is_arch_support_pdl
+
+    finalized = moe_finalize_fuse_shared(
+        gemm2_out,
+        expanded_idx_to_permuted_idx,
+        expert_scale_factor,
+        shared_expert_output,
+        top_k,
+        enable_pdl=is_arch_support_pdl(),
+    )
+    return _forward_with_allreduce_fusion(
+        norm_module, finalized, residual, None, weight, use_attn_tp_group
+    )
+
+
 class RMSNorm(MultiPlatformOp):
     def __init__(
         self,
@@ -562,6 +616,29 @@ def forward_with_allreduce_fusion(
             self, x, residual, post_residual_addition, self.weight, use_attn_tp_group
         )
 
+    def forward_with_moe_finalize_allreduce_fusion(
+        self,
+        gemm2_out: torch.Tensor,
+        residual: torch.Tensor,
+        expanded_idx_to_permuted_idx: torch.Tensor,
+        expert_scale_factor: torch.Tensor,
+        shared_expert_output: Optional[torch.Tensor],
+        top_k: int,
+        use_attn_tp_group: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Fuse the deferred MoE finalize into this layer's AR + residual + RMSNorm."""
+        return _forward_with_moe_finalize_allreduce_fusion(
+            self,
+            gemm2_out,
+            residual,
+            self.weight,
+            expanded_idx_to_permuted_idx,
+            expert_scale_factor,
+            shared_expert_output,
+            top_k,
+            use_attn_tp_group,
+        )
+
 
 class LayerNorm(MultiPlatformOp):
     def __init__(