[main] flashcomm_v2 optim solution

Author: Levi-JQ

vllm_ascend/ascend_config.py

@@ -92,6 +92,25 @@ def __init__(self, vllm_config):
                 raise AssertionError(
                     "oproj_tensor_parallel_size is only supported in pd scenario and can only be used in D node."
                 )
+        self.flashcomm2_oproj_tensor_parallel_size = additional_config.get(
+            "flashcomm2_oproj_tensor_parallel_size", None)
+        if self.flashcomm2_oproj_tensor_parallel_size is not None:
+            global_tp_size = vllm_config.parallel_config.tensor_parallel_size
+            logger.info(
+                f"Enable Flashcomm2 with flashcomm2_oproj_tensor_parallel_size={self.flashcomm2_oproj_tensor_parallel_size} and global_tp_size={global_tp_size}"
+            )
+            if self.oproj_tensor_parallel_size is not None:
+                raise AssertionError(
+                    "flashcomm2_oproj_tensor_parallel_size cannot be enabled simultaneously with oproj_tensor_parallel_size"
+                )
+            if global_tp_size <= self.flashcomm2_oproj_tensor_parallel_size:
+                raise AssertionError(
+                    "flashcomm2_oproj_tensor_parallel_size ({self.flashcomm2_oproj_tensor_parallel_size}) cannot exceed global tensor parallel size ({global_tp_size})"
+                )
+            if global_tp_size % self.flashcomm2_oproj_tensor_parallel_size != 0:
+                raise AssertionError(
+                    "Global tensor parallel size ({global_tp_size}) must be divisible by flashcomm2_oproj_tensor_parallel_size ({self.flashcomm2_oproj_tensor_parallel_size})"
+                )
 
 
 class TorchairGraphConfig:

vllm_ascend/ascend_forward_context.py

@@ -12,6 +12,7 @@
 
 import vllm_ascend.envs as envs_ascend
 from vllm_ascend.utils import enable_sp
+from vllm_ascend.utils import flashcomm2_enable
 
 
 class FusedMoEState(Enum):
@@ -109,12 +110,18 @@ def set_ascend_forward_context(
         sp_enabled = enable_sp() and \
             tp_world_size > 1 and \
             num_tokens is not None and num_tokens > 1000
+        
+        flashcomm_v2_enabled = flashcomm2_enable() and \
+            tp_world_size > 1 and \
+            num_tokens is not None
 
-        if sp_enabled:
+        if sp_enabled or flashcomm_v2_enabled:
             pad_size = (tp_world_size -
                         (num_tokens % tp_world_size)) % tp_world_size
             forward_context.pad_size = pad_size
+
         forward_context.sp_enabled = sp_enabled
+        forward_context.flashcomm_v2_enabled = flashcomm_v2_enabled
 
         # set this for rope forward_oot using
         forward_context.is_first_layer = True

vllm_ascend/distributed/parallel_state.py

@@ -2,8 +2,11 @@
 
 import torch
 from vllm.config import ParallelConfig
-from vllm.distributed.parallel_state import (GroupCoordinator, get_world_group,
+from vllm.distributed.parallel_state import (GroupCoordinator, get_world_group, get_tp_group, 
                                              init_model_parallel_group)
+import vllm_ascend.envs as envs_ascend
+from vllm.logger import logger
+from vllm_ascend.utils import flashcomm2_enable, oproj_tp_enable
 
 import vllm_ascend.envs as envs_ascend
 from vllm_ascend.ascend_config import get_ascend_config
@@ -13,6 +16,9 @@
 _MLP_TP: Optional[GroupCoordinator] = None
 _OTP: Optional[GroupCoordinator] = None
 _LMTP: Optional[GroupCoordinator] = None
+_FLASHCOMM2_OTP: Optional[GroupCoordinator] = None
+_FLASHCOMM2_ODP: Optional[GroupCoordinator] = None
+
 
 
 def get_mc2_group() -> GroupCoordinator:
@@ -25,12 +31,18 @@ def get_otp_group() -> GroupCoordinator:
         "output tensor parallel group is not initialized")
     return _OTP
 
-
 def get_lmhead_tp_group() -> GroupCoordinator:
     assert _LMTP is not None, (
         "lm head tensor parallel group is not initialized")
     return _LMTP
 
+def get_flashcomm2_otp_group() -> GroupCoordinator:
+    return _FLASHCOMM2_OTP
+
+def get_flashcomm2_odp_group() -> GroupCoordinator:
+    assert _FLASHCOMM2_ODP is not None, (
+        "output data parallel group for flashcomm2 is not initialized")
+    return _FLASHCOMM2_ODP
 
 def get_mlp_tp_group() -> GroupCoordinator:
     assert _MLP_TP is not None, ("mlp group is not initialized")
@@ -110,6 +122,39 @@ def init_ascend_model_parallel(parallel_config: ParallelConfig, ):
                                           get_world_group().local_rank,
                                           backend,
                                           group_name="lmheadtp")
