refactor linear

tests/ut/ops/test_linear.py

@@ -7,8 +7,7 @@
 
 from vllm_ascend import ascend_config
 from vllm_ascend.distributed import parallel_state
-from vllm_ascend.ops.linear import (AscendColumnParallelLinear,
-                                    AscendMergedColumnParallelLinear,
+from vllm_ascend.ops.linear import (AscendMergedColumnParallelLinear,
                                     AscendRowParallelLinear)
 
 
@@ -56,8 +55,7 @@ def test_mlp_optimize(self):
             output_size=8,
             prefix="down_proj",
         )
-        self.assertEqual(linear.comm_group, parallel_state._MLP_TP)
-        self.assertEqual(linear.forward_type, "mlp_tp")
+        self.assertEqual(linear.custom_op.comm_group, parallel_state._MLP_TP)
 
         input_tensor = torch.randn(16, 8)
         linear(input_tensor)
@@ -71,34 +69,23 @@ def test_oproj_tp(self):
             output_size=8,
             prefix="o_proj",
         )
-        self.assertEqual(linear.comm_group, parallel_state._OTP)
-        self.assertEqual(linear.forward_type, "oproj_tp")
+        self.assertEqual(linear.custom_op.comm_group, parallel_state._OTP)
 
         input_tensor = torch.randn(16, 8)
         linear(input_tensor)
 
 
-class TestAscendColumnParallelLinear(BaseLinearTest):
-
-    def test_mlp_tp_init(self):
-        linear = AscendColumnParallelLinear(
-            input_size=16,
-            output_size=8,
-            prefix="down_proj",
-        )
-        self.assertEqual(linear.comm_group, parallel_state._MLP_TP)
-
-
 class TestAscendMergedColumnParallelLinear(BaseLinearTest):
 
     def test_merged_mlp_tp_init(self):
+        os.environ["VLLM_ASCEND_ENABLE_MLP_OPTIMIZE"] = "1"
+
         linear = AscendMergedColumnParallelLinear(
             input_size=16,
             output_sizes=[8, 8],
             prefix="gate_up_proj",
         )
-        self.assertEqual(linear.comm_group, parallel_state._MLP_TP)
-        self.assertEqual(linear.forward_type, "mlp_tp")
+        self.assertEqual(linear.custom_op.comm_group, parallel_state._MLP_TP)
 
 
 if __name__ == '__main__':