optimize linear keepx quant (#10332)

Co-authored-by: Pan Zhaowu <panzhaowu@baidu.com>

Author: phlrain

paddlenlp/transformers/deepseek_v2/fp8_linear.py

@@ -395,7 +395,7 @@ def forward(ctx, x, w1, w2):
         )
 
         # ===== o3 = deep_gemm(o2_fp8, w2_t_fp8) =====
-        w2_fp8, w2_sacle, w2_t_fp8, w2_t_scale = kitchen_quant(
+        _, _, w2_t_fp8, w2_t_scale = kitchen_quant(
             w2, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_transpose=True
         )
         o3 = paddle.empty([o2_fp8.shape[0], w2_t_fp8.shape[0]], dtype=o2.dtype)
@@ -426,8 +426,7 @@ def forward(ctx, x, w1, w2):
             # w1_fp8,
             # w1_sacle,
             o1,
-            w2_fp8,
-            w2_sacle,
+            w2,
             paddle.to_tensor(x_orig_shape, dtype="int64", place=paddle.CPUPlace()),
         )
         return o3
@@ -438,9 +437,13 @@ def backward(ctx, do3):
         do3_orig_shape = do3.shape
         do3 = do3.reshape([-1, do3_orig_shape[-1]])
 
-        x_t_fp8, x_t_scale, w1, o1, w2_fp8, w2_sacle, x_orig_shape = ctx.saved_tensor()
+        x_t_fp8, x_t_scale, w1, o1, w2, x_orig_shape = ctx.saved_tensor()
         x_orig_shape = x_orig_shape.numpy()
 
+        w2_fp8, w2_scale = kitchen_quant(
+            w2, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_transpose=False
+        )
+
         # ===== [recompute] o2 = swiglu(o1) =====
         # TODO: [Fusion] swiglu + transpose + quant
         o2 = swiglu(o1)
@@ -454,7 +457,7 @@ def backward(ctx, do3):
             do3, backend=kitchen.ops.Backend.CUTLASS, is_1d_scaled=True, return_transpose=False
         )
         do2 = paddle.empty([do3_fp8.shape[0], w2_fp8.shape[0]], do3.dtype)
-        deep_gemm.gemm_fp8_fp8_bf16_nt((do3_fp8, do3_scale), (w2_fp8, w2_sacle), do2)
+        deep_gemm.gemm_fp8_fp8_bf16_nt((do3_fp8, do3_scale), (w2_fp8, w2_scale), do2)
 
         # ===== dw2 = deep_gemm(o2_t_fp8, do3_t_fp8)
         if o2_t.shape[-1] % 128 != 0 or o2_t.shape[-1] % 512 != 0:
@@ -549,3 +552,25 @@ def __init__(self, config: DeepseekV2Config, hidden_size=None, intermediate_size
 
     def forward(self, x):
         return Fuse_FFN_FP8_Func.apply(x, self.w1, self.w2)
+
+
+class FusedFP8DeepseekV2MLP(paddle.nn.Layer):
+    def __init__(self, config: DeepseekV2Config, hidden_size=None, intermediate_size=None, is_moe=False):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
+        self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
+
+        self.w1 = self.create_parameter(
+            shape=[4 * self.hidden_size, self.intermediate_size * 2],
+            dtype="bfloat16",
+            is_bias=False,
+        )
+        self.w2 = self.create_parameter(
+            shape=[4 * self.intermediate_size, self.hidden_size],
+            dtype="bfloat16",
+            is_bias=False,
+        )
+
+    def forward(self, x):
+        return Fuse_FFN_FP8_Func.apply(x, self.w1, self.w2)