From ce6c241ddcc722ee7f26fb4db132461150dad18e Mon Sep 17 00:00:00 2001 From: Taoyu Zhu Date: Tue, 8 Sep 2026 15:21:49 +0800 Subject: [PATCH] [ROCm][Perf][M3] Fused allreduce+GemmaRMSNorm fast path (#54787) Signed-off-by: Taoyu Zhu Co-authored-by: Shanshan Shen <467638484@qq.com> --- vllm/_aiter_ops.py | 80 ++----------------- .../aiter_custom_all_reduce.py | 34 ++++++++ .../layers/fused_allreduce_gemma_rms_norm.py | 36 +++++++++ 3 files changed, 77 insertions(+), 73 deletions(-) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 2055fad22f..9bb018b3d3 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -981,33 +981,12 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( residual: torch.Tensor, weight: torch.Tensor, epsilon: float, + gemma_norm: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: aiter_ar = rocm_aiter_ops.get_aiter_allreduce() assert aiter_ar is not None, "aiter allreduce must be initialized" ca = aiter_ar.aiter_ca - - total_bytes = input_.numel() * input_.element_size() - hidden_dim = input_.shape[-1] - token_num = input_.numel() // hidden_dim - if input_.dtype in (torch.bfloat16, torch.float16): - pack_size = 16 // input_.element_size() - hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024 - else: - hidden_ok = False - token_ok = token_num <= 80 - world_size = ca.world_size - full_nvlink = ca.fully_connected - - if world_size == 2: - size_ok = True - elif full_nvlink and world_size <= 4: - size_ok = total_bytes < 256 * 1024 - elif full_nvlink and world_size <= 8: - size_ok = total_bytes < 128 * 1024 - else: - size_ok = False - - use_1stage = hidden_ok and token_ok and size_ok + use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_) result = ca.custom_fused_ar_rms( input_, @@ -1015,6 +994,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( weight, epsilon, use_1stage=use_1stage, + gemma_norm=gemma_norm, ) assert result is not None return result[0], result[1] @@ -1025,6 +1005,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_fake( residual: torch.Tensor, weight: torch.Tensor, epsilon: float, + gemma_norm: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: return torch.empty_like(input_), torch.empty_like(residual) @@ -1036,38 +1017,12 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( epsilon: float, group_size: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Fused AllReduce + RMSNorm + per-group FP8 quant. - - Mirrors the eligibility logic of ``_rocm_aiter_fused_allreduce_rmsnorm_impl`` - for the 1-stage vs 2-stage AITER kernel dispatch (both variants run inside - AITER, the only choice we make here is the launcher to call into). - """ + """Fused AllReduce + RMSNorm + per-group FP8 quant.""" aiter_ar = rocm_aiter_ops.get_aiter_allreduce() assert aiter_ar is not None, "aiter allreduce must be initialized" ca = aiter_ar.aiter_ca - total_bytes = input_.numel() * input_.element_size() - hidden_dim = input_.shape[-1] - token_num = input_.numel() // hidden_dim - if input_.dtype in (torch.bfloat16, torch.float16): - pack_size = 16 // input_.element_size() - hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024 - else: - hidden_ok = False - token_ok = token_num <= 80 - world_size = ca.world_size - full_nvlink = ca.fully_connected - - if world_size == 2: - size_ok = True - elif full_nvlink and world_size <= 4: - size_ok = total_bytes < 256 * 1024 - elif full_nvlink and world_size <= 8: - size_ok = total_bytes < 128 * 1024 - else: - size_ok = False - - use_1stage = hidden_ok and token_ok and size_ok + use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_) result = ca.fused_ar_rms_per_group_quant( input_, @@ -1119,28 +1074,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl( assert aiter_ar is not None, "aiter allreduce must be initialized" ca = aiter_ar.aiter_ca - total_bytes = input_.numel() * input_.element_size() - hidden_dim = input_.shape[-1] - token_num = input_.shape[0] - if input_.dtype in (torch.bfloat16, torch.float16): - pack_size = 16 // input_.element_size() - hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024 - else: - hidden_ok = False - token_ok = token_num <= 80 - world_size = ca.world_size - full_nvlink = ca.fully_connected - - if world_size == 2: - size_ok = True - elif full_nvlink and world_size <= 4: - size_ok = total_bytes < 256 * 1024 - elif full_nvlink and world_size <= 8: - size_ok = total_bytes < 128 * 1024 - else: - size_ok = False - - use_1stage = hidden_ok and token_ok and size_ok + use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_) result = ca.fused_ar_rms_per_group_quant( input_, diff --git a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py index 63e06b77a7..be2c901281 100644 --- a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py +++ b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py @@ -56,6 +56,40 @@ class AiterCustomAllreduce: def custom_all_reduce(self, inp: torch.Tensor) -> torch.Tensor | None: return self._impl.custom_all_reduce(inp) + def use_1stage_fused_ar_rms(self, inp: torch.Tensor) -> bool: + """Whether AITER's fused allreduce+RMSNorm runs as its one-stage kernel. + + Mirrors the launcher contract of aiter's ``fused_allreduce_rmsnorm`` + (csrc/include/custom_all_reduce.cuh): rows of 16-byte packs, at most + 1024 packs per row, at most 80 tokens, and the byte cap of the + one-stage custom allreduce for this TP size and topology. Outside it + the fused op runs the two-stage variant (cross-device reduce-scatter + + local norm), which is slower than an explicit ``all_reduce`` + norm, + so callers that can fall back should require this. Capture-static: + depends only on shape, dtype, TP size and topology. + """ + hidden_dim = inp.shape[-1] + # Token cap first: prefill-sized inputs leave here with one comparison. + if inp.numel() // hidden_dim > 80: + return False + if inp.dtype not in (torch.bfloat16, torch.float16): + return False + pack_size = 16 // inp.element_size() + if hidden_dim % pack_size != 0 or hidden_dim // pack_size > 1024: + return False + ca = self._impl + world_size = ca.world_size + if world_size == 2: + return True + if not ca.fully_connected: + return False + total_bytes = inp.numel() * inp.element_size() + if world_size <= 4: + return total_bytes < 256 * 1024 + if world_size <= 8: + return total_bytes < 128 * 1024 + return False + def capture(self): return self._impl.capture() diff --git a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py index 921c721c54..7046bfc991 100644 --- a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py +++ b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py @@ -23,6 +23,7 @@ from vllm.distributed.parallel_state import ( get_tp_group, ) from vllm.model_executor.layers.layernorm import GemmaRMSNorm +from vllm.platforms import current_platform MiB = 1024 * 1024 @@ -111,6 +112,30 @@ def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool return True, max_token_num +def _can_use_aiter_fused_ar_rms(hidden_states: torch.Tensor) -> bool: + if not current_platform.is_rocm(): + return False + from vllm._aiter_ops import rocm_aiter_ops + + if not rocm_aiter_ops.is_custom_all_reduce_enabled(): + return False + if ( + hidden_states.dim() != 2 + or not hidden_states.is_contiguous() + or hidden_states.dtype not in _FI_SUPPORTED_DTYPES + ): + return False + aiter_ar = rocm_aiter_ops.get_aiter_allreduce() + if aiter_ar is None or aiter_ar.disabled: + return False + total_bytes = hidden_states.numel() * hidden_states.element_size() + if total_bytes > aiter_ar.effective_max_size(): + return False + if not aiter_ar.should_custom_ar(hidden_states): + return False + return aiter_ar.use_1stage_fused_ar_rms(hidden_states) + + def fused_allreduce_gemma_rms_norm( hidden_states: torch.Tensor, residual: torch.Tensor, @@ -149,6 +174,17 @@ def fused_allreduce_gemma_rms_norm( ) return norm_out, hidden_states + if _can_use_aiter_fused_ar_rms(hidden_states): + from vllm._aiter_ops import rocm_aiter_ops + + return rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()( + input_=hidden_states, + residual=residual, + weight=norm.weight, + epsilon=norm.variance_epsilon, + gemma_norm=True, + ) + # Fallback: explicit all-reduce + GemmaRMSNorm (matches the unfused model). reduced = tensor_model_parallel_all_reduce(hidden_states) return norm(reduced, residual)