diff --git a/vllm/compilation/passes/fusion/collective_fusion.py b/vllm/compilation/passes/fusion/collective_fusion.py index 55a5a2e5df..a9b64adcb3 100644 --- a/vllm/compilation/passes/fusion/collective_fusion.py +++ b/vllm/compilation/passes/fusion/collective_fusion.py @@ -53,7 +53,7 @@ class GEMMReduceScatterPattern(BasePattern): gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter( mul, mm_weight, - "avg", + "sum", scatter_dim=0, group_name=self.tp.device_group.group_name, ) @@ -150,7 +150,7 @@ class ScaledMMReduceScatterPattern(BasePattern): mat2, scale_a, scale_b, - "avg", + "sum", scatter_dim, # orig_scatter_dim scatter_dim, # scatter_dim_after_maybe_reshape self.tp.device_group.group_name, @@ -285,7 +285,7 @@ class CutlassScaledMMReduceScatterPattern(BasePattern): mat2, scale_a, scale_b, - "avg", + "sum", scatter_dim, # orig_scatter_dim scatter_dim, # scatter_dim_after_maybe_reshape self.tp.device_group.group_name,