[Perf] Fuse GLM-5.x MoE finalize into the next RMSNorm's all-reduce

tests/kernels/core/test_fused_allreduce_rms_norm.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests ``fused_allreduce_rms_norm`` on a MoE output whose top-k reduction is
+still open (``MoEOutput`` holding an ``UnfinalizedMoEOutput``).
+
+It must match the unfused path: finalize the routed output, add the shared
+expert output, then ``RMSNorm(all_reduce(.), residual)``. Covers the FlashInfer
+MoE-finalize all-reduce kernel (TP>1) and the eager fallback.
+"""
+
+import pytest
+import torch
+from torch.multiprocessing import spawn
+
+from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
+from vllm.model_executor.layers.fused_moe.moe_output import (
+    MoEOutput,
+    UnfinalizedMoEOutput,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.models.common.ops import fused_allreduce_rms_norm as fused_ar_norm
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+from vllm.utils.torch_utils import set_random_seed
+
+TOP_K = 8
+
+
+@ensure_current_vllm_config()
+def _worker_moe_finalize_ar_norm(
+    local_rank, world_size, port, num_tokens, hidden_size, fused, seed
+):
+    device = torch.device(f"cuda:{local_rank}")
+    torch.accelerator.set_device_index(device)
+    init_test_distributed_environment(
+        world_size, 1, local_rank, port, local_rank=local_rank
+    )
+    if fused:
+        import flashinfer.comm as flashinfer_comm
+
+        kernel = flashinfer_comm.trtllm_moe_finalize_allreduce_fusion
+        calls = []
+
+        def counted_kernel(**kwargs):
+            calls.append(kwargs)
+            kernel(**kwargs)
+
+        flashinfer_comm.trtllm_moe_finalize_allreduce_fusion = counted_kernel
+    else:
+        fused_ar_norm.flashinfer_trtllm_fused_allreduce_norm = None
+    dtype = torch.bfloat16
+
+    set_random_seed(seed)
+    norm = RMSNorm(hidden_size, eps=1e-5).cuda().to(dtype)
+    with torch.no_grad():
+        norm.weight.normal_(mean=1.0, std=0.1)
+    residual = 4 * torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
+
+    # Per-rank partial MoE outputs: GEMM2 rows are permuted into a padded
+    # buffer, and the permute map sends each (token, k) to a distinct row.
+    torch.manual_seed(seed + 1000 + local_rank)
+    num_rows = 2 * num_tokens * TOP_K + 16
+    gemm2 = torch.randn(num_rows, hidden_size, dtype=dtype, device=device)
+    permuted_idx = (
+        torch.randperm(num_rows, device=device)[: num_tokens * TOP_K]
+        .to(torch.int32)
+        .view(num_tokens, TOP_K)
+    )
+    weights = torch.rand(num_tokens, TOP_K, device=device).to(dtype)
+    shared = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
+
+    # Reference: the unfused path (FP32 top-k sum, shared add, AR, RMSNorm).
+    routed = (gemm2[permuted_idx.long()].float() * weights.float()[..., None]).sum(1)
+    partial = shared + routed.to(dtype)
+    ref_out, ref_res = norm(tensor_model_parallel_all_reduce(partial), residual.clone())
+
+    moe_output = MoEOutput(
+        routed=UnfinalizedMoEOutput(
+            gemm2_permuted=gemm2,
+            expert_weights=weights,
+            expanded_idx_to_permuted_idx=permuted_idx,
+        ),
+        shared_output=shared,
+    )
+    out, res = fused_ar_norm.fused_allreduce_rms_norm(
+        moe_output, residual.clone(), norm
+    )
+    torch.accelerator.synchronize()
+    if fused:
+        assert len(calls) == 1
+
+    # The fused kernel sums the top-k and the ranks in BF16.
+    torch.testing.assert_close(res, ref_res, atol=0.25, rtol=2e-2)
+    torch.testing.assert_close(out, ref_out, atol=6e-2, rtol=2e-2)
+    cleanup_dist_env_and_memory()
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("world_size", [1, 2, 4])
+@pytest.mark.parametrize("num_tokens", [1, 4, 64])
+@pytest.mark.parametrize("hidden_size", [2048, 6144])
+@pytest.mark.parametrize("fused", [True, False])
+def test_fused_allreduce_rms_norm_moe_output(
+    world_size, num_tokens, hidden_size, fused
+):
+    if world_size == 1 and fused:
+        pytest.skip("TP==1 has no fused path")
+    if current_platform.device_count() < world_size:
+        pytest.skip(f"Need >= {world_size} GPUs")
+    spawn(
+        _worker_moe_finalize_ar_norm,
+        args=(world_size, str(get_open_port()), num_tokens, hidden_size, fused, 0),
+        nprocs=world_size,
+        join=True,
+    )

vllm/models/common/ops/fused_allreduce_rms_norm.py

@@ -4,35 +4,55 @@
 
 This recovers a fusion that vLLM's torch.compile passes would normally do but
 that doesn't fire for models running eager (or under a breakable CUDA graph).
+A MoE output whose top-k reduction is still open (an ``UnfinalizedMoEOutput``)
+has that reduction and the shared-expert add folded into the same kernel.
 """
 
 import torch
 
 from vllm.distributed import (
+    get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
+    get_tp_group,
     tensor_model_parallel_all_reduce,
 )
 from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
     _AR_RESIDUAL_RMS_NORM,