+    
+    if flashcomm2_enable():
+        flashcomm2_otp_size = get_ascend_config().flashcomm2_oproj_tensor_parallel_size
+        global_tp_size = get_tp_group().world_size
+        num_oproj_tensor_parallel_groups: int = (global_tp_size // flashcomm2_otp_size)
+
+        global _FLASHCOMM2_OTP
+        global _FLASHCOMM2_ODP
+
+        _FLASHCOMM2_OTP = None
+        _FLASHCOMM2_ODP = get_tp_group()
+
+        if flashcomm2_otp_size > 1:
+            otp_group_ranks = []
+            odp_group_ranks = [[] for _ in range(flashcomm2_otp_size)]
+            dp_group_index = torch.distributed.get_rank() // global_tp_size
+
+            for i in range(num_oproj_tensor_parallel_groups):
+                ranks = []
+                for j in range(flashcomm2_otp_size):
+                    rank_idx = dp_group_index * global_tp_size + i + j * num_oproj_tensor_parallel_groups
+                    ranks.append(rank_idx)
+                    odp_group_ranks[j].append(rank_idx)
+                otp_group_ranks.append(ranks)
+
+            _FLASHCOMM2_OTP = init_model_parallel_group(otp_group_ranks,
+                                    get_world_group().local_rank,
+                                    backend,
+                                    group_name="flashcomm2_otp")
+            _FLASHCOMM2_ODP = init_model_parallel_group(odp_group_ranks,
+                                        get_world_group().local_rank,
+                                        backend,
+                                        group_name="flashcomm2_odp")
 
 
 def get_mlp_tensor_model_parallel_world_size():
@@ -142,3 +187,13 @@ def destroy_ascend_model_parallel():
     if _OTP:
         _OTP.destroy()
     _OTP = None
+    
+    global _FLASHCOMM2_OTP
+    if _FLASHCOMM2_OTP and get_ascend_config().flashcomm2_oproj_tensor_parallel_size != 1:
+        _FLASHCOMM2_OTP.destroy()  
+        _FLASHCOMM2_OTP = None
+
+    global _FLASHCOMM2_ODP
+    if _FLASHCOMM2_ODP and get_ascend_config().flashcomm2_oproj_tensor_parallel_size != 1:
+        _FLASHCOMM2_ODP.destroy()  
+        _FLASHCOMM2_ODP = None

vllm_ascend/ops/linear.py

@@ -275,6 +275,7 @@ def forward(
         self,
         input_,
         is_prefill: bool = True,
+        is_force_scatter: bool = False,
     ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
         if self.custom_op is not None:
             return self.custom_op.apply(input_)

vllm_ascend/ops/linear_op.py

@@ -46,9 +46,10 @@
 from vllm.distributed import split_tensor_along_last_dim
 from vllm.distributed.parallel_state import get_tp_group
 
-from vllm_ascend.distributed.parallel_state import (get_mlp_tp_group,
+from vllm.forward_context import get_forward_context
+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.utils import (dense_optim_enable, enable_sp,
+from vllm_ascend.utils import (dense_optim_enable, enable_sp, flashcomm2_enable, get_flashcomm2_reorgnized_batch_ids,
                                matmul_allreduce_enable, mlp_tp_enable,
                                oproj_tp_enable)
 
@@ -311,6 +312,104 @@ def update_attrs(self):
         self.input_size_per_partition = self.layer.input_size_per_partition
 
 
+class Flashcomm2OProjRowParallelOp(CustomRowParallelOp):
+
+    def __init__(self, layer):
+        super().__init__(layer)
+        self.forward_type = "flashcomm2_oproj_tp"
+        self.odp_group = get_flashcomm2_odp_group()
+        self.odp_size = self.odp_group.world_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()
+
+    @property
+    def comm_group(self):
+        return get_flashcomm2_otp_group()
+
+    def apply_impl(
+        self,
+        input_: torch.Tensor,
+    ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
+        # Handle input parallelism - split or use as-is
+        if self.input_is_parallel:
+            input_parallel = input_
+        else:
+            tp_rank = self.tp_rank
+            splitted_input = split_tensor_along_last_dim(
+                input_, num_partitions=self.tp_size)
+            input_parallel = splitted_input[tp_rank].contiguous()
+
+        # padding for all-to-all
+        forward_context = get_forward_context()
+        num_padding_tokens = forward_context.pad_size
+        if num_padding_tokens > 0:
+            input_parallel = nn.functional.pad(input_parallel, (0, 0, 0, num_padding_tokens))
+        
+        # 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 = input_parallel.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 = input_parallel.view(chunk_num, batch_size_per_chunk, input_parallel.shape[1])
+        reorganized_chunks = chunked[self.group_indices]
+        send_buf = reorganized_chunks.flatten(1, 2)
+
+        # all-to-all operation parameters
+        all2all_tp_size = self.odp_size
+        local_intermediate_size = input_parallel.size(1)
+        chunk_size = input_parallel.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=input_parallel.dtype,
+            device=input_parallel.device)
+        
+        # Perform all-to-all communication
+        dist.all_to_all_single(recv_buf, send_buf, group=self.odp_group.device_group)
+
+        input_parallel = recv_buf.view(
+            all2all_tp_size,  
+            chunk_size,       
+            -1                
+        ).transpose(0, 1).reshape(chunk_size, -1) 
+
+        # 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,
+                                                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 forward_context.flashcomm1_ds_prefill:
+            # flashcomm1 not enabled
+            output = get_tp_group().all_gather(output, 0)
+            if num_padding_tokens > 0:
+                output = output[:-num_padding_tokens]
+
+        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.input_is_parallel = self.layer.input_is_parallel
+        self.input_size_per_partition = self.layer.input_size_per_partition
+
+
 class MatmulAllreduceRowParallelOp(CustomRowParallelOp):
     _HCOMM_INFO = None
 
@@ -437,17 +536,19 @@ def get_row_parallel_op(
     disable_tp, prefix, layer
 ) -> Tuple[Optional[Union[MLPRowParallelOp, OProjRowParallelOp,
                           MatmulAllreduceRowParallelOp,
-                          SequenceRowParallelOp]], int, int]:
+                          SequenceRowParallelOp, Flashcomm2OProjRowParallelOp]], int, int]:
     if disable_tp:
         return None, 0, 1
 
     custom_op: Optional[Union[MLPRowParallelOp, OProjRowParallelOp,
                               MatmulAllreduceRowParallelOp,
-                              SequenceRowParallelOp]] = None
+                              SequenceRowParallelOp, Flashcomm2OProjRowParallelOp]] = 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 "o_proj" in prefix and flashcomm2_enable():
+        custom_op = Flashcomm2OProjRowParallelOp(layer)
     elif matmul_allreduce_enable():
         custom_op = MatmulAllreduceRowParallelOp(layer)
     elif enable_sp():

