diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index b86018a755..1aac4b2bec 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -222,6 +222,25 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module): ] +class TestAllReduceGemmaRMSNormStaticQuantFP8Model( + TestAllReduceRMSNormStaticQuantFP8Model +): + def __init__( + self, + hidden_size=16, + token_num=16, + eps=1e-6, + dtype: torch.dtype = torch.float16, + ): + 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) + + def ops_in_model_before(self): + return [torch.ops.vllm.all_reduce.default] + + class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module): """Exercises the new ROCm AITER AR+RMS+per-group-FP8-quant patterns. @@ -416,6 +435,15 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module): reason="Not supported on ROCm platform", ), ), + pytest.param( + TestAllReduceGemmaRMSNormStaticQuantFP8Model, + True, + False, + marks=pytest.mark.skipif( + current_platform.is_rocm(), + reason="Not supported on ROCm platform", + ), + ), pytest.param( TestAllReduceRMSNormStaticQuantFP8Model, False, @@ -606,7 +634,10 @@ 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 is TestAllReduceGemmaRMSNormModel: + if test_model_cls in ( + TestAllReduceGemmaRMSNormModel, + TestAllReduceGemmaRMSNormStaticQuantFP8Model, + ): fused_op = torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default fused_nodes = list(find_op_nodes(fused_op, backend.graph_post_pass)) assert fused_nodes diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index ab40002892..906dea8101 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -752,7 +752,12 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern): return allreduce[4], allreduce[2] pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass + pattern, + replacement, + self.get_inputs(), + pm.fwd_only, + pm_pass, + extra_check=_norm_input_weight_dtype_match, ) @@ -941,7 +946,12 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern): return allreduce[4], allreduce[2], allreduce[5] pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass + pattern, + replacement, + self.get_inputs(), + pm.fwd_only, + pm_pass, + extra_check=_norm_input_weight_dtype_match, )