[ROCm][Perf] Fix Gemma AR+RMS prefill fallback (#58547)

Signed-off-by: Mikko Tukiainen <Mikko.Tukiainen@amd.com>
Signed-off-by: Mikko Tukiainen <mikko.tukiainen@amd.com>

.buildkite/test_areas/distributed.yaml

@@ -522,6 +522,7 @@ steps:
       - tests/distributed/test_rocm_aiter_custom_ar.py
       - tests/distributed/test_rocm_quick_reduce.py
       - tests/distributed/test_quick_all_reduce.py
+      - tests/distributed/test_aiter_rms_fallback.py
       - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
       - tests/utils.py
       - examples/rl/rlhf_async_new_apis.py
@@ -539,6 +540,7 @@ steps:
       - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
       - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
       - pytest -v -s tests/distributed/test_quick_all_reduce.py
+      - pytest -v -s tests/distributed/test_aiter_rms_fallback.py
 
 - label: ":nvidia: (B200) Distributed"
   key: distributed-tests-2xb200

tests/compile/passes/distributed/test_fusion_all_reduce.py

@@ -177,10 +177,12 @@ class TestAllReduceGemmaRMSNormModel(torch.nn.Module):
         token_num=16,
         eps=1e-6,
         dtype: torch.dtype = torch.float16,
+        use_aiter: bool = False,
     ):
         super().__init__()
         self.hidden_size = hidden_size
         self.eps = eps
+        self.use_aiter = use_aiter
         self.norm = [GemmaRMSNorm(hidden_size, eps) for _ in range(4)]
         # Non-trivial weight (~Gemma range) so (1 + w) exercises the scale path.
         for n in self.norm:
@@ -210,6 +212,8 @@ class TestAllReduceGemmaRMSNormModel(torch.nn.Module):
         return [torch.ops.vllm.all_reduce.default]
 
     def ops_in_model_after(self):
+        if self.use_aiter:
+            return [rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()]
         return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]
 
 
@@ -383,6 +387,32 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
         ]
 
 
