| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Fused all-reduce + residual-add + RMSNorm for eager model paths. |
| 4 | |
| 5 | This recovers a fusion that vLLM's torch.compile passes would normally do but |
| 6 | that doesn't fire for models running eager (or under a breakable CUDA graph). |
| 7 | """ |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm.distributed import ( |
| 12 | get_tensor_model_parallel_world_size, |
| 13 | tensor_model_parallel_all_reduce, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import ( |
| 16 | _AR_RESIDUAL_RMS_NORM, |
| 17 | _can_use_flashinfer, |
| 18 | flashinfer_trtllm_fused_allreduce_norm, |
| 19 | ) |
| 20 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 21 | |
| 22 | |
| 23 | def fused_allreduce_rms_norm( |
| 24 | hidden_states: torch.Tensor, |
| 25 | residual: torch.Tensor, |
| 26 | norm: RMSNorm, |
| 27 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 28 | """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer. |
| 29 | |
| 30 | ``hidden_states`` is the per-rank *partial* output of a row-parallel linear |
| 31 | run with ``reduce_results=False``; ``norm`` is the RMSNorm applied right |
| 32 | after. Returns ``(normed_output, new_residual)``, equivalent to |
| 33 | ``norm(all_reduce(hidden_states), residual)``. Falls back to an explicit |
| 34 | all-reduce + RMSNorm when the flashinfer fast path is unavailable. |
| 35 | """ |
| 36 | tp_size = get_tensor_model_parallel_world_size() |
| 37 | if tp_size == 1: |
| 38 | return norm(hidden_states, residual) |
| 39 | |
| 40 | if flashinfer_trtllm_fused_allreduce_norm is not None: |
| 41 | ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size) |
| 42 | if ok: |
| 43 | norm_out = torch.empty_like(hidden_states) |
| 44 | # With norm_out provided, the kernel writes the new residual |
| 45 | # (all_reduce(hidden_states) + residual) into the hidden_states |
| 46 | # buffer and the normalized result into norm_out. |
| 47 | flashinfer_trtllm_fused_allreduce_norm( |
| 48 | allreduce_in=hidden_states, |
| 49 | residual=residual, |
| 50 | rms_gamma=norm.weight, |
| 51 | rms_eps=norm.variance_epsilon, |
| 52 | world_size=tp_size, |
| 53 | weight_bias=0.0, # standard RMSNorm (Gemma would use 1.0) |
| 54 | launch_with_pdl=True, |
| 55 | fp32_acc=True, |
| 56 | max_token_num=max_token_num, |
| 57 | pattern_code=_AR_RESIDUAL_RMS_NORM, |
| 58 | norm_out=norm_out, |
| 59 | ) |
| 60 | return norm_out, hidden_states |
| 61 | |
| 62 | reduced = tensor_model_parallel_all_reduce(hidden_states) |
| 63 | return norm(reduced, residual) |