[Core][MoE] Add explicit expert-parallel topology

tests/config/test_explicit_expert_parallel.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.config import ParallelConfig
+from vllm.distributed.parallel_state import _get_ep_group_ranks
+
+
+def test_explicit_ep_groups_preserve_tp_lanes() -> None:
+    all_ranks = torch.arange(8).reshape(1, 2, 1, 1, 4)
+
+    groups = _get_ep_group_ranks(
+        all_ranks,
+        data_parallel_size=2,
+        prefill_context_model_parallel_size=1,
+        tensor_model_parallel_size=4,
+        expert_parallel_size=2,
+    )
+
+    assert groups == [[0, 4], [1, 5], [2, 6], [3, 7]]
+
+
+def test_explicit_ep_disables_flattened_sequence_parallel() -> None:
+    config = ParallelConfig(
+        tensor_parallel_size=4,
+        data_parallel_size=2,
+        enable_expert_parallel=True,
+        expert_parallel_size=2,
+    )
+
+    assert not config.use_sequence_parallel_moe
+
+
+@pytest.mark.parametrize(
+    "overrides",
+    [
+        {"enable_expert_parallel": False},
+        {"pipeline_parallel_size": 2},
+        {"prefill_context_parallel_size": 2},
+        {"expert_parallel_size": 4},
+        {"all2all_backend": "deepep_low_latency"},
+        {"enable_eplb": True},
+        {"enable_elastic_ep": True},
+    ],
+)
+def test_explicit_ep_rejects_unsupported_topologies(overrides) -> None:
+    kwargs = {
+        "tensor_parallel_size": 4,
+        "data_parallel_size": 2,
+        "enable_expert_parallel": True,
+        "expert_parallel_size": 2,
+    }
+    kwargs.update(overrides)
+
+    with pytest.raises(ValueError):
+        ParallelConfig(**kwargs)

vllm/config/parallel.py

@@ -164,6 +164,13 @@ class ParallelConfig:
     """Whether the deployed model is MoE (if known)."""
     enable_expert_parallel: bool = False
     """Use expert parallelism instead of tensor parallelism for MoE layers."""
