[ROCM] Fix AITER Fused AllReduce RMSNorm for Transformers Backend (#49673)
Signed-off-by: BadrBasowid <badr.basowid@gmail.com>
BadrBasowid committed
Jul 24, 2026 at 22:16 UTC
8eac21a602d4959eb8478c430d7fe1ca5b57311f
2 files changed
+18
-4
tests/compile/fusions_e2e/test_tp2_ar_rms.py
+16
-2
@@ -181,8 +181,13 @@ def test_tp2_ar_rms_fp4_fusions(
181
182
@multi_gpu_test(num_gpus=2)
183
@pytest.mark.parametrize(
184
- "model_name, matches_fn, model_kwargs, hf_overrides",
185
- [llama3_8b, qwen3_a3b, gpt_oss_20b],
184
+ "model_name, matches_fn, model_kwargs, hf_overrides, model_impl",
185
+ [
186
+ (*llama3_8b, "auto"),
187
+ (*llama3_8b, "transformers"),
188
+ (*qwen3_a3b, "auto"),
189
+ (*gpt_oss_20b, "auto"),
190
+ ],
191
)
192
@pytest.mark.parametrize(
193
"attn_backend",
@@ -202,17 +207,26 @@ def test_tp2_ar_rms_fusions(
207
matches_fn: Callable[[int], Matches],
208
model_kwargs: dict,
209
hf_overrides: Callable[[int], dict],
210
+ model_impl: str,
211
attn_backend: AttentionBackendCase,
212
n_layers: int,
213
custom_ops: str,
214
inductor_graph_partition: bool,
215
run_e2e_fusion_test,
216
):
217
+ if model_impl == "transformers" and not current_platform.is_rocm():
218
+ pytest.skip("Transformers 3D AR+RMS regression is ROCm-only")
219
+
220
matches = matches_fn(n_layers)
221
+ if model_impl == "transformers":
222
+ # TODO(BadrBasowid): Match the vLLM backend's fusion count once the
223
+ # separate residual add and RMSNorm operations are fused.
224
+ matches = matches._replace(aiter_ar_rms_fusion=1)
225
226
# Reduce size of model and skip weight loading time
227
model_kwargs["hf_overrides"] = hf_overrides(n_layers)
228
model_kwargs["load_format"] = "dummy"
229
+ model_kwargs["model_impl"] = model_impl
230
model_kwargs["max_model_len"] = 1024
231
model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
232
model_kwargs["disable_custom_all_reduce"] = False
vllm/_aiter_ops.py
+2
-2
@@ -764,7 +764,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
764
765
total_bytes = input_.numel() * input_.element_size()
766
hidden_dim = input_.shape[-1]
767
- token_num = input_.shape[0]
767
+ token_num = input_.numel() // hidden_dim
768
if input_.dtype in (torch.bfloat16, torch.float16):
769
pack_size = 16 // input_.element_size()
770
hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
@@ -824,7 +824,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl(
824
825
total_bytes = input_.numel() * input_.element_size()
826
hidden_dim = input_.shape[-1]
827
- token_num = input_.shape[0]
827
+ token_num = input_.numel() // hidden_dim
828
if input_.dtype in (torch.bfloat16, torch.float16):
829
pack_size = 16 // input_.element_size()
830
hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024