vllm_ascend/quantization/quant_config.py

@@ -36,12 +36,14 @@
 from vllm.model_executor.utils import set_weight_attrs
 
 from vllm_ascend.distributed.parallel_state import (get_mlp_tp_group,
-                                                    get_otp_group)
+                                                    get_otp_group,
+                                                    get_flashcomm2_otp_group)
 from vllm_ascend.ops.fused_moe import AscendUnquantizedFusedMoEMethod
-from vllm_ascend.utils import (ASCEND_QUANTIZATION_METHOD, mlp_tp_enable,
+from vllm_ascend.utils import (ASCEND_QUANTIZATION_METHOD, flashcomm2_enable, mlp_tp_enable,
                                oproj_tp_enable)
 
 from .utils import get_quant_method
+from vllm_ascend.ascend_config import get_ascend_config
 
 
 @register_quantization_config(ASCEND_QUANTIZATION_METHOD)
@@ -301,6 +303,11 @@ def apply(
                 tp_rank = get_otp_group().rank_in_group
             elif layer.prefix.find("down_proj") != -1 and mlp_tp_enable():
                 tp_rank = get_mlp_tp_group().rank_in_group
+            elif layer.prefix.find("o_proj") != -1 and flashcomm2_enable():
+                if get_ascend_config().flashcomm2_oproj_tensor_parallel_size == 1:
+                    tp_rank = 0
+                else:
+                    tp_rank = get_flashcomm2_otp_group().rank_in_group
             else:
                 tp_rank = get_tensor_model_parallel_rank()
         else:

vllm_ascend/utils.py

@@ -642,3 +642,25 @@ def npu_stream_switch(target_stream: torch.npu.Stream,
         return nullcontext()
     assert target_stream is not None
     return torch.npu.stream(target_stream)
+
+
+def flashcomm2_enable() -> bool:
+    return get_ascend_config().flashcomm2_oproj_tensor_parallel_size is not None
+
+
+def get_flashcomm2_reorgnized_batch_ids(global_tp_size) -> list[list[int]]:
+    # Reorganize batch_ids so that, after the all2all and reduce-scatter operation, each batch_id corresponds to the rank_id within the DP domain. 
+    # For example, when DP = [0, 1, 2, ..., 15] and flashcomm2_oproj_tensor_parallel_size = 2, 
+    # the reorganized batch_ids will be [[batch0, batch8], [batch1, batch9], ..., [batch7, batch15]].
+    flashcomm2_otp_size = get_ascend_config().flashcomm2_oproj_tensor_parallel_size
+    num_oproj_tensor_parallel_groups: int = (global_tp_size // flashcomm2_otp_size)
+
+    reorgnized_batch_ids = []
+    for i in range(num_oproj_tensor_parallel_groups):
+        ranks = []
+        for j in range(flashcomm2_otp_size):
+            rank_idx = i + j * num_oproj_tensor_parallel_groups
+            ranks.append(rank_idx)
+        reorgnized_batch_ids.append(ranks)
+
+    return reorgnized_batch_ids