[Bugfix] Guard mixed-dtype allreduce RMSNorm quant fusions (#48330)
Signed-off-by: hcenteno <hugo.centeno@estudiantat.upc.edu>
Hugo Centeno committed
Jul 12, 2026 at 11:39 UTC
5f8e73cb8b8d41f7a2a5168cddf5b772888fa991
2 files changed
+44
-3
tests/compile/passes/distributed/test_fusion_all_reduce.py
+32
-1
@@ -222,6 +222,25 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
222
]
223
224
225
+class TestAllReduceGemmaRMSNormStaticQuantFP8Model(
226
+ TestAllReduceRMSNormStaticQuantFP8Model
227
+):
228
+ def __init__(
229
+ self,
230
+ hidden_size=16,
231
+ token_num=16,
232
+ eps=1e-6,
233
+ dtype: torch.dtype = torch.float16,
234
+ ):
235
+ super().__init__(hidden_size, token_num, eps, dtype)
236
+ self.norm = [GemmaRMSNorm(hidden_size, eps) for _ in range(4)]
237
+ for norm in self.norm:
238
+ norm.weight.requires_grad_(False)
239
+
240
+ def ops_in_model_before(self):
241
+ return [torch.ops.vllm.all_reduce.default]
242
+
243
+
244
class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
245
"""Exercises the new ROCm AITER AR+RMS+per-group-FP8-quant patterns.
246
@@ -416,6 +435,15 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
435
reason="Not supported on ROCm platform",
436
),
437
),
438
+ pytest.param(
439
+ TestAllReduceGemmaRMSNormStaticQuantFP8Model,
440
+ True,
441
+ False,
442
+ marks=pytest.mark.skipif(
443
+ current_platform.is_rocm(),
444
+ reason="Not supported on ROCm platform",
445
+ ),
446
+ ),
447
pytest.param(
448
TestAllReduceRMSNormStaticQuantFP8Model,
449
False,
@@ -606,7 +634,10 @@ def all_reduce_fusion_pass_on_test_model(
634
)
635
backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
636
backend.check_after_ops(model.ops_in_model_after())
609
- if test_model_cls is TestAllReduceGemmaRMSNormModel:
637
+ if test_model_cls in (
638
+ TestAllReduceGemmaRMSNormModel,
639
+ TestAllReduceGemmaRMSNormStaticQuantFP8Model,
640
+ ):
641
fused_op = torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default
642
fused_nodes = list(find_op_nodes(fused_op, backend.graph_post_pass))
643
assert fused_nodes
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+12
-2
@@ -752,7 +752,12 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
752
return allreduce[4], allreduce[2]
753
754
pm.register_replacement(
755
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
755
+ pattern,
756
+ replacement,
757
+ self.get_inputs(),
758
+ pm.fwd_only,
759
+ pm_pass,
760
+ extra_check=_norm_input_weight_dtype_match,
761
)
762
763
@@ -941,7 +946,12 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
946
return allreduce[4], allreduce[2], allreduce[5]
947
948
pm.register_replacement(
944
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
949
+ pattern,
950
+ replacement,
951
+ self.get_inputs(),
952
+ pm.fwd_only,
953
+ pm_pass,
954
+ extra_check=_norm_input_weight_dtype_match,
955
)
956
957