diff --git a/tests/kernels/core/test_fused_allreduce_rms_norm.py b/tests/kernels/core/test_fused_allreduce_rms_norm.py new file mode 100644 index 000000000000..95f33af9b9ef --- /dev/null +++ b/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, + ) diff --git a/vllm/models/common/ops/fused_allreduce_rms_norm.py b/vllm/models/common/ops/fused_allreduce_rms_norm.py index 443c8f3eec3b..a5e15c31ddca 100644 --- a/vllm/models/common/ops/fused_allreduce_rms_norm.py +++ b/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 diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index b42ba28c8df8..407dc60c8c12 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -18,6 +18,18 @@ from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( + TrtLlmBf16ExpertsMonolithic, +) +from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( + TrtLlmNvFp4ExpertsModular, + TrtLlmNvFp4ExpertsMonolithic, +) +from vllm.model_executor.layers.fused_moe.moe_output import ( + MoEOutput, + UnfinalizedMoEOutput, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner import _unpack from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, @@ -37,7 +49,10 @@ make_empty_intermediate_tensors_factory, make_layers, ) -from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm +from vllm.models.common.ops.fused_allreduce_rms_norm import ( + fused_allreduce_rms_norm, + moe_finalize_allreduce_max_num_tokens, +) from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_padding_mask, @@ -53,6 +68,14 @@ from .glm52_low_latency_gemm import enable_glm52_low_latency_gemm +# TRTLLM-Gen experts that can stop after GEMM2 in a single launch (the +# ``do_finalize=False`` path). +_DEFERRING_EXPERTS = ( + TrtLlmNvFp4ExpertsMonolithic, + TrtLlmNvFp4ExpertsModular, + TrtLlmBf16ExpertsMonolithic, +) + class DeepseekV32DecoderLayer(torch.nn.Module): def __init__( @@ -119,13 +142,47 @@ def __init__( ) self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + def defer_moe_finalize(self) -> None: + """Leave the MoE top-k reduction and shared-expert add to the consumer. + + The consumer is the ``fused_allreduce_rms_norm`` that follows this layer + (the next layer's input norm or the final norm), which then finalizes, + all-reduces and normalizes in one kernel. + """ + if self.use_sequence_parallel or not isinstance(self.mlp, DeepseekV2MoE): + return + experts = self.mlp.experts + experts_cls = getattr(experts.routed_experts.quant_method, "experts_cls", None) + # The deferred output carries no routed scaling: defer only when the + # router folds routed_scaling_factor into the routing weights, so the + # experts have no factor of their own to apply. + if ( + experts_cls not in _DEFERRING_EXPERTS + or experts.routed_scaling_factor != 1.0 + ): + return + max_num_tokens = moe_finalize_allreduce_max_num_tokens( + self.hidden_size, self.input_layernorm.weight.dtype + ) + if max_num_tokens > 0: + experts.moe_config.defer_moe_finalize(max_num_tokens) + + def _forward_moe_unfinalized(self, hidden_states: torch.Tensor) -> MoEOutput: + """``mlp`` with the routed top-k reduction and all-reduce left open.""" + # The runner's custom op returns tensors only, so run its body directly. + shared_output, routed = _unpack( + self.mlp.experts._forward_impl(hidden_states, hidden_states, hidden_states) + ) + assert isinstance(routed, UnfinalizedMoEOutput) + return MoEOutput(routed=routed, shared_output=shared_output) + def forward( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | MoEOutput, residual: torch.Tensor | None, attn_in: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor | MoEOutput, torch.Tensor]: full_num_tokens = positions.shape[0] if residual is None: @@ -161,6 +218,12 @@ def forward( if self.use_sequence_parallel and isinstance(self.mlp, DeepseekV2MoE): hidden_states = self.mlp(hidden_states, already_sequence_parallel=True) + elif isinstance( + self.mlp, DeepseekV2MoE + ) and self.mlp.experts.moe_config.should_defer_moe_finalize( + hidden_states.shape[0] + ): + hidden_states = self._forward_moe_unfinalized(hidden_states) else: hidden_states = self.mlp(hidden_states) return hidden_states, residual @@ -235,6 +298,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.aux_hidden_state_layers = tuple[int, ...]() self.num_redundant_experts = parallel_config.eplb_config.num_redundant_experts + # A layer's output feeds the next layer's input norm or the final norm, + # both of which can finalize a MoE output, except on a non-last PP rank, + # whose last layer ships its output as IntermediateTensors. + local_layers = self.layers[self.start_layer : self.end_layer] + if not get_pp_group().is_last_rank: + local_layers = local_layers[:-1] + for layer in local_layers: + layer.defer_moe_finalize() + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -437,6 +509,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): if self.config.model_type == "glm_moe_dsa": enable_glm52_low_latency_gemm(self, vllm_config.model_config.dtype) + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + super().set_aux_hidden_state_layers(layers) + # An aux hidden state reads the previous layer's output as a tensor. + for idx in layers: + if self.model.start_layer < idx < self.model.end_layer: + mlp = self.model.layers[idx - 1].mlp + if isinstance(mlp, DeepseekV2MoE): + mlp.experts.moe_config.limit_deferred_moe_finalize(0) + def set_moe_parameters(self): # Same as the base, but keyed on the MoE block type rather than the # decoder-layer type (DeepseekV32DecoderLayer is a plain nn.Module). diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index d4dee7b9d462..c6e4fd587862 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -98,6 +98,8 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: topk_indices_buffer=topk_indices_buffer, index_group_builder=index_group_builder, ) + # The block's output feeds the final norm below. + self.mtp_block.defer_moe_finalize() def forward( self,