[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>
Taoyu Zhu committed
Sep 8, 2026 at 15:21 UTC
ce6c241ddcc722ee7f26fb4db132461150dad18e
3 files changed
+77
-73
vllm/_aiter_ops.py
+7
-73
@@ -981,33 +981,12 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
981
residual: torch.Tensor,
982
weight: torch.Tensor,
983
epsilon: float,
984
+ gemma_norm: bool = False,
985
) -> tuple[torch.Tensor, torch.Tensor]:
986
aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
987
assert aiter_ar is not None, "aiter allreduce must be initialized"
988
ca = aiter_ar.aiter_ca
988
-
989
- total_bytes = input_.numel() * input_.element_size()
990
- hidden_dim = input_.shape[-1]
991
- token_num = input_.numel() // hidden_dim
992
- if input_.dtype in (torch.bfloat16, torch.float16):
993
- pack_size = 16 // input_.element_size()
994
- hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
995
- else:
996
- hidden_ok = False
997
- token_ok = token_num <= 80
998
- world_size = ca.world_size
999
- full_nvlink = ca.fully_connected
1000
-
1001
- if world_size == 2:
1002
- size_ok = True
1003
- elif full_nvlink and world_size <= 4:
1004
- size_ok = total_bytes < 256 * 1024
1005
- elif full_nvlink and world_size <= 8:
1006
- size_ok = total_bytes < 128 * 1024
1007
- else:
1008
- size_ok = False
1009
-
1010
- use_1stage = hidden_ok and token_ok and size_ok
989
+ use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_)
990
991
result = ca.custom_fused_ar_rms(
992
input_,
@@ -1015,6 +994,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
994
weight,
995
epsilon,
996
use_1stage=use_1stage,
997
+ gemma_norm=gemma_norm,
998
)
999
assert result is not None
1000
return result[0], result[1]
@@ -1025,6 +1005,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_fake(
1005
residual: torch.Tensor,
1006
weight: torch.Tensor,
1007
epsilon: float,
1008
+ gemma_norm: bool = False,
1009
) -> tuple[torch.Tensor, torch.Tensor]:
1010
return torch.empty_like(input_), torch.empty_like(residual)
1011
@@ -1036,38 +1017,12 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl(
1017
epsilon: float,
1018
group_size: int,
1019
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1039
- """Fused AllReduce + RMSNorm + per-group FP8 quant.
1040
-
1041
- Mirrors the eligibility logic of ``_rocm_aiter_fused_allreduce_rmsnorm_impl``
1042
- for the 1-stage vs 2-stage AITER kernel dispatch (both variants run inside
1043
- AITER, the only choice we make here is the launcher to call into).
1044
- """
1020
+ """Fused AllReduce + RMSNorm + per-group FP8 quant."""
1021
aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
1022
assert aiter_ar is not None, "aiter allreduce must be initialized"
1023
ca = aiter_ar.aiter_ca
1024
1049
- total_bytes = input_.numel() * input_.element_size()
1050
- hidden_dim = input_.shape[-1]
1051
- token_num = input_.numel() // hidden_dim
1052
- if input_.dtype in (torch.bfloat16, torch.float16):
1053
- pack_size = 16 // input_.element_size()
1054
- hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
1055
- else:
1056
- hidden_ok = False
1057
- token_ok = token_num <= 80
1058
- world_size = ca.world_size
1059
- full_nvlink = ca.fully_connected
1060
-
1061
- if world_size == 2:
1062
- size_ok = True
1063
- elif full_nvlink and world_size <= 4:
1064
- size_ok = total_bytes < 256 * 1024
1065
- elif full_nvlink and world_size <= 8:
1066
- size_ok = total_bytes < 128 * 1024
1067
- else:
1068
- size_ok = False
1069
-
1070
- use_1stage = hidden_ok and token_ok and size_ok
1025
+ use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_)
1026
1027
result = ca.fused_ar_rms_per_group_quant(
1028
input_,
@@ -1119,28 +1074,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl(
1074
assert aiter_ar is not None, "aiter allreduce must be initialized"
1075
ca = aiter_ar.aiter_ca
1076
1122
- total_bytes = input_.numel() * input_.element_size()
1123
- hidden_dim = input_.shape[-1]
1124
- token_num = input_.shape[0]
1125
- if input_.dtype in (torch.bfloat16, torch.float16):
1126
- pack_size = 16 // input_.element_size()
1127
- hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
1128
- else:
1129
- hidden_ok = False
1130
- token_ok = token_num <= 80
1131
- world_size = ca.world_size
1132
- full_nvlink = ca.fully_connected
1133
-
1134
- if world_size == 2:
1135
- size_ok = True
1136
- elif full_nvlink and world_size <= 4:
1137
- size_ok = total_bytes < 256 * 1024
1138
- elif full_nvlink and world_size <= 8:
1139
- size_ok = total_bytes < 128 * 1024
1140
- else:
1141
- size_ok = False
1142
-
1143
- use_1stage = hidden_ok and token_ok and size_ok
1077
+ use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_)
1078
1079
result = ca.fused_ar_rms_per_group_quant(
1080
input_,
vllm/distributed/device_communicators/aiter_custom_all_reduce.py
+34
@@ -56,6 +56,40 @@ class AiterCustomAllreduce:
56
def custom_all_reduce(self, inp: torch.Tensor) -> torch.Tensor | None:
57
return self._impl.custom_all_reduce(inp)
58
59
+ def use_1stage_fused_ar_rms(self, inp: torch.Tensor) -> bool:
60
+ """Whether AITER's fused allreduce+RMSNorm runs as its one-stage kernel.
61
+
62
+ Mirrors the launcher contract of aiter's ``fused_allreduce_rmsnorm``
63
+ (csrc/include/custom_all_reduce.cuh): rows of 16-byte packs, at most
64
+ 1024 packs per row, at most 80 tokens, and the byte cap of the
65
+ one-stage custom allreduce for this TP size and topology. Outside it
66
+ the fused op runs the two-stage variant (cross-device reduce-scatter
67
+ + local norm), which is slower than an explicit ``all_reduce`` + norm,
68
+ so callers that can fall back should require this. Capture-static:
69
+ depends only on shape, dtype, TP size and topology.
70
+ """
71
+ hidden_dim = inp.shape[-1]
72
+ # Token cap first: prefill-sized inputs leave here with one comparison.
73
+ if inp.numel() // hidden_dim > 80:
74
+ return False
75
+ if inp.dtype not in (torch.bfloat16, torch.float16):
76
+ return False
77
+ pack_size = 16 // inp.element_size()
78
+ if hidden_dim % pack_size != 0 or hidden_dim // pack_size > 1024:
79
+ return False
80
+ ca = self._impl
81
+ world_size = ca.world_size
82
+ if world_size == 2:
83
+ return True
84
+ if not ca.fully_connected:
85
+ return False
86
+ total_bytes = inp.numel() * inp.element_size()
87
+ if world_size <= 4:
88
+ return total_bytes < 256 * 1024
89
+ if world_size <= 8:
90
+ return total_bytes < 128 * 1024
91
+ return False
92
+
93
def capture(self):
94
return self._impl.capture()
95
vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py
+36
@@ -23,6 +23,7 @@ from vllm.distributed.parallel_state import (
23
get_tp_group,
24
)
25
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
26
+from vllm.platforms import current_platform
27
28
MiB = 1024 * 1024
29
@@ -111,6 +112,30 @@ def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool
112
return True, max_token_num
113
114
115
+def _can_use_aiter_fused_ar_rms(hidden_states: torch.Tensor) -> bool:
116
+ if not current_platform.is_rocm():
117
+ return False
118
+ from vllm._aiter_ops import rocm_aiter_ops
119
+
120
+ if not rocm_aiter_ops.is_custom_all_reduce_enabled():
121
+ return False
122
+ if (
123
+ hidden_states.dim() != 2
124
+ or not hidden_states.is_contiguous()
125
+ or hidden_states.dtype not in _FI_SUPPORTED_DTYPES
126
+ ):
127
+ return False
128
+ aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
129
+ if aiter_ar is None or aiter_ar.disabled:
130
+ return False
131
+ total_bytes = hidden_states.numel() * hidden_states.element_size()
132
+ if total_bytes > aiter_ar.effective_max_size():
133
+ return False
134
+ if not aiter_ar.should_custom_ar(hidden_states):
135
+ return False
136
+ return aiter_ar.use_1stage_fused_ar_rms(hidden_states)
137
+
138
+
139
def fused_allreduce_gemma_rms_norm(
140
hidden_states: torch.Tensor,
141
residual: torch.Tensor,
@@ -149,6 +174,17 @@ def fused_allreduce_gemma_rms_norm(
174
)
175
return norm_out, hidden_states
176
177
+ if _can_use_aiter_fused_ar_rms(hidden_states):
178
+ from vllm._aiter_ops import rocm_aiter_ops
179
+
180
+ return rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()(
181
+ input_=hidden_states,
182
+ residual=residual,
183
+ weight=norm.weight,
184
+ epsilon=norm.variance_epsilon,
185
+ gemma_norm=True,
186
+ )
187
+
188
# Fallback: explicit all-reduce + GemmaRMSNorm (matches the unfused model).
189
reduced = tensor_model_parallel_all_reduce(hidden_states)
190
return norm(reduced, residual)