+class TestAiterAllReduceGemmaRMSNormGroupQuantFP8Model(
+    TestAiterAllReduceRMSNormGroupQuantFP8Model
+):
+    """Gemma norms feeding the per-group FP8 quant sites.
+
+    The fused AR+RMS+quant op has no ``gemma_norm`` mode, so these sites must
+    keep matching the quant-fused patterns with the pre-added ``1 + weight``.
+    The Gemma AR+RMS-only patterns match a prefix of the same subgraph, so if
+    they were registered ahead of the quant-fused ones they would consume the
+    ``all_reduce`` and strand ``rocm_aiter_group_fp8_quant`` as a separate
+    kernel -- which the ``fully_replaced=True`` check below would catch.
+    """
+
+    def __init__(
+        self,
+        hidden_size=128,
+        token_num=16,
+        eps=1e-6,
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        super().__init__(hidden_size, token_num, eps, dtype)
+        self.norm = [GemmaRMSNorm(hidden_size, eps) for _ in range(4)]
+        for norm in self.norm:
+            norm.weight.requires_grad_(False)
+
+
 class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
     def __init__(
         self, hidden_size=16, token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16
@@ -450,10 +480,10 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
         pytest.param(
             TestAllReduceGemmaRMSNormModel,
             False,
-            False,
+            IS_AITER_FOUND,
             marks=pytest.mark.skipif(
-                current_platform.is_rocm(),
-                reason="Not supported on ROCm platform",
+                current_platform.is_rocm() and not IS_AITER_FOUND,
+                reason="AITER is required on ROCm",
             ),
         ),
         pytest.param(
@@ -646,7 +676,10 @@ def all_reduce_fusion_pass_on_test_model(
         )
 
         token_num = batch_size * seq_len
-        if test_model_cls is TestAllReduceRMSNormModel:
+        if test_model_cls in (
+            TestAllReduceRMSNormModel,
+            TestAllReduceGemmaRMSNormModel,
+        ):
             model = test_model_cls(
                 hidden_size, token_num, dtype=dtype, use_aiter=use_aiter
             )
@@ -667,9 +700,13 @@ def all_reduce_fusion_pass_on_test_model(
         )
         backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
         backend.check_after_ops(model.ops_in_model_after())
-        if test_model_cls in (
-            TestAllReduceGemmaRMSNormModel,
-            TestAllReduceGemmaRMSNormStaticQuantFP8Model,
+        if (
+            test_model_cls
+            in (
+                TestAllReduceGemmaRMSNormModel,
+                TestAllReduceGemmaRMSNormStaticQuantFP8Model,
+            )
+            and not use_aiter
         ):
             fused_op = torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default
             fused_nodes = list(find_op_nodes(fused_op, backend.graph_post_pass))
@@ -684,6 +721,13 @@ def all_reduce_fusion_pass_on_test_model(
 @pytest.mark.parametrize("hidden_size", [128])
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
+@pytest.mark.parametrize(
+    "test_model_cls",
+    [
+        TestAiterAllReduceRMSNormGroupQuantFP8Model,
+        TestAiterAllReduceGemmaRMSNormGroupQuantFP8Model,
+    ],
+)
 @pytest.mark.skipif(
     not current_platform.is_rocm(),
     reason="ROCm AITER AR+RMS+per-group-FP8-quant fusion is ROCm-only",
@@ -695,6 +739,7 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
     hidden_size: int,
     dtype: torch.dtype,
     enable_rms_norm_custom_op: bool,
+    test_model_cls: type[torch.nn.Module],
     monkeypatch: pytest.MonkeyPatch,
 ):
     """Sibling of ``test_all_reduce_fusion_pass_replace`` for the new
@@ -710,6 +755,10 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
       residual, DSv3.2 indexer fan-out; parametrized over
       ``rocm_aiter_group_fp8_quant``
       producer).
+
+    Parametrized over plain and Gemma norms: the Gemma AR+RMS-only patterns
+    match a prefix of these subgraphs, so this also pins their registration
+    order behind the quant-fused ones.
     """
     with monkeypatch.context() as m:
         m.setenv("VLLM_ROCM_USE_AITER", "1")
@@ -730,7 +779,7 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
             args=(
                 num_processes,
                 master_port,
-                TestAiterAllReduceRMSNormGroupQuantFP8Model,
+                test_model_cls,
                 batch_size,
                 seq_len,
                 hidden_size,

tests/distributed/test_aiter_rms_fallback.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+import types
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm._aiter_ops as aiter_ops
+import vllm.distributed as vllm_distributed
+import vllm.kernels  # noqa: F401
+
+
+def test_rocm_aiter_gemma_rmsnorm_falls_back_after_allreduce(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    class FakeAiterCA:
+        world_size = 8
+        fully_connected = True
+
+        def custom_fused_ar_rms(self, *args, **kwargs):
+            return None
+
+    class FakeAiterAllReduce:
+        aiter_ca = FakeAiterCA()
+
+        def use_1stage_fused_ar_rms(self, inp):
+            return False
+
+    class FakeQuickReduce:
+        disabled = False
+
+        def should_quick_allreduce(self, inp):
+            raise AssertionError("Gemma norm must not use fused QR+RMSNorm")
+
+    device_comm = SimpleNamespace(qr_comm=FakeQuickReduce())
+    monkeypatch.setattr(
+        vllm_distributed,
+        "get_tp_group",
+        lambda: SimpleNamespace(device_communicator=device_comm),
+    )
+    monkeypatch.setattr(
+        aiter_ops.rocm_aiter_ops,
+        "get_aiter_allreduce",
+        lambda: FakeAiterAllReduce(),
+    )
+    monkeypatch.setattr(
+        vllm_distributed,
+        "tensor_model_parallel_all_reduce",
+        lambda inp: inp + 2,
+    )
+
+    calls = []
+    fake_aiter = types.SimpleNamespace()
+
+    def fake_rmsnorm2d_fwd_with_add(
+        out,
+        inp,
+        residual,
+        residual_out,
+        weight,
+        epsilon,
+        *,
+        gemma_norm,
+    ):
+        calls.append((inp, weight, epsilon, gemma_norm))
+        residual_out.copy_(inp + residual)
+        out.copy_(residual_out)
+
+    fake_aiter.rmsnorm2d_fwd_with_add = fake_rmsnorm2d_fwd_with_add
+    monkeypatch.setitem(sys.modules, "aiter", fake_aiter)
+
+    inp = torch.randn(8192, 16, dtype=torch.bfloat16)
+    residual = torch.randn_like(inp)
+    weight = torch.randn(16, dtype=torch.bfloat16)
+
+    out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl(
+        inp, residual, weight, 1e-6, gemma_norm=True
+    )
+
+    assert len(calls) == 1
+    assert calls[0][0] is not inp
+    assert calls[0][1] is weight
+    assert calls[0][2] == 1e-6
+    assert calls[0][3] is True
+    torch.testing.assert_close(residual_out, inp + 2 + residual)
+    torch.testing.assert_close(out, residual_out)

vllm/_aiter_ops.py

@@ -1086,7 +1086,10 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
         hidden_dim = input_.shape[-1]
         row_size = hidden_dim * input_.element_size()
         fused_qr_rmsnorm_ok = (
-            qr_comm is not None
+            # aiter.qr_all_reduce_rmsnorm does not implement Gemma's
+            # (1 + weight) scaling.
+            not gemma_norm
+            and qr_comm is not None
             and not getattr(qr_comm, "disabled", True)
             and hasattr(qr_comm, "should_quick_allreduce")
             and qr_comm.should_quick_allreduce(input_)
@@ -1133,7 +1136,16 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
         use_1stage=use_1stage,
         gemma_norm=gemma_norm,
     )
-    assert result is not None
+    if result is None:
+        from vllm.distributed import tensor_model_parallel_all_reduce
+
+        return torch.ops.vllm_aiter.fused_add_rms_norm(
+            tensor_model_parallel_all_reduce(input_),
+            residual,
+            weight,
+            epsilon,
+            gemma_norm,
+        )
     return result[0], result[1]
 
 

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

@@ -1149,6 +1149,8 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
 
 # TODO: make BasePattern to inherit from VllmPatternReplacement
 class AiterAllreduceFusedRMSNormPattern(BasePattern, VllmPatternReplacement):
+    gemma_norm = False
+
     def __init__(
         self,
         epsilon: float,
@@ -1170,7 +1172,8 @@ class AiterAllreduceFusedRMSNormPattern(BasePattern, VllmPatternReplacement):
             input: torch.Tensor, weight: torch.Tensor
         ) -> tuple[torch.Tensor, torch.Tensor]:
             allreduce_output = tensor_model_parallel_all_reduce(input)
-            rms = vllm.ir.ops.rms_norm(allreduce_output, weight, self.epsilon)
+            norm_weight = weight.float() + 1.0 if self.gemma_norm else weight
+            rms = vllm.ir.ops.rms_norm(allreduce_output, norm_weight, self.epsilon)
 
             return rms, allreduce_output
 
@@ -1187,13 +1190,22 @@ class AiterAllreduceFusedRMSNormPattern(BasePattern, VllmPatternReplacement):
                 residual=residual,
                 weight=weight.to(input.dtype),
                 epsilon=self.epsilon,
+                gemma_norm=self.gemma_norm,
             )
             return allreduce[0], allreduce[1]
 
         return _replacement
 
 
+class AiterAllreduceFusedGemmaRMSNormPattern(AiterAllreduceFusedRMSNormPattern):
+    """Match all-reduce followed by GemmaRMSNorm without a residual."""
+
+    gemma_norm = True
+
+
 class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement):
+    gemma_norm = False
+
     def __init__(
         self,
         epsilon: float,
@@ -1216,8 +1228,9 @@ class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement):
             residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
         ) -> tuple[torch.Tensor, torch.Tensor]:
             allreduce_output = tensor_model_parallel_all_reduce(input)
+            norm_weight = weight.float() + 1.0 if self.gemma_norm else weight
             rms, residual = vllm.ir.ops.fused_add_rms_norm(
-                allreduce_output, residual, weight, self.epsilon
+                allreduce_output, residual, norm_weight, self.epsilon
             )
             return rms, residual
 
@@ -1233,12 +1246,19 @@ class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement):
                 residual=residual,
                 weight=weight.to(input.dtype),
                 epsilon=self.epsilon,
+                gemma_norm=self.gemma_norm,
             )
             return allreduce[0], allreduce[1]
 
         return _replacement
 
 
+class AiterAllreduceFusedAddGemmaRMSNormPattern(AiterAllreduceFusedAddRMSNormPattern):
+    """Match all-reduce followed by fused-add GemmaRMSNorm."""
+
+    gemma_norm = True
+
+
 class AiterAllreduceFusedAddRMSNormOutputOnlyPattern(
     AiterAllreduceFusedAddRMSNormPattern
 ):
@@ -1267,6 +1287,14 @@ class AiterAllreduceFusedAddRMSNormOutputOnlyPattern(
         return _replacement
 
 
+class AiterAllreduceFusedAddGemmaRMSNormOutputOnlyPattern(
+    AiterAllreduceFusedAddRMSNormOutputOnlyPattern
+):
+    """Match fused-add GemmaRMSNorm when its residual output is dead."""
+
+    gemma_norm = True
+
+
 class AiterAllreduceFusedRMSNormGroupQuantFP8Pattern(
     BasePattern, VllmPatternReplacement
 ):
@@ -1608,7 +1636,7 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
             # AR+RMS-only fusion runs first and consumes the all_reduce node,
             # leaving the trailing quant op stranded as an unfused kernel.
             # Register larger subgraphs first (DeepSeek indexer fan-out, then
-            # quant-only AR+RMS+quant, then AR+RMS-only).
+            # quant-only AR+RMS+quant, then Gemma AR+RMS, then AR+RMS-only).
             if supports_per_group_quant:
                 self.register(
                     AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern(
@@ -1632,6 +1660,30 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
                     )
                 )
 
+            # Must stay above the non-Gemma AR+RMS patterns below, which would
+            # otherwise match first.
+            self.register(
+                AiterAllreduceFusedGemmaRMSNormPattern(
+                    epsilon,
+                    self.model_dtype,
+                    self.device,
+                )
+            )
+            self.register(
+                AiterAllreduceFusedAddGemmaRMSNormPattern(
+                    epsilon,
+                    self.model_dtype,
+                    self.device,
+                )
+            )
+            self.register(
+                AiterAllreduceFusedAddGemmaRMSNormOutputOnlyPattern(
+                    epsilon,
+                    self.model_dtype,
+                    self.device,
+                )
+            )
+
             self.register(
                 AiterAllreduceFusedRMSNormPattern(
                     epsilon,

vllm/kernels/aiter_ops.py

@@ -136,6 +136,7 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_impl(
     residual: torch.Tensor,
     weight: torch.Tensor,
     variance_epsilon: float,
+    gemma_norm: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     from aiter import rmsnorm2d_fwd_with_add
 
@@ -157,6 +158,7 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_impl(
         residual_out,  # residual output
         weight,
         variance_epsilon,
+        gemma_norm=gemma_norm,
     )
     return out.reshape(x_shape), residual_out.reshape(residual_shape)
 
@@ -166,6 +168,7 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_fake(
     residual: torch.Tensor,
     weight: torch.Tensor,
     variance_epsilon: float,
+    gemma_norm: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     residual_out = torch.empty_like(residual)
     out = torch.empty_like(x)