init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

vllm_ascend/ops/linear_op.py

@@ -17,44 +17,75 @@ This file extends the functionality of linear operations by encapsulating custom
 communication groups and forward functions into classes (linear ops).
 
 Current class inheritance structure:
-CustomTensorParallelOp
+CustomLinearOp
 ├── CustomColumnParallelOp
 │   ├── MLPColumnParallelOp
-│   ├── DenseOptimMergedColumnParallelOp
-│   └── DenseOptimQKVParallelOp
+│   ├── SequenceColumnParallelOp
+│   ├── Flashcomm2OshardQKVParallelOp
 └── CustomRowParallelOp
-    ├── MLPRowParallelOp
-    ├── OProjRowParallelOp
-    ├── MatmulAllreduceRowParallelOp
-    └── DenseOptimRowParallelOp
-
+│   ├── MLPRowParallelOp
+│   ├── OProjRowParallelOp
+|   ├── Flashcomm2OProjRowParallelOp
+│   ├── MatmulAllreduceRowParallelOp
+│   └── SequenceRowParallelOp
+└── CustomReplicatedOp
 How to extend a new linear op? Taking column parallel op as an example:
 1. Inherit from CustomColumnParallelOp and create a new class MyColumnParallelOp
-2. [Optional] The default communication group is the TP group. If a custom communication group is needed, override the comm_group method
+2. [Optional] The default communication group is the TP group. If a custom communication group is needed,
+   override the comm_group method
 3. Override the apply method according to requirements, which will replace the original linear.forward
