[scheduler] support cp

Signed-off-by: LookAround <[email protected]>

Author: LookAround0301

vllm/config/parallel.py

@@ -67,6 +67,8 @@ class ParallelConfig:
     """Number of pipeline parallel groups."""
     tensor_parallel_size: int = 1
     """Number of tensor parallel groups."""
+    context_parallel_size: int = 1
+    """Number of context parallel groups."""
     data_parallel_size: int = 1
     """Number of data parallel groups. MoE layers will be sharded according to
     the product of the tensor parallel size and data parallel size."""
@@ -349,7 +351,7 @@ def __post_init__(self) -> None:
 
         # Continue with the rest of the initialization
         self.world_size = self.pipeline_parallel_size * \
-            self.tensor_parallel_size
+            self.tensor_parallel_size * self.context_parallel_size
 
         if self.data_parallel_size_local > self.data_parallel_size:
             raise ValueError(

vllm/distributed/parallel_state.py

@@ -982,6 +982,24 @@ def get_pp_group() -> GroupCoordinator:
     return _PP
 
 
+_CP: Optional[GroupCoordinator] = None
+
+
+def get_cp_group() -> GroupCoordinator:
+    assert _CP is not None, ("context parallel group is not initialized")
+    return _CP
+
+
+def get_context_model_parallel_world_size():
+    """Return world size for the tensor model parallel group."""
+    return get_cp_group().world_size
+
+
+def get_context_model_parallel_rank():
+    """Return my rank for the tensor model parallel group."""
+    return get_cp_group().rank_in_group
+
+
 @deprecated("`get_pipeline_model_parallel_group` has been replaced with "
             "`get_pp_group` and may be removed in v0.12. Please use "
             "`get_pp_group` instead.")
@@ -1088,6 +1106,7 @@ def init_distributed_environment(world_size: int = -1,
 def initialize_model_parallel(
     tensor_model_parallel_size: int = 1,
     pipeline_model_parallel_size: int = 1,
+    context_model_parallel_size: int = 1,
     decode_context_model_parallel_size: Optional[int] = 1,
     backend: Optional[str] = None,
 ) -> None:
@@ -1138,7 +1157,7 @@ def initialize_model_parallel(
     # last dimension, then reshape to 2D, then unbind the last dimension
     all_ranks = torch.arange(world_size).reshape(
         -1, data_parallel_size, pipeline_model_parallel_size,
-        tensor_model_parallel_size)  # noqa
+        context_model_parallel_size, tensor_model_parallel_size)  # noqa
 
     # Build the tensor model-parallel groups.
     global _TP
@@ -1174,7 +1193,7 @@ def initialize_model_parallel(
     global _PP
     assert _PP is None, (
         "pipeline model parallel group is already initialized")
-    group_ranks = all_ranks.transpose(2, 3).reshape(
+    group_ranks = all_ranks.transpose(2, 4).reshape(
         -1, pipeline_model_parallel_size).unbind(0)
     group_ranks = [x.tolist() for x in group_ranks]
     _PP = init_model_parallel_group(group_ranks,
@@ -1185,7 +1204,7 @@ def initialize_model_parallel(
     global _DP
     assert _DP is None, ("data parallel group is already initialized")
     group_ranks = all_ranks.transpose(1,
-                                      3).reshape(-1,
+                                      4).reshape(-1,
                                                  data_parallel_size).unbind(0)
     group_ranks = [x.tolist() for x in group_ranks]
     _DP = init_model_parallel_group(group_ranks,
@@ -1196,23 +1215,34 @@ def initialize_model_parallel(
     global _EP
     assert _EP is None, ("expert parallel group is already initialized")
     group_ranks = all_ranks.transpose(1, 2).reshape(
-        -1, data_parallel_size * tensor_model_parallel_size).unbind(0)
+        -1, data_parallel_size * tensor_model_parallel_size * context_model_parallel_size).unbind(0)
     group_ranks = [x.tolist() for x in group_ranks]
     _EP = init_model_parallel_group(group_ranks,
                                     get_world_group().local_rank,
                                     backend,
                                     group_name="ep")
 
+    global _CP
+    assert _CP is None, ("context parallel group is already initialized")
+    group_ranks = all_ranks.transpose(3, 4).reshape(
+        -1, context_model_parallel_size).unbind(0)
+    group_ranks = [x.tolist() for x in group_ranks]
+    _CP = init_model_parallel_group(group_ranks,
+                                    get_world_group().local_rank,
+                                    backend,
+                                    group_name="cp")
+
     logger.info(
         "rank %s in world size %s is assigned as "
-        "DP rank %s, PP rank %s, TP rank %s, EP rank %s", rank, world_size,
+        "DP rank %s, PP rank %s, TP rank %s, EP rank %s, CP rank %s", rank, world_size,
         _DP.rank_in_group, _PP.rank_in_group, _TP.rank_in_group,
-        _EP.rank_in_group)
+        _EP.rank_in_group, _CP.rank_in_group)
 
 
 def ensure_model_parallel_initialized(
     tensor_model_parallel_size: int,
     pipeline_model_parallel_size: int,
+    context_model_parallel_size: int,
     decode_context_model_parallel_size: Optional[int] = 1,
     backend: Optional[str] = None,
 ) -> None:
@@ -1225,6 +1255,7 @@ def ensure_model_parallel_initialized(
     if not model_parallel_is_initialized():
         initialize_model_parallel(tensor_model_parallel_size,
                                   pipeline_model_parallel_size,
+                                  context_model_parallel_size,
                                   decode_context_model_parallel_size, backend)
         return
 
@@ -1238,6 +1269,11 @@ def ensure_model_parallel_initialized(
         "pipeline parallel group already initialized, but of unexpected size. "
         f"got: {pp_world_size=} vs. "
         f"wanted: {pipeline_model_parallel_size=}")
+    cp_world_size = get_cp_group().world_size
+    assert (cp_world_size == context_model_parallel_size), (
+        "context parallel group already initialized, but of unexpected size: "
+        f"{cp_world_size=} vs. "
+        f"{context_model_parallel_size=}")
 
 
 def prepare_communication_buffer_for_model(model: torch.nn.Module):
@@ -1345,6 +1381,11 @@ def destroy_model_parallel():
         _EP.destroy()
     _EP = None
 
+    global _CP
+    if _CP:
+        _CP.destroy()
+    _CP = None
+
 
 def destroy_distributed_environment():
     global _WORLD, _NODE_COUNT

vllm/engine/arg_utils.py

@@ -318,6 +318,7 @@ class EngineArgs:
     tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
     decode_context_parallel_size: int = \
         ParallelConfig.decode_context_parallel_size
+    context_parallel_size: int = ParallelConfig.context_parallel_size
     data_parallel_size: int = ParallelConfig.data_parallel_size
     data_parallel_rank: Optional[int] = None
     data_parallel_start_rank: Optional[int] = None
@@ -653,6 +654,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         parallel_group.add_argument(
             "--decode-context-parallel-size", "-dcp",
             **parallel_kwargs["decode_context_parallel_size"])
+        parallel_group.add_argument(
+            "--context-parallel-size", "-cp",
+            **parallel_kwargs["context_parallel_size"])
         parallel_group.add_argument("--data-parallel-size", "-dp",
                                     **parallel_kwargs["data_parallel_size"])
         parallel_group.add_argument(
@@ -1310,6 +1314,7 @@ def create_engine_config(
         parallel_config = ParallelConfig(
             pipeline_parallel_size=self.pipeline_parallel_size,
             tensor_parallel_size=self.tensor_parallel_size,
+            context_parallel_size=self.context_parallel_size,
             data_parallel_size=self.data_parallel_size,
             data_parallel_rank=self.data_parallel_rank or 0,
             data_parallel_external_lb=data_parallel_external_lb,
@@ -1369,7 +1374,7 @@ def create_engine_config(
             long_prefill_token_threshold=self.long_prefill_token_threshold,
             disable_hybrid_kv_cache_manager=self.
             disable_hybrid_kv_cache_manager,
-            async_scheduling=self.async_scheduling,
+            async_scheduling=self.async_scheduling
         )
 
         if not model_config.is_multimodal_model and self.default_mm_loras:

vllm/model_executor/layers/fused_moe/config.py

@@ -7,7 +7,7 @@
 
 import vllm.envs as envs
 from vllm.config import ParallelConfig
-from vllm.distributed import get_dp_group, get_tensor_model_parallel_rank
+from vllm.distributed import get_dp_group, get_tensor_model_parallel_rank, get_context_model_parallel_rank
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     GroupShape)
@@ -566,9 +566,11 @@ def biased_moe_quant_config(
 @dataclass
 class FusedMoEParallelConfig:
     tp_size: int
+    cp_size: int
     dp_size: int
     ep_size: int
     tp_rank: int
+    cp_rank: int
     dp_rank: int
     ep_rank: int
 
@@ -594,15 +596,15 @@ def use_deepep_ll_kernels(self):
                 and envs.VLLM_ALL2ALL_BACKEND == "deepep_low_latency")
 
     @staticmethod
-    def make(tp_size_: int, dp_size_: int,
+    def make(tp_size_: int, dp_size_: int, cp_size_: int,
              vllm_parallel_config: ParallelConfig) -> "FusedMoEParallelConfig":
         """
         Determine MoE parallel configuration. Based on the input `tp_size_`,
         `dp_size_` and vllm's parallel config, determine what
         level's of parallelism to use in the fused moe layer.
 
         Args:
-            tp_size_ (int): `tp_size` passed into the FusedMoE constructor.
+            tp_size_ (int): `tp_size` pa use_ep = (dp_size_ * tp_size_ssed into the FusedMoE constructor.
             dp_size_ (int): `dp_size` passed into the FusedMoE constructor.
             vllm_parallel_config (ParallelConfig): vLLM's parallel config
                 object which contains the `enable_expert_parallel` flag.
@@ -675,16 +677,20 @@ def flatten_tp_across_dp(dp_rank: int):
             tp_rank = dp_rank * tp_size_ + tp_rank
             return tp_size, tp_rank
 
-        use_ep = (dp_size_ * tp_size_ > 1
+        use_ep = (dp_size_ * tp_size_ * cp_size_ > 1
                   and vllm_parallel_config.enable_expert_parallel)
 
         dp_size = dp_size_
         dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
         tp_size, tp_rank = flatten_tp_across_dp(dp_rank)
+        cp_size = cp_size_
+        cp_rank = get_context_model_parallel_rank() if cp_size_ > 1 else 0
 
         if not use_ep:
             return FusedMoEParallelConfig(tp_size=tp_size,
                                           tp_rank=tp_rank,
+                                          cp_size=cp_size,
+                                          cp_rank=cp_rank,
                                           dp_size=dp_size,
                                           dp_rank=dp_rank,
                                           ep_size=1,
@@ -694,10 +700,12 @@ def flatten_tp_across_dp(dp_rank: int):
         assert use_ep
         # In EP, each device owns a set of experts fully. There is no tensor
         # parallel update tp_size, tp_rank, ep_size and ep_rank to reflect that.
-        ep_size = tp_size
-        ep_rank = tp_rank
+        ep_size = tp_size * cp_size
+        ep_rank = tp_rank + tp_size * cp_rank
         return FusedMoEParallelConfig(tp_size=1,
                                       tp_rank=0,
+                                      cp_size=1,
+                                      cp_rank=0,
                                       dp_size=dp_size,
                                       dp_rank=dp_rank,
                                       ep_size=ep_size,

vllm/model_executor/layers/fused_moe/layer.py

@@ -15,6 +15,7 @@
 from vllm.config.parallel import ExpertPlacementStrategy
 from vllm.distributed import (get_dp_group, get_ep_group,
                               get_tensor_model_parallel_world_size,
+                              get_context_model_parallel_world_size,
                               tensor_model_parallel_all_reduce)
 from vllm.distributed.eplb.eplb_state import EplbState
 from vllm.forward_context import ForwardContext, get_forward_context
@@ -828,6 +829,7 @@ def __init__(
         tp_size: Optional[int] = None,
         ep_size: Optional[int] = None,
         dp_size: Optional[int] = None,
+        cp_size: Optional[int] = None,
         prefix: str = "",
         custom_routing_function: Optional[Callable] = None,
         scoring_func: str = "softmax",
@@ -849,6 +851,8 @@ def __init__(
                     get_tensor_model_parallel_world_size())
         dp_size_ = (dp_size
                     if dp_size is not None else get_dp_group().world_size)
+        cp_size_ = (cp_size
+                    if cp_size is not None else get_context_model_parallel_world_size())
 
         self.is_sequence_parallel = is_sequence_parallel
         if self.is_sequence_parallel:
@@ -859,6 +863,7 @@ def __init__(
             FusedMoEParallelConfig.make(
                 tp_size_=tp_size_,
                 dp_size_=dp_size_,
+                cp_size_=cp_size_,
                 vllm_parallel_config=vllm_config.parallel_config))
 
         self.global_num_experts = num_experts + num_redundant_experts