+    expert_parallel_size: int | None = None
+    """Experimental explicit expert-parallel size.
+
+    When unset, EP retains the standard flattened ``TP * PCP * DP`` size.
+    The initial explicit topology supports ``EP == DP`` with ``PCP == PP == 1``
+    and preserves tensor parallelism within each expert.
+    """
     enable_ep_weight_filter: bool = False
     """Skip non-local expert weights during model loading when expert
     parallelism is active.  Each rank only reads its own expert shard from
@@ -682,6 +689,7 @@ def use_sequence_parallel_moe(self) -> bool:
                 "nixl_ep",
             )
             and self.enable_expert_parallel
+            and self.expert_parallel_size is None
             and self.tensor_parallel_size > 1
             and self.data_parallel_size > 1
         )
@@ -904,6 +912,38 @@ def __post_init__(self) -> None:
 
         self.data_parallel_index = self.data_parallel_rank
 
+        if self.expert_parallel_size is not None:
+            if not self.enable_expert_parallel:
+                raise ValueError(
+                    "expert_parallel_size requires enable_expert_parallel=True."
+                )
+            if self.pipeline_parallel_size != 1:
+                raise ValueError(
+                    "Explicit expert_parallel_size does not yet support "
+                    "pipeline parallelism."
+                )
+            if self.prefill_context_parallel_size != 1:
+                raise ValueError(
+                    "Explicit expert_parallel_size does not yet support "
+                    "prefill context parallelism."
+                )
+            if self.expert_parallel_size != self.data_parallel_size:
+                raise ValueError(
+                    "The initial hybrid TP x EP implementation requires "
+                    "expert_parallel_size == data_parallel_size, but got "
+                    f"{self.expert_parallel_size} and {self.data_parallel_size}."
+                )
+            if self.all2all_backend != "allgather_reducescatter":
+                raise ValueError(
+                    "Explicit expert_parallel_size initially supports only "
+                    "allgather_reducescatter."
+                )
+            if self.enable_eplb or self.enable_elastic_ep:
+                raise ValueError(
+                    "Explicit expert_parallel_size does not yet support EPLB "
+                    "or elastic EP."
+                )
+
         if self.distributed_executor_backend == "external_launcher":
             os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
             logger.info("Disabling V1 multiprocessing for external launcher.")

vllm/distributed/parallel_state.py

@@ -1748,6 +1748,28 @@ def init_distributed_environment(
             _INNER_DP_WORLD = _WORLD
 
 
+def _get_ep_group_ranks(
+    all_ranks: torch.Tensor,
+    *,
+    data_parallel_size: int,
+    prefill_context_model_parallel_size: int,
+    tensor_model_parallel_size: int,
+    expert_parallel_size: int | None,
+) -> list[list[int]]:
+    if expert_parallel_size is not None:
+        # Preserve each TP group and form EP groups across DP ranks at the
+        # same TP lane. TP4/DP2 yields [0,4], [1,5], [2,6], [3,7].
+        group_ranks = all_ranks.transpose(1, 4).reshape(-1, expert_parallel_size)
+    else:
+        group_ranks = all_ranks.transpose(1, 2).reshape(
+            -1,
+            data_parallel_size
+            * prefill_context_model_parallel_size
+            * tensor_model_parallel_size,
+        )
+    return [ranks.tolist() for ranks in group_ranks.unbind(0)]
+
+
 def initialize_model_parallel(
     tensor_model_parallel_size: int = 1,
     pipeline_model_parallel_size: int = 1,
@@ -1924,17 +1946,13 @@ def initialize_model_parallel(
     assert _EP is None, "expert parallel group is already initialized"
     # Don't create EP group for dense models.
     if config.model_config is None or config.model_config.is_moe:
-        group_ranks = (
-            all_ranks.transpose(1, 2)
-            .reshape(
-                -1,
-                data_parallel_size
-                * prefill_context_model_parallel_size
-                * tensor_model_parallel_size,
-            )
-            .unbind(0)
+        group_ranks = _get_ep_group_ranks(
+            all_ranks,
+            data_parallel_size=data_parallel_size,
+            prefill_context_model_parallel_size=(prefill_context_model_parallel_size),
+            tensor_model_parallel_size=tensor_model_parallel_size,
+            expert_parallel_size=parallel_config.expert_parallel_size,
         )
-        group_ranks = [x.tolist() for x in group_ranks]
         use_all2all = parallel_config.use_all2all
         if enable_elastic_ep:
             _EP = _init_stateless_group(

vllm/engine/arg_utils.py

@@ -493,6 +493,7 @@ class EngineArgs:
     data_parallel_multi_port_external_lb: bool = False
     data_parallel_backend: DataParallelBackend = ParallelConfig.data_parallel_backend
     enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
+    expert_parallel_size: int | None = ParallelConfig.expert_parallel_size
     enable_ep_weight_filter: bool = ParallelConfig.enable_ep_weight_filter
     moe_backend: MoEBackend = KernelConfig.moe_backend
     linear_backend: LinearBackend = KernelConfig.linear_backend
@@ -1134,6 +1135,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             "-ep",
             **parallel_kwargs["enable_expert_parallel"],
         )
+        parallel_group.add_argument(
+            "--expert-parallel-size",
+            **parallel_kwargs["expert_parallel_size"],
+        )
         parallel_group.add_argument(
             "--enable-ep-weight-filter",
             **parallel_kwargs["enable_ep_weight_filter"],
@@ -2241,6 +2246,7 @@ def create_engine_config(
             data_parallel_hybrid_lb=self.data_parallel_hybrid_lb,
             is_moe_model=model_config.is_moe,
             enable_expert_parallel=self.enable_expert_parallel,
+            expert_parallel_size=self.expert_parallel_size,
             enable_ep_weight_filter=self.enable_ep_weight_filter,
             all2all_backend=self.all2all_backend,
             enable_elastic_ep=self.enable_elastic_ep,