fix weight_only_linear_pass

PaddlePaddle · Dec 25, 2023 · 97e9192 · 97e9192
1 parent 54d4bf3
commit 97e9192
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/paddle/fluid/pir/transforms/fusion/fused_weight_only_linear_pass.cc b/paddle/fluid/pir/transforms/fusion/fused_weight_only_linear_pass.cc
@@ -96,9 +96,14 @@ class FusedWeightOnlyLinearPattern
           return getSMVersion();
         });
 
+    const auto &group_size_attr = res.Attr(
+        [](const pir::drr::MatchContext &match_ctx) -> int { return -1; });
+
     const auto &weight_quantize =
         res.Op(paddle::dialect::WeightQuantizeOp::name(),
-               {{"algo", weight_only_int8_attr}, {"arch", arch_attr}});
+               {{"algo", weight_only_int8_attr},
+                {"arch", arch_attr},
+                {"group_size", group_size_attr}});
     weight_quantize({&res.Tensor("w")},
                     {&res.Tensor("quanted_weight_tensor"),
                      &res.Tensor("weight_scale_tensor")});
@@ -110,7 +115,9 @@ class FusedWeightOnlyLinearPattern
 
     const auto &weight_only_linear =
         res.Op(paddle::dialect::WeightOnlyLinearOp::name(),
-               {{"weight_dtype", weight_dtype_attr}, {"arch", arch_attr}});
+               {{"weight_dtype", weight_dtype_attr},
+                {"arch", arch_attr},
+                {"group_size", group_size_attr}});
     weight_only_linear({&res.Tensor("x"),
                         &res.Tensor("quanted_weight_tensor"),
                         &res.Tensor("bias"),