From abad32a718d2f9e3af7b18818eb474322eaadfc4 Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Mon, 12 Oct 2026 22:53:08 +0300 Subject: [PATCH] [ROCm][Perf] Fix Gemma AR+RMS prefill fallback (#58547) Signed-off-by: Mikko Tukiainen Signed-off-by: Mikko Tukiainen --- .buildkite/test_areas/distributed.yaml | 2 + .../distributed/test_fusion_all_reduce.py | 65 ++++++++++++-- tests/distributed/test_aiter_rms_fallback.py | 89 +++++++++++++++++++ vllm/_aiter_ops.py | 16 +++- .../passes/fusion/allreduce_rms_fusion.py | 58 +++++++++++- vllm/kernels/aiter_ops.py | 3 + 6 files changed, 220 insertions(+), 13 deletions(-) create mode 100644 tests/distributed/test_aiter_rms_fallback.py diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 0b1a73136b..f541c42c93 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.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 diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index 1c99cc2533..9e2b969575 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/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, diff --git a/tests/distributed/test_aiter_rms_fallback.py b/tests/distributed/test_aiter_rms_fallback.py new file mode 100644 index 0000000000..4a0fc67873 --- /dev/null +++ b/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) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 0ed8490618..8af999906f 100644 --- a/vllm/_aiter_ops.py +++ b/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] diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index b2961debce..59757deb26 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/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, diff --git a/vllm/kernels/aiter_ops.py b/vllm/kernels/aiter_ops.py index 402c9cac17..9817f5adc9 100644 --- a/vllm/kernels/aiter_ops.py +++ b/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)