[CI] fix

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

Author: zzhx1

tests/ut/torchair/models/test_torchair_deepseek_v2.py

@@ -118,7 +118,6 @@ def mock_distributed():
             patch("vllm_ascend.torchair.models.torchair_deepseek_v2.get_pp_group",
                   return_value=Mock(is_first_rank=False, is_last_rank=False)), \
             patch("vllm_ascend.torchair.ops.torchair_fused_moe.get_current_vllm_config", return_value=mock_vllm_config), \
-            patch("vllm_ascend.distributed.parallel_state.get_mlp_tp_group", return_value=mlp_tp_group), \
             patch.dict("vllm.distributed.parallel_state.__dict__", _TP=tp_group, _EP=ep_group, _DP=dp_group,
                        _PP=pp_group), \
             patch.dict("vllm_ascend.distributed.parallel_state.__dict__", _MC2=ep_group):
@@ -202,10 +201,6 @@ def test_torchair_deepseek_v2_mlp(mock_distributed, base_config):
                                 quant_config=None)
     assert isinstance(mlp.act_fn, TorchairDeepseekV2SiluAndMul)
 
-    x = torch.randn(2, 4, 128)
-    output = mlp(x)
-    assert output.shape == (2, 4, 128)
-
     with patch(
             "vllm_ascend.torchair.models.torchair_deepseek_v2.QuantizationConfig"
     ) as mock_quant_config:

vllm_ascend/ops/linear.py

@@ -195,24 +195,25 @@ def forward(
         input_,
         is_prefill: bool = True,
     ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
+        # Choose different forward function according to the type of TP group
+        if self.forward_type == "oproj_tp":
+            return self._forward_oproj_tp(input_)
+        elif self.forward_type == "mlp_tp":
+            return self._forward_mlp_tp(input_)
+        else:
+            return super().forward(input_)
 
+    # enable custom MLP tensor parallel
+    def _forward_mlp_tp(self, input_: torch.Tensor) -> torch.Tensor:
+        
         if self.input_is_parallel:
             input_parallel = input_
         else:
             splitted_input = split_tensor_along_last_dim(
                 input_, num_partitions=self.tp_size)
             input_parallel = splitted_input[self.tp_rank].contiguous()
+        
         assert self.quant_method is not None
-        # Choose different forward function according to the type of TP group
-        if self.forward_type == "oproj_tp":
-            return self._forward_oproj_tp(input_parallel)
-        elif self.forward_type == "mlp_tp":
-            return self._forward_mlp_tp(input_parallel)
-        else:
-            return RowParallelLinear.forward(self, input_)
-
-    #
-    def _forward_mlp_tp(self, input_parallel: torch.Tensor) -> torch.Tensor:
         bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
         output_parallel = self.quant_method.apply(self,
                                                   input_parallel,
@@ -224,12 +225,19 @@ def _forward_mlp_tp(self, input_parallel: torch.Tensor) -> torch.Tensor:
             return output
         return output, output_bias
 
-    # enable oproj tp forward function
+    # enable custom Oproj tensor parallel
     def _forward_oproj_tp(
         self,
-        input_parallel: torch.Tensor,
+        input_: torch.Tensor,
     ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
-
+        
+        if self.input_is_parallel:
+            input_parallel = input_
+        else:
+            splitted_input = split_tensor_along_last_dim(
+                input_, num_partitions=self.tp_size)
+            input_parallel = splitted_input[self.tp_rank].contiguous()
+        
         # Prepare tensors for all-to-all communication
         local_batch_size = input_parallel.size(0)
         chunk_size = self.input_size_per_partition
@@ -254,6 +262,7 @@ def _forward_oproj_tp(
 
         # Only fuse bias add for rank 0 to avoid duplicate bias addition in TP>1
         bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
+        assert self.quant_method is not None
         output_parallel = self.quant_method.apply(self,
                                                   input_parallel,
                                                   bias=bias_)
@@ -267,30 +276,6 @@ def _forward_oproj_tp(
             return output
         return output, output_bias
 
-    # original forward function of RowParallelLinear
-    # def _forward_normal(
-    #     self,
-    #     input_parallel: torch.Tensor,
-    # ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
-
-    #     # Matrix multiply with quantized method
-    #     assert self.quant_method is not None
-    #     # Only fuse bias add for rank 0 to avoid duplicate bias addition in TP>1
-    #     bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
-    #     output_parallel = self.quant_method.apply(self,
-    #                                               input_parallel,
-    #                                               bias=bias_)
-
-    #     if self.reduce_results and self.tp_size > 1:
-    #         output = tensor_model_parallel_all_reduce(output_parallel)
-    #     else:
-    #         output = output_parallel
-    #     # Handle bias return based on configuration
-    #     output_bias = self.bias if self.skip_bias_add else None
-    #     if not self.return_bias:
-    #         return output
-    #     return output, output_bias
-
 
 class AscendMergedColumnParallelLinear(MergedColumnParallelLinear):
     """Packed linear layers with column parallelism.
@@ -349,8 +334,7 @@ def forward(
         if self.forward_type == "mlp_tp":
             return self._forward_mlp_tp(input_)
         else:
-            # same as origin ColumnParallelLinear forward
-            return MergedColumnParallelLinear.forward(self, input_)
+            return super().forward(input_)
 
     def _forward_mlp_tp(
         self,