[MoE] Preserve TP sharding with explicit EP

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)

tests/kernels/moe/test_hybrid_tp_ep_config.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import torch
+
+import vllm.config as config_module
+import vllm.distributed as distributed
+import vllm.model_executor.layers.fused_moe.config as moe_config
+import vllm.model_executor.layers.fused_moe.runner.moe_runner as moe_runner
+from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
+from vllm.model_executor.layers.fused_moe.expert_map_manager import (
+    determine_expert_map,
+)
+from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
+
+
+def test_explicit_ep_preserves_tensor_parallelism(monkeypatch) -> None:
+    monkeypatch.setattr(
+        moe_config,
+        "get_dp_group",
+        lambda: SimpleNamespace(world_size=2, rank_in_group=1),
+    )
+    monkeypatch.setattr(
+        moe_config,
+        "get_ep_group",
+        lambda: SimpleNamespace(world_size=2, rank_in_group=1),
+    )
+    monkeypatch.setattr(
+        moe_config,
+        "get_pcp_group",
+        lambda: SimpleNamespace(world_size=1, rank_in_group=0),
+    )
+    monkeypatch.setattr(moe_config, "get_tensor_model_parallel_rank", lambda: 3)
+    parallel_config = SimpleNamespace(
+        enable_expert_parallel=True,
+        expert_parallel_size=2,
+        all2all_backend="allgather_reducescatter",
+        enable_eplb=False,
+    )
+
+    result = FusedMoEParallelConfig.make(
+        tp_size_=4,
+        pcp_size_=1,
+        dp_size_=2,
+        sp_size_=1,
+        vllm_parallel_config=parallel_config,
+    )
+
+    assert result.tp_size == 4
+    assert result.tp_rank == 3
+    assert result.ep_size == 2
+    assert result.ep_rank == 1
+    assert result.dp_size == 2
+    assert result.dp_rank == 1
+
+
+def test_explicit_ep_owns_half_of_global_experts() -> None:
+    local_count, expert_map, _ = determine_expert_map(
+        ep_size=2,
+        ep_rank=1,
+        global_num_experts=3584,
+    )
+
+    assert local_count == 1792
+    assert expert_map is not None
+    assert torch.count_nonzero(expert_map >= 0).item() == 1792
+    assert torch.all(expert_map[:1792] == -1)
+    assert torch.equal(expert_map[1792:], torch.arange(1792, dtype=torch.int32))
+
+
+def test_weight_filter_uses_explicit_ep_ownership(monkeypatch) -> None:
+    parallel_config = SimpleNamespace(
+        enable_expert_parallel=True,
+        enable_ep_weight_filter=True,
+        enable_eplb=False,
+        expert_parallel_size=2,
+        data_parallel_size=2,
+        tensor_parallel_size=4,
+        prefill_context_parallel_size=1,
+        expert_placement_strategy="linear",
+    )
+    monkeypatch.setattr(
+        config_module,
+        "get_current_vllm_config",
+        lambda: SimpleNamespace(parallel_config=parallel_config),
+    )
+    monkeypatch.setattr(
+        distributed,
+        "get_ep_group",
+        lambda: SimpleNamespace(world_size=2, rank_in_group=1),
+    )
+    loader = DefaultModelLoader.__new__(DefaultModelLoader)
+    loader.local_expert_ids = None
+    model_config = SimpleNamespace(is_moe=True, get_num_experts=lambda: 3584)
+
+    loader._init_ep_weight_filter(model_config)
+
+    assert loader.local_expert_ids == set(range(1792, 3584))
+
+
+def test_hybrid_ep_keeps_late_tp_reduction(monkeypatch) -> None:
+    calls = []
+
+    def fake_all_reduce(states):
+        calls.append(states)
+        return states + 1
+
+    monkeypatch.setattr(moe_runner, "tensor_model_parallel_all_reduce", fake_all_reduce)
+    runner = SimpleNamespace(
+        moe_config=SimpleNamespace(
+            is_sequence_parallel=False,
+            skip_final_all_reduce=False,
+            tp_size=4,
+            ep_size=2,
+        ),
+        _fused_output_is_reduced=False,
+    )
+    states = torch.zeros(2, 4)
+
+    result = MoERunner._maybe_reduce_final_output(
+        runner,
+        states,
+        trunc_size=None,
+        output_is_reduced=False,
+    )
+
+    assert calls == [states]
+    assert torch.equal(result, states + 1)

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,

vllm/model_executor/layers/fused_moe/config.py

@@ -8,7 +8,12 @@
 
 from vllm.config import ParallelConfig, SchedulerConfig
 from vllm.config.kernel import MoEBackend
-from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank
+from vllm.distributed import (
+    get_dp_group,
+    get_ep_group,
+    get_pcp_group,
+    get_tensor_model_parallel_rank,
+)
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
@@ -1214,6 +1219,22 @@ def make(
         dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
         pcp_size = pcp_size_
         pcp_rank = get_pcp_group().rank_in_group if pcp_size > 1 else 0
+        if use_ep and vllm_parallel_config.expert_parallel_size is not None:
+            ep_group = get_ep_group()
+            return FusedMoEParallelConfig(
+                tp_size=tp_size_,
+                tp_rank=(get_tensor_model_parallel_rank() if tp_size_ > 1 else 0),
+                pcp_size=pcp_size,
+                pcp_rank=pcp_rank,
+                dp_size=dp_size,
+                dp_rank=dp_rank,
+                ep_size=ep_group.world_size,
+                ep_rank=ep_group.rank_in_group,
+                sp_size=1,
+                use_ep=True,
+                all2all_backend=vllm_parallel_config.all2all_backend,
+                enable_eplb=vllm_parallel_config.enable_eplb,
+            )
         tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
             tp_size_, dp_size_, dp_rank, pcp_size_, pcp_rank
         )

vllm/model_executor/model_loader/default_loader.py

@@ -383,18 +383,23 @@ def _init_ep_weight_filter(self, model_config: ModelConfig) -> None:
         #   ep_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank
         from vllm.distributed import (
             get_dp_group,
+            get_ep_group,
             get_pcp_group,
             get_tensor_model_parallel_rank,
         )
 
         dp_size = parallel_config.data_parallel_size
         tp_size = parallel_config.tensor_parallel_size
         pcp_size = parallel_config.prefill_context_parallel_size
-        dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
-        tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0
-        pcp_rank = get_pcp_group().rank_in_group if pcp_size > 1 else 0
-        ep_size = dp_size * pcp_size * tp_size
-        ep_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank
+        if parallel_config.expert_parallel_size is not None:
+            ep_size = get_ep_group().world_size
+            ep_rank = get_ep_group().rank_in_group
+        else:
+            dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
+            tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0
+            pcp_rank = get_pcp_group().rank_in_group if pcp_size > 1 else 0
+            ep_size = dp_size * pcp_size * tp_size
+            ep_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank
 
         self.local_expert_ids = compute_local_expert_ids(
             num_experts,