-4. Add selection logic for MyColumnParallelOp in the get_column_parallel_op method, typically based on prefix and configuration judgments
-Row parallel op follows a similar approach - inherit from RowColumnParallelOp and register the new class in get_row_parallel_op.
+4. Add selection logic for MyColumnParallelOp in the get_column_parallel_op method, typically based on
+   prefix and configuration judgments
+Row parallel op follows a similar approach - inherit from RowColumnParallelOp and register the new class in
+get_row_parallel_op.
 """
 
-from typing import Optional, Tuple, Union
+from functools import lru_cache
+from types import SimpleNamespace
 
+import regex as re
 import torch
 import torch.distributed as dist
+import torch.nn.functional as F
 import torch_npu
+from torch import nn
 from torch.distributed import ProcessGroup
 from torch.nn.parameter import Parameter
-from vllm.distributed import split_tensor_along_last_dim
+from vllm.distributed import (
+    split_tensor_along_last_dim,
+    tensor_model_parallel_all_reduce,
+    tensor_model_parallel_reduce_scatter,
+)
 from vllm.distributed.parallel_state import get_tp_group
+from vllm.logger import logger
+from vllm.model_executor.models.utils import extract_layer_index
 
-from vllm_ascend.distributed.parallel_state import (get_mlp_tp_group,
-                                                    get_otp_group)
-from vllm_ascend.utils import (dense_optim_enable, enable_sp,
-                               matmul_allreduce_enable, mlp_tp_enable,
-                               oproj_tp_enable)
+from vllm_ascend.ascend_config import get_ascend_config
+from vllm_ascend.ascend_forward_context import _EXTRA_CTX
+from vllm_ascend.device.device_op import DeviceOperator
+from vllm_ascend.distributed.parallel_state import (
+    get_flashcomm2_odp_group,
+    get_flashcomm2_otp_group,
+    get_mlp_tp_group,
+    get_otp_group,
+)
+from vllm_ascend.ops.flashcomm2_oshard_manager import flashcomm2_oshard_manager
+from vllm_ascend.utils import (
+    enable_dsa_cp,
+    enable_dsa_cp_with_layer_shard,
+    enable_sp,
+    flashcomm2_enable,
+    get_flashcomm2_reorgnized_batch_ids,
+    get_weight_prefetch_method,
+    is_vl_model,
+    matmul_allreduce_enable,
+    mlp_tp_enable,
+    oproj_tp_enable,
+    shared_expert_dp_enabled,
+)
 
 
-class CustomTensorParallelOp:
-
+class CustomLinearOp:
     def __init__(self, layer):
         self.layer = layer
         self.bias = None
@@ -96,8 +127,7 @@ class CustomTensorParallelOp:
         return output, output_bias
 
 
-class CustomColumnParallelOp(CustomTensorParallelOp):
-
+class CustomColumnParallelOp(CustomLinearOp):
     def __init__(self, layer):
         super().__init__(layer)
         self.gather_output = None
@@ -107,8 +137,7 @@ class CustomColumnParallelOp(CustomTensorParallelOp):
         self.gather_output = self.layer.gather_output
 
 
-class CustomRowParallelOp(CustomTensorParallelOp):
-
+class CustomRowParallelOp(CustomLinearOp):
     def __init__(self, layer):
         super().__init__(layer)
         self.reduce_results = None
@@ -123,15 +152,35 @@ class CustomRowParallelOp(CustomTensorParallelOp):
 
     def apply(self, input_):
         output, output_bias = self.apply_impl(input_)
-        if dense_optim_enable():
-            torch.ops.vllm.maybe_prefetch_mlp_gate_up_proj(output, self.prefix)
+        weight_prefetch_method = get_weight_prefetch_method()
+        weight_prefetch_method.maybe_prefetch_mlp_weight_preprocess(
+            weight_prefetch_method.MLP_GATE_UP, output, self.prefix
+        )
+
         if not self.return_bias:
             return output
         return output, output_bias
 
+    def get_input_parallel(self, input_: torch.Tensor) -> torch.Tensor:
+        if self.input_is_parallel:
+            return input_
+
+        split_input = split_tensor_along_last_dim(input_, num_partitions=self.tp_size)
+        return split_input[self.tp_rank].contiguous()
+
+
+class CustomReplicatedOp(CustomLinearOp):
+    def apply_impl(self, input_):
+        bias = self.bias if not self.skip_bias_add else None
+        assert self.quant_method is not None
+
+        output = self.quant_method.apply(self.layer, input_, bias)
+        output_bias = self.bias if self.skip_bias_add else None
+
+        return output, output_bias
+
 
 class MLPColumnParallelOp(CustomColumnParallelOp):
-
     def __init__(self, layer):
         super().__init__(layer)
 
@@ -142,7 +191,7 @@ class MLPColumnParallelOp(CustomColumnParallelOp):
     def apply_impl(
         self,
         input_: torch.Tensor,
-    ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
+    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
         bias = self.bias if not self.skip_bias_add else None
         # Matrix multiply.
         assert self.quant_method is not None
@@ -153,71 +202,7 @@ class MLPColumnParallelOp(CustomColumnParallelOp):
         return output, output_bias
 
 
-class SequenceMergedColumnParallelOp(CustomColumnParallelOp):
-
-    def apply_impl(
-        self, input_: torch.Tensor
-    ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
-        """Linear layer with column parallelism.
-
-        Implemented multiple optimization projects for dense models, such as FlashComm and
-        communication-computation fusion.
-        """
-
-        bias = self.bias if not self.skip_bias_add else None
-
-        # Matrix multiply.
-        assert self.quant_method is not None
-
-        input_ = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(input_, True)
-        output_parallel = self.quant_method.apply(self.layer, input_, bias)
-
-        if self.gather_output:
-            # All-gather across the partitions.
-            output = self.comm_group.all_gather(output_parallel)
-        else:
-            output = output_parallel
-        output_bias = self.bias if self.skip_bias_add else None
-        return output, output_bias
-
-
-class SequenceQKVParallelOp(CustomColumnParallelOp):
-
-    def __init__(self, layer, prefix):
-        super().__init__(layer)
-        self.prefix = prefix
-
-    def apply_impl(
-        self, input_: torch.Tensor
-    ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
-        """Linear layer with column parallelism.
-
-        Implemented multiple optimization projects for dense models, such as FlashComm and
-        communication-computation fusion.
-        """
-
-        bias = self.bias if not self.skip_bias_add else None
-
-        # Matrix multiply.
-        assert self.quant_method is not None
-
-        layer_num = self.prefix.split('.')[2]
-
-        input_ = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
-            input_, layer_num != '0')
-        output_parallel = self.quant_method.apply(self.layer, input_, bias)
-
-        if self.gather_output:
-            # All-gather across the partitions.
-            output = self.comm_group.all_gather(output_parallel)
-        else:
-            output = output_parallel
-        output_bias = self.bias if self.skip_bias_add else None
-        return output, output_bias
-
-
 class MLPRowParallelOp(CustomRowParallelOp):
-
     def __init__(self, layer):
         super().__init__(layer)
 
@@ -225,30 +210,52 @@ class MLPRowParallelOp(CustomRowParallelOp):
     def comm_group(self):
         return get_mlp_tp_group()
 
-    def apply_impl(
-        self, 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()
+    def apply_impl(self, input_: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        input_parallel = self.get_input_parallel(input_)
 
         assert self.quant_method is not None
-        bias_ = None if (self.tp_rank > 0
-                         or self.skip_bias_add) else self.layer.bias
-        output_parallel = self.quant_method.apply(self.layer,
-                                                  input_parallel,
-                                                  bias=bias_)
+        bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.layer.bias
+        output_parallel = self.quant_method.apply(self.layer, input_parallel, bias=bias_)
         output = self.comm_group.reduce_scatter(output_parallel, 0)
 
         output_bias = self.bias if self.skip_bias_add else None
         return output, output_bias
 
 
-class OProjRowParallelOp(CustomRowParallelOp):
+class DSV4OProjColumnParallelOp(CustomColumnParallelOp):
+    @property
+    def comm_group(self):
+        return get_otp_group()
 
+    def apply_impl(
+        self,
+        input_: torch.Tensor,
+    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        bias = self.bias if not self.skip_bias_add else None
+        assert self.quant_method is not None
+        output_parallel = self.quant_method.apply(self.layer, input_, bias)
+        output_bias = self.bias if self.skip_bias_add else None
+        return output_parallel, output_bias
+
+
+class DSV4OProjRowParallelOp(CustomRowParallelOp):
+    @property
+    def comm_group(self):
+        return get_otp_group()
+
+    def apply_impl(
+        self,
+        input_: torch.Tensor,
+    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        input_parallel = self.get_input_parallel(input_)
+        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.layer, input_parallel, bias=bias_)
+        output_bias = self.bias if self.skip_bias_add else None
+        return output_parallel, output_bias
+
+
+class OProjRowParallelOp(CustomRowParallelOp):
     def __init__(self, layer):
         super().__init__(layer)
 
@@ -259,14 +266,8 @@ class OProjRowParallelOp(CustomRowParallelOp):
     def apply_impl(
         self,
         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()
+    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        input_parallel = self.get_input_parallel(input_)
 
         # Prepare tensors for all-to-all communication
         local_batch_size = input_parallel.size(0)
@@ -275,27 +276,19 @@ class OProjRowParallelOp(CustomRowParallelOp):
 
         # Reshape tensor for efficient cross-device transfer:
         # [batch, dim] -> [tp_size, batch, chunk] -> flattened
-        send_buf = (input_parallel.reshape(-1,
-                                           self.tp_size, chunk_size).transpose(
-                                               0, 1).contiguous().view(-1))
+        send_buf = input_parallel.reshape(-1, self.tp_size, chunk_size).transpose(0, 1).contiguous().view(-1)
 
         # Create receive buffer
-        recv_buf = torch.empty(total_batch_size * chunk_size,
-                               dtype=input_parallel.dtype,
-                               device=input_parallel.device)
+        recv_buf = torch.empty(total_batch_size * chunk_size, dtype=input_parallel.dtype, device=input_parallel.device)
 
         # Perform all-to-all communication
-        dist.all_to_all_single(recv_buf,
-                               send_buf,
-                               group=self.comm_group.device_group)
+        dist.all_to_all_single(recv_buf, send_buf, group=self.comm_group.device_group)
         input_parallel = recv_buf.view(total_batch_size, chunk_size)
 
         # 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.layer,
-                                                  input_parallel,
-                                                  bias=bias_)
+        output_parallel = self.quant_method.apply(self.layer, input_parallel, bias=bias_)
 
         # otp-specific: Combine partial results across devices
         output = self.comm_group.reduce_scatter(output_parallel, dim=0)
@@ -311,6 +304,119 @@ class OProjRowParallelOp(CustomRowParallelOp):
         self.input_size_per_partition = self.layer.input_size_per_partition
 
 
+class Flashcomm2OProjRowParallelOp(CustomRowParallelOp):
+    def __init__(self, layer):
+        super().__init__(layer)
+        self.odp_group = get_flashcomm2_odp_group()
+        self.odp_size = self.odp_group.world_size
+        self.otp_size = get_ascend_config().flashcomm2_oproj_tensor_parallel_size
+        self.reorgnized_batch_ids = get_flashcomm2_reorgnized_batch_ids(get_tp_group().world_size)
+        self.group_indices = torch.tensor(self.reorgnized_batch_ids).npu()
+        self.layer._quant_comm_config = {}
+
+    @property
+    def comm_group(self):
+        return get_flashcomm2_otp_group()
+
+    @property
+    def tp_rank(self):
+        if get_ascend_config().flashcomm2_oproj_tensor_parallel_size == 1:
+            return 0
+        return self.comm_group.rank_in_group
+
+    @property
+    def tp_size(self):
+        if get_ascend_config().flashcomm2_oproj_tensor_parallel_size == 1:
+            return 1
+        return self.comm_group.world_size
+
+    def apply_impl(
+        self,
+        input_: torch.Tensor,
+    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        """Linear layer for Flashcomm2.
+        Input.ahspe = [batchsize*seqlength, headnum*headdim/TP]
+        Output.shape = [(batchsize*seqlength+padsize)/TP, hiddensize]
+        """
+        # Handle input parallelism - split or use as-is
+        input_parallel = self.get_input_parallel(input_)
+
+        # padding for all-to-all
+        num_padding_tokens = _EXTRA_CTX.pad_size
+        if num_padding_tokens > 0:
+            input_parallel = nn.functional.pad(input_parallel, (0, 0, 0, num_padding_tokens))
+
+        def otp_maybe_quant_comm(x):
+            # Reorganize the tensor so that the batch id and rank id correspond to each other.
+            chunk_num = len(self.reorgnized_batch_ids) * len(self.reorgnized_batch_ids[0])
+            batch_size = x.size(0)
+
+            assert batch_size % chunk_num == 0, f"Batch_size({batch_size}) must be divisible by chunk_num({chunk_num})"
+
+            batch_size_per_chunk = batch_size // chunk_num
+            # Indices of reorganized tensor
+            chunked = x.view(chunk_num, batch_size_per_chunk, x.shape[1])
+            if self.otp_size != 1:
+                chunked = chunked[self.group_indices]
+            send_buf = chunked.flatten(1, 2)
+
+            # all-to-all operation parameters
+            all2all_tp_size = self.odp_size
+            local_intermediate_size = x.size(1)
+            chunk_size = x.size(0) // all2all_tp_size
+            total_intermediate_size = local_intermediate_size * all2all_tp_size
+
+            # Create receive buffer
+            recv_buf = torch.empty(total_intermediate_size * chunk_size, dtype=x.dtype, device=x.device)
+
+            # Perform all-to-all communication
+            dist.all_to_all_single(recv_buf, send_buf, group=self.odp_group.device_group)
+
+            return recv_buf.view(all2all_tp_size, chunk_size, -1).transpose(0, 1).reshape(chunk_size, -1)
+
+        if not hasattr(self, "_quant_comm_config"):
+            self.layer._quant_comm_config = {}
+        self.layer._quant_comm_config["communication_fn"] = otp_maybe_quant_comm
+        actual_quant_method = getattr(self.quant_method, "quant_method", self.quant_method)
+        from vllm_ascend.quantization.methods.w8a8_static import AscendW8A8LinearMethod
+
+        if not isinstance(actual_quant_method, AscendW8A8LinearMethod):
+            # Check if w8a8 quantization is enabled. If not, communicate immediately.
+            input_parallel = otp_maybe_quant_comm(input_parallel)
+
+        # Matrix multiply.
+        assert self.quant_method is not None
+        # Only fuse bias add into GEMM for rank 0 (this ensures that
+        # bias will not get added more than once in TP>1 case)
+        bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
+
+        output_parallel = self.quant_method.apply(self.layer, input_parallel, bias=bias_)
+        # output_parallel shape: [bs/(TP/flashcomm2_otp_size), hiddenstate]
+        if self.tp_size > 1:
+            # flashcomm2 with reduce-scatter
+            output = self.comm_group.reduce_scatter(output_parallel, dim=0)
+        else:
+            output = output_parallel
+
+        if not _EXTRA_CTX.flash_comm_v1_enabled:
+            # flashcomm1 not enabled
+            output = get_tp_group().all_gather(output, 0)
+            if num_padding_tokens > 0:
+                output = output[:-num_padding_tokens]
+
+        # Handle bias return based on configuration
+        output_bias = self.bias if self.skip_bias_add else None
+
+        return output, output_bias
+
+    def update_attrs(self):
+        super().update_attrs()
+        self.input_is_parallel = self.layer.input_is_parallel
+        self.input_size_per_partition = self.layer.input_size_per_partition
+        if flashcomm2_oshard_manager.flashcomm2_oshard_enable():
+            flashcomm2_oshard_manager.register_layer(self.layer, prefetch_step=1)
+
+
 class MatmulAllreduceRowParallelOp(CustomRowParallelOp):
     _HCOMM_INFO = None
 
@@ -318,28 +424,18 @@ class MatmulAllreduceRowParallelOp(CustomRowParallelOp):
         super().__init__(layer)
         self.hcomm_info = self.get_hcomm_info(self.comm_group.device_group)
 
-    def apply_impl(
-        self, 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()
+    def apply_impl(self, input_: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        input_parallel = self.get_input_parallel(input_)
         """Calculate the output tensor of forward by considering
         fusing communication and computation."""
         bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
         if self.reduce_results and self.tp_size > 1:
-            output = torch_npu.npu_mm_all_reduce_base(input_parallel,
-                                                      self.weight_t,
-                                                      self.hcomm_info,
-                                                      bias=bias_)
+            output = torch_npu.npu_mm_all_reduce_base(
+                input_parallel, self.layer.weight.t(), self.hcomm_info, bias=bias_
+            )
         else:
             assert self.quant_method is not None
-            output = self.quant_method.apply(self.layer,
-                                             input_parallel,
-                                             bias=bias_)
+            output = self.quant_method.apply(self.layer, input_parallel, bias=bias_)
 
         output_bias = self.bias if self.skip_bias_add else None
         return output, output_bias
@@ -353,107 +449,363 @@ class MatmulAllreduceRowParallelOp(CustomRowParallelOp):
         rank = torch.distributed.get_rank(group)
         if torch.__version__ > "2.0":
             global_rank = torch.distributed.get_global_rank(group, rank)
-            cls._HCOMM_INFO = group._get_backend(
-                torch.device("npu")).get_hccl_comm_name(global_rank)
+            cls._HCOMM_INFO = group._get_backend(torch.device("npu")).get_hccl_comm_name(global_rank)
         else:
             cls._HCOMM_INFO = group.get_hccl_comm_name(rank)
         return cls._HCOMM_INFO
 
-    def update_attrs(self):
-        super().update_attrs()
-        self.weight_t = self.layer.weight.t()
 
-
-class SequenceRowParallelOp(CustomRowParallelOp):
-
-    def __init__(self, layer, prefix):
-        super().__init__(layer)
-        self.prefix = prefix
-
-    def apply_impl(
-        self, input_: torch.Tensor
-    ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
+class SequenceColumnParallelOp(CustomColumnParallelOp):
+    def apply_impl(self, input_: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
         """Linear layer with column parallelism.
 
         Implemented multiple optimization projects for dense models, such as FlashComm and
         communication-computation fusion.
         """
 
-        if self.input_is_parallel:
-            input_parallel = input_
+        bias = self.bias if not self.skip_bias_add else None
+
+        # Matrix multiply.
+        assert self.quant_method is not None
+        need_all_gather = not (extract_layer_index(self.layer.prefix) == 0 and is_vl_model() and "attn" in self.prefix)
+        input_ = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(input_, label=need_all_gather)
+        output_parallel = self.quant_method.apply(self.layer, input_, bias)
+
+        if self.gather_output:
+            # All-gather across the partitions.
+            output = self.comm_group.all_gather(output_parallel)
         else:
-            splitted_input = split_tensor_along_last_dim(
-                input_, num_partitions=self.tp_size)
-            input_parallel = splitted_input[self.tp_rank].contiguous()
+            output = output_parallel
+        output_bias = self.bias if self.skip_bias_add else None
+        return output, output_bias
+
+
+class Flashcomm2OshardQKVParallelOp(CustomColumnParallelOp):
+    def __init__(self, layer):
+        super().__init__(layer)
+
+    def apply_impl(self, input_: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        """Column-parallel linear with FlashComm2 OShard optimization."""
+
+        bias = self.bias if not self.skip_bias_add else None
+
+        # Matrix multiply.
+        assert self.quant_method is not None
+
+        if enable_sp():
+            input_ = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(input_, True)
+
+        # Trigger async broadcast before matmul to overlap communication.
+        flashcomm2_oshard_manager.trigger_broadcast_for_layer(self.layer.prefix)
+
+        output_parallel = self.quant_method.apply(self.layer, input_, bias)
+        if self.gather_output and self.tp_size > 1:
+            # All-gather across the partitions.
+            output = self.comm_group.all_gather(output_parallel)
+        else:
+            output = output_parallel
+        output_bias = self.bias if self.skip_bias_add else None
+        return output, output_bias
+
+
+class SequenceRowParallelOp(CustomRowParallelOp):
+    def __init__(self, layer):
+        super().__init__(layer)
+        self.unique_prefix = None
+
+    def apply_impl(self, input_: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        """Linear layer with column parallelism.
+
+        Implemented multiple optimization projects for dense models, such as FlashComm and
+        communication-computation fusion.
+        """
+        input_parallel = self.get_input_parallel(input_)
 
         assert self.quant_method is not None
         bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
 
         if self.tp_size == 1 or not self.reduce_results:
-            output = self.quant_method.apply(self.layer,
-                                             input_parallel,
-                                             bias=bias_)
+            output = self.quant_method.apply(self.layer, input_parallel, bias=bias_)
         else:
-            output_parallel = self.quant_method.apply(self.layer,
-                                                      input_parallel,
-                                                      bias=bias_)
-            output = torch.ops.vllm.maybe_pad_and_reduce(output_parallel)
+            output = torch.ops.vllm.matmul_and_reduce(input_parallel, self.unique_prefix)
 
         output_bias = self.bias if self.skip_bias_add else None
         return output, output_bias
 
+    def matmul_and_reduce(self, input_parallel: torch.Tensor, bias_: Parameter | None) -> torch.Tensor:
+        assert self.quant_method is not None
+        try:
+            flash_comm_v1_enabled = _EXTRA_CTX.flash_comm_v1_enabled
+            mmrs_fusion = _EXTRA_CTX.mmrs_fusion
+        except AssertionError:
+            flash_comm_v1_enabled = False
+            mmrs_fusion = False
+            logger.debug(
+                "matmul_and_reduce: _EXTRA_CTX access failed (profile_run?), "
+                "using defaults: flash_comm_v1=False, mmrs_fusion=False",
+            )
+
+        x = input_parallel
+
+        if not flash_comm_v1_enabled:
+            output_parallel = self.layer.quant_method.apply(self.layer, x, bias=bias_)
+            return tensor_model_parallel_all_reduce(output_parallel)
+
+        pad_size = _EXTRA_CTX.pad_size
+        dsa_cp_attn_out = enable_dsa_cp() and ("o_proj" in self.layer.prefix or "wo_b" in self.layer.prefix)
+        if pad_size > 0 and not dsa_cp_attn_out:
+            x = F.pad(x, (0, 0, 0, pad_size))
+
+        world_size = self.layer.tp_size
+        hcom_name = get_tp_group().device_group._get_backend(torch.device("npu")).get_hccl_comm_name(self.layer.tp_rank)
+
+        from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+        from vllm_ascend.quantization.method_adapters import AscendLinearMethod
+        from vllm_ascend.quantization.methods import AscendW8A8LinearMethod
+
+        # For unquant
+        if mmrs_fusion and isinstance(self.layer.quant_method, UnquantizedLinearMethod):
+            output = DeviceOperator.npu_mm_reduce_scatter_base(
+                x,
+                self.layer.weight.t(),
+                hcom_name,
+                world_size,
+                reduce_op="sum",
+                bias=None,
+                comm_turn=0,
+            )
+            if bias_ is not None:
+                output.add_(bias_)
+        # For w8a8 quant
+        elif mmrs_fusion and (
+            isinstance(self.layer.quant_method, AscendLinearMethod)
+            and isinstance(self.layer.quant_method.quant_method, AscendW8A8LinearMethod)
+        ):
+            if x.dtype != torch.int8:
+                x_quant = torch.ops.vllm.quantize(
+                    x,
+                    self.layer.aclnn_input_scale,
+                    self.layer.aclnn_input_scale_reciprocal,
+                    self.layer.aclnn_input_offset,
+                )
+            else:
+                x_quant = x
+            quant_bias = self.layer.quant_bias
+            deq_scale = self.layer.deq_scale
+            output_dtype = torch.bfloat16
+            output = DeviceOperator.npu_mm_reduce_scatter_base(
+                x_quant,
+                self.layer.weight,
+                hcom_name,
+                world_size,
+                reduce_op="sum",
+                bias=None,
+                comm_turn=0,
+                x2_scale=deq_scale,
+                output_dtype=output_dtype,
+            )
+            output = torch.add(output, torch.mul(quant_bias, deq_scale).to(self.layer.params_dtype))
+        else:
+            output_parallel = self.layer.quant_method.apply(self.layer, x, bias=bias_)
+            output = tensor_model_parallel_reduce_scatter(output_parallel, 0)
+
+        return output
+
     def update_attrs(self):
         super().update_attrs()
         self.input_is_parallel = self.layer.input_is_parallel
         self.reduce_results = self.layer.reduce_results
+        self.unique_prefix = self.layer.unique_prefix
 
 
-def get_column_parallel_op(
-    disable_tp, prefix, layer
-) -> Tuple[Optional[Union[MLPColumnParallelOp, SequenceMergedColumnParallelOp,
-                          SequenceQKVParallelOp]], int, int]:
-    if disable_tp:
+class ShardedCPRowParallelOp(CustomRowParallelOp):
+    @property
+    def comm_group(self):
+        # fake comm group to bypass tp logic
+        return SimpleNamespace(world_size=1, rank_in_group=0, device_group=None)
+
+    def apply_impl(
+        self,
+        input_,
+    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
+        assert self.quant_method is not None
+        output = self.quant_method.apply(self.layer, input_, bias_)
+        output_bias = self.bias if self.skip_bias_add else None
+        if not self.return_bias:
+            return output
+        return output, output_bias
+
+    def update_attrs(self):
+        super().update_attrs()
+        self.layer.reduce_results = False
+
+
+class ShardedCPColumnParallelOp(CustomColumnParallelOp):
+    @property
+    def comm_group(self):
+        # fake comm group to bypass tp logic
+        return SimpleNamespace(world_size=1, rank_in_group=0, device_group=None)
+
+    def apply_impl(
+        self,
+        input_,
+    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        bias = self.bias if not self.skip_bias_add else None
+        assert self.quant_method is not None
+        output = self.quant_method.apply(self.layer, input_, bias)
+        output_bias = self.bias if self.skip_bias_add else None
+        if not self.return_bias:
+            return output
+        return output, output_bias
+
+
+def _get_column_parallel_op(
+    prefix, layer
+) -> (
+    MLPColumnParallelOp
+    | DSV4OProjColumnParallelOp
+    | SequenceColumnParallelOp
+    | ShardedCPColumnParallelOp
+    | Flashcomm2OshardQKVParallelOp
+    | None
+):
+    if enable_dsa_cp() and ("q_b_proj" in prefix or "kv_b_proj" in prefix):
+        return ShardedCPColumnParallelOp(layer)
+    if "wo_a" in prefix and oproj_tp_enable():
+        return DSV4OProjColumnParallelOp(layer)
+    if "gate_up_proj" in prefix and mlp_tp_enable() and not is_moe_layer(prefix):
+        return MLPColumnParallelOp(layer)
+    if flashcomm2_oshard_manager.flashcomm2_oshard_enable():
+        if any(p in prefix for p in ("qkv_proj", "conv1d", "query_key_value")):
+            return Flashcomm2OshardQKVParallelOp(layer)
+    if enable_sp():
+        if "shared_expert" in prefix:
+            return None
+        sp_column_prefix = [
+            "gate_up_proj",  # first MLP of most LLMs
+            "in_proj",  # gated deltanet of Qwen3 Next
+            "qkv_proj",  # qkv linear of most LLMs
+            "conv1d",  # gated deltanet of Qwen3 Next
+            "query_key_value",  # qkv linear of Bailing
+        ]
+        for a_prefix in sp_column_prefix:
+            if a_prefix in prefix:
+                return SequenceColumnParallelOp(layer)
+
+    return None
+
+
+def _get_row_parallel_op(
+    prefix, layer
+) -> (
+    MLPRowParallelOp
+    | OProjRowParallelOp
+    | DSV4OProjRowParallelOp
+    | Flashcomm2OProjRowParallelOp
+    | MatmulAllreduceRowParallelOp
+    | SequenceRowParallelOp
+    | ShardedCPRowParallelOp
+    | None
+):
+    if "wo_b" in prefix and oproj_tp_enable():
+        return DSV4OProjRowParallelOp(layer)
+    if enable_dsa_cp_with_layer_shard() and "o_proj" in prefix:
+        return ShardedCPRowParallelOp(layer)
+    if "down_proj" in prefix and mlp_tp_enable() and not is_moe_layer(prefix):
+        return MLPRowParallelOp(layer)
+    if "o_proj" in prefix and oproj_tp_enable():
+        return OProjRowParallelOp(layer)
+    if matmul_allreduce_enable():
+        return MatmulAllreduceRowParallelOp(layer)
+    if flashcomm2_enable():
+        if "o_proj" in prefix or "out_proj" in prefix:
+            return Flashcomm2OProjRowParallelOp(layer)
+    if enable_sp():
+        if "shared_expert" in prefix:
+            return None
+        sp_row_prefixes = [
+            "o_proj",  # attn output linear of most LLMs
+            "out_proj",  # attn output linear of Qwen3 Next
+            "down_proj",  # second MLP of most LLMs
+            "attention.dense",  # attn output linear of Bailing
+            "wo_b",  # attn output linear of v4
+        ]
+        for a_prefix in sp_row_prefixes:
+            if a_prefix in prefix:
+                return SequenceRowParallelOp(layer)
+
+    return None
+
+
+def get_parallel_op(disable_tp, prefix, layer, direct):
+    if (
+        disable_tp
+        or ("shared_experts" in prefix and shared_expert_dp_enabled())
+        or ("shared_expert" in prefix and shared_expert_dp_enabled())
+    ):
         return None, 0, 1
+    custom_op: (
+        MLPColumnParallelOp
+        | DSV4OProjColumnParallelOp
+        | SequenceColumnParallelOp
+        | MLPRowParallelOp
+        | OProjRowParallelOp
+        | DSV4OProjRowParallelOp
+        | Flashcomm2OProjRowParallelOp
+        | Flashcomm2OshardQKVParallelOp
+        | MatmulAllreduceRowParallelOp
+        | SequenceRowParallelOp
+        | ShardedCPRowParallelOp
+        | ShardedCPColumnParallelOp
+        | None
+    ) = None
+    if direct == "row":
+        custom_op = _get_row_parallel_op(prefix, layer)
 
-    custom_op: Optional[Union[
-        MLPColumnParallelOp,
-        SequenceMergedColumnParallelOp,
-        SequenceQKVParallelOp,
-    ]] = None
-    if "gate_up_proj" in prefix and mlp_tp_enable():
-        custom_op = MLPColumnParallelOp(layer)
-    elif "gate_up_proj" in prefix and enable_sp():
-        custom_op = SequenceMergedColumnParallelOp(layer)
-    elif enable_sp():
-        custom_op = SequenceQKVParallelOp(layer, prefix)
+    if direct == "column":
+        custom_op = _get_column_parallel_op(prefix, layer)
 
     if custom_op is not None:
+        logger.debug(
+            "get_parallel_op: prefix=%s, direct=%s -> %s (tp_rank=%d, tp_size=%d)",
+            prefix,
+            direct,
+            type(custom_op).__name__,
+            custom_op.tp_rank,
+            custom_op.tp_size,
+        )
         return custom_op, custom_op.tp_rank, custom_op.tp_size
 
     return None, get_tp_group().rank_in_group, get_tp_group().world_size
 
 
-def get_row_parallel_op(
-    disable_tp, prefix, layer
-) -> Tuple[Optional[Union[MLPRowParallelOp, OProjRowParallelOp,
-                          MatmulAllreduceRowParallelOp,
-                          SequenceRowParallelOp]], int, int]:
+def get_replicated_op(disable_tp, prefix, layer) -> tuple[CustomReplicatedOp | None, int | None, int | None]:
     if disable_tp:
-        return None, 0, 1
+        return None, None, None
 
-    custom_op: Optional[Union[MLPRowParallelOp, OProjRowParallelOp,
-                              MatmulAllreduceRowParallelOp,
-                              SequenceRowParallelOp]] = None
-    if "down_proj" in prefix and mlp_tp_enable():
-        custom_op = MLPRowParallelOp(layer)
-    elif "o_proj" in prefix and oproj_tp_enable():
-        custom_op = OProjRowParallelOp(layer)
-    elif matmul_allreduce_enable():
-        custom_op = MatmulAllreduceRowParallelOp(layer)
-    elif enable_sp():
-        custom_op = SequenceRowParallelOp(layer, prefix)
+    custom_op = CustomReplicatedOp(layer)
+    return custom_op, custom_op.tp_rank, custom_op.tp_size
 
-    if custom_op is not None:
-        return custom_op, custom_op.tp_rank, custom_op.tp_size
 
-    return None, get_tp_group().rank_in_group, get_tp_group().world_size
+def is_moe_layer(prefix: str) -> bool:
+    @lru_cache(maxsize=1)
+    def get_moe_params():
+        from vllm.config import get_current_vllm_config
+
+        vllm_config = get_current_vllm_config()
+        config = vllm_config.model_config.hf_text_config
+        n_routed_experts = getattr(config, "n_routed_experts", 0)
+        first_k_dense_replace = getattr(config, "first_k_dense_replace", float("inf"))
+        moe_layer_freq = getattr(config, "moe_layer_freq", 1)
+        return n_routed_experts, first_k_dense_replace, moe_layer_freq
+
+    match = re.search(r"layers\.(\d+)\.", prefix)
+    if match is None:
+        return False
+    layer_idx = int(match.group(1))
+
+    n_routed_experts, first_k_dense_replace, moe_layer_freq = get_moe_params()
+
+    return n_routed_experts is not None and layer_idx >= first_k_dense_replace and layer_idx % moe_layer_freq == 0