[scheduler] support cp

Signed-off-by: LookAround <lixushi@huawei.com>

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