[Feat] flashcomm2+oshard Generalized (#4723)

### What this PR does / why we need it?
[FlashComm2](https://gitcode.com/ascend-tribe/ascend-inference-cluster/blob/main/FlashComm/FlashComm2%E5%A4%A7%E6%A8%A1%E5%9E%8B%E6%8E%A8%E7%90%86%E4%B8%AD%E4%BB%A5%E5%AD%98%E6%8D%A2%E4%BC%A0%E7%9A%84%E9%80%9A%E4%BF%A1%E4%BC%98%E5%8C%96%E6%8A%80%E6%9C%AF.pdf)
introduces redundant storage of the o_proj matrix, which imposes
pressure on GPU memory. We propose the FlashComm2+Oshard approach by
integrating the shared linear layer feature (#2931). This approach
distributes weights layer-by-layer to each GPU and accesses the o_proj
of each layer via asynchronous broadcast operations, thereby alleviating
memory pressure while achieving nearly lossless performance compared to
the original FlashComm2. This PR implements a generalized
FlashComm2+Oshard solution.

Using following env to support flashcomm2 with oshard

```shell
export VLLM_ASCEND_FLASHCOMM2_PARALLEL_SIZE=1
--additional-config '{
  "layer_sharding": ["o_proj"]
}'
```

### How was this patch tested?

- vLLM version: v0.12.0
- vLLM main:
ad32e3e19c

---------

Signed-off-by: Levi-JQ <yujinqi2@huawei.com>
Co-authored-by: Levi-JQ <yujinqi2@huawei.com>

vllm_ascend/ops/linear_op.py

@@ -21,6 +21,7 @@ CustomLinearOp
 ├── CustomColumnParallelOp
 │   ├── MLPColumnParallelOp
 │   ├── SequenceColumnParallelOp
+│   ├── Flashcomm2OshardQKVParallelOp
 └── CustomRowParallelOp
 │   ├── MLPRowParallelOp
 │   ├── OProjRowParallelOp
@@ -60,6 +61,7 @@ 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_sp, flashcomm2_enable,
                                get_flashcomm2_reorgnized_batch_ids,
                                matmul_allreduce_enable, mlp_tp_enable,
@@ -400,6 +402,9 @@ class Flashcomm2OProjRowParallelOp(CustomRowParallelOp):
         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):
@@ -479,6 +484,39 @@ class SequenceColumnParallelOp(CustomColumnParallelOp):
         return output, output_bias
 
 
+class Flashcomm2OshardQKVParallelOp(CustomColumnParallelOp):
+
+    def __init__(self, layer):
+        super().__init__(layer)
+
+    def apply_impl(
+        self, input_: torch.Tensor
+    ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
+        """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):
@@ -657,12 +695,15 @@ class ShardedCPColumnParallelOp(CustomColumnParallelOp):
 def _get_column_parallel_op(
     prefix, layer
 ) -> Optional[Union[MLPColumnParallelOp, SequenceColumnParallelOp,
-                    ShardedCPColumnParallelOp]]:
+                    ShardedCPColumnParallelOp, Flashcomm2OshardQKVParallelOp]]:
     if enable_dsa_cp() and ("q_b_proj" in prefix or "kv_b_proj" in prefix):
         return ShardedCPColumnParallelOp(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
@@ -719,6 +760,7 @@ def get_parallel_op(disable_tp, prefix, layer, direct):
     custom_op: Optional[Union[MLPColumnParallelOp, SequenceColumnParallelOp,
                               MLPRowParallelOp, OProjRowParallelOp,
                               Flashcomm2OProjRowParallelOp,
+                              Flashcomm2OshardQKVParallelOp,
                               MatmulAllreduceRowParallelOp,
                               SequenceRowParallelOp, ShardedCPRowParallelOp,
                               ShardedCPColumnParallelOp]] = None