+    MiB,
     _can_use_flashinfer,
+    _max_token_num,
     flashinfer_trtllm_fused_allreduce_norm,
 )
+from vllm.model_executor.layers.fused_moe.moe_output import (
+    MoEOutput,
+    UnfinalizedMoEOutput,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.platforms import current_platform
+
+# FlashInfer's MoE-finalize all-reduce is one-shot only and warns past this many
+# tokens (kOneShotMaxToken in trtllm_moe_allreduce_fusion.cuh).
+_MOE_FINALIZE_ONE_SHOT_MAX_TOKENS = 128
 
 
 def fused_allreduce_rms_norm(
-    hidden_states: torch.Tensor,
+    hidden_states: torch.Tensor | MoEOutput,
     residual: torch.Tensor,
     norm: RMSNorm,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer.
 
     ``hidden_states`` is the per-rank *partial* output of a row-parallel linear
-    run with ``reduce_results=False``; ``norm`` is the RMSNorm applied right
-    after. Returns ``(normed_output, new_residual)``, equivalent to
+    run with ``reduce_results=False``, or of a MoE layer whose reduction was
+    left open (``MoEOutput``, whose buffers the fused kernel path releases);
+    ``norm`` is the RMSNorm applied right after.
+    Returns ``(normed_output, new_residual)``, equivalent to
     ``norm(all_reduce(hidden_states), residual)``. Falls back to an explicit
     all-reduce + RMSNorm when the flashinfer fast path is unavailable.
     """
+    if isinstance(hidden_states, MoEOutput):
+        return _moe_finalize_allreduce_rms_norm(hidden_states, residual, norm)
+
     tp_size = get_tensor_model_parallel_world_size()
     if tp_size == 1:
         return norm(hidden_states, residual)
@@ -61,3 +81,136 @@ def fused_allreduce_rms_norm(
 
     reduced = tensor_model_parallel_all_reduce(hidden_states)
     return norm(reduced, residual)
+
+
+def moe_finalize_allreduce_max_num_tokens(hidden_size: int, dtype: torch.dtype) -> int:
+    """Most tokens a deferred MoE output can be finalized for by one fused
+    finalize + all-reduce + RMSNorm kernel, or 0 if that kernel is unavailable.
+
+    FlashInfer's MoE-finalize all-reduce is one-shot only and runs on the
+    TRT-LLM workspace, so this stays within both its own one-shot bound and the
+    size up to which vLLM picks the one-shot all-reduce anyway. Used by a MoE
+    layer to decide whether, and up to which batch size, to leave its top-k
+    reduction to ``fused_allreduce_rms_norm``. Creates (or reuses) that
+    workspace, so a layer only defers when the fused kernel will run.
+    """
+    tp_size = get_tensor_model_parallel_world_size()
+    if tp_size == 1 or flashinfer_trtllm_fused_allreduce_norm is None:
+        return 0
+    from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
+        _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB,
+    )
+    from vllm.distributed.device_communicators.flashinfer_all_reduce import (
+        _resolve_fi_ar_backend,
+        get_fi_ar_workspace,
+    )
+
+    backend, _ = _resolve_fi_ar_backend(get_tp_group().cpu_group)
+    capability = current_platform.get_device_capability()
+    if backend != "trtllm" or capability is None:
+        return 0
+    one_shot_mb = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get(capability.to_int(), {}).get(
+        tp_size
+    )
+    max_token_num = _max_token_num(tp_size, hidden_size, dtype)
+    if one_shot_mb is None or max_token_num is None:
+        return 0
+    workspace = get_fi_ar_workspace(
+        world_size=tp_size,
+        rank=get_tensor_model_parallel_rank(),
+        max_token_num=max_token_num,
+        hidden_dim=hidden_size,
+        dtype=dtype,
+        group=get_tp_group().cpu_group,
+    )
+    if workspace is None or workspace.backend != "trtllm":
+        return 0
+    element_size = torch.tensor([], dtype=dtype).element_size()
+    return min(
+        _MOE_FINALIZE_ONE_SHOT_MAX_TOKENS,
+        int(one_shot_mb * MiB) // (hidden_size * element_size),
+        max_token_num,
+    )
+
+
+def _moe_finalize_allreduce_rms_norm(
+    moe_output: MoEOutput,
+    residual: torch.Tensor,
+    norm: RMSNorm,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    routed = moe_output.routed
+    if not isinstance(routed, UnfinalizedMoEOutput):
+        return fused_allreduce_rms_norm(
+            _add_shared_output(routed, moe_output), residual, norm
+        )
+
+    # The fused kernel reads the routing weights in the activation dtype.
+    expert_weights = routed.expert_weights.to(residual.dtype)
+    tp_size = get_tensor_model_parallel_world_size()
+    if tp_size > 1 and flashinfer_trtllm_fused_allreduce_norm is not None:
+        from vllm.distributed.device_communicators.flashinfer_all_reduce import (
+            flashinfer_comm,
+            get_fi_ar_workspace,
+        )
+
+        # The kernel sizes its all-reduce by the token output, i.e. residual.
+        ok, max_token_num = _can_use_flashinfer(residual, tp_size)
+        workspace = (
+            get_fi_ar_workspace(
+                world_size=tp_size,
+                rank=get_tensor_model_parallel_rank(),
+                max_token_num=max_token_num,
+                hidden_dim=residual.shape[-1],
+                dtype=residual.dtype,
+                group=get_tp_group().cpu_group,
+            )
+            if ok
+            else None
+        )
+        if workspace is not None and workspace.backend == "trtllm":
+            norm_out = torch.empty_like(residual)
+            residual_out = torch.empty_like(residual)
+            flashinfer_comm.trtllm_moe_finalize_allreduce_fusion(
+                allreduce_in=routed.gemm2_permuted,
+                residual_in=residual,
+                norm_weight=norm.weight,
+                expanded_idx_to_permuted_idx=routed.expanded_idx_to_permuted_idx,
+                norm_out=norm_out,
+                residual_out=residual_out,
+                quant_out=None,
+                scale_out=None,
+                workspace_ptrs=workspace.workspace_tensor,
+                launch_with_pdl=True,
+                world_rank=workspace.rank,
+                world_size=workspace.world_size,
+                eps=norm.variance_epsilon,
+                shared_expert_output=moe_output.shared_output,
+                expert_scale_factor=expert_weights,
+                routed_scaling_factor=moe_output.routed_scaling_factor,
+            )
+            # The caller's layer loop holds moe_output until the consuming layer
+            # returns. Free the padded GEMM2 rows now, so the layer's later
+            # allocations can't be split around them (in a CUDA graph pool that
+            # costs a new segment per graph for the indexer logits).
+            del moe_output.routed, moe_output.shared_output, routed
+            return norm_out, residual_out
+
+    # Layers defer only when the fused kernel is available (see
+    # moe_finalize_allreduce_max_num_tokens); other callers finalize unfused here.
+    permuted_idx = routed.expanded_idx_to_permuted_idx
+    rows = routed.gemm2_permuted[permuted_idx.clamp(min=0).long()]
+    weights = torch.where(permuted_idx >= 0, expert_weights.float(), 0.0)
+    routed_output = (rows.float() * weights.unsqueeze(-1)).sum(dim=1)
+    return fused_allreduce_rms_norm(
+        _add_shared_output(routed_output.to(residual.dtype), moe_output),
+        residual,
+        norm,
+    )
+
+
+def _add_shared_output(routed: torch.Tensor, moe_output: MoEOutput) -> torch.Tensor:
+    if moe_output.routed_scaling_factor != 1.0:
+        routed = routed * moe_output.routed_scaling_factor
+    if moe_output.shared_output is None:
+        return routed
+    return moe_output.shared_output + routed