[ROCm][Perf][M3] Fused allreduce+GemmaRMSNorm fast path (#54787)

Signed-off-by: Taoyu Zhu <zhutaoyu97@gmail.com>
Co-authored-by: Shanshan Shen <467638484@qq.com>

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_,

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()
 

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)