[NPU] Fix HunyuanImage3 MoE EP+SP parallel mapping (#4344)

Signed-off-by: Celeste-jq <591998922@qq.com>

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+
+import vllm_omni.diffusion.distributed.parallel_state as parallel_state
+
+
+class _FakeGroup:
+    def __init__(
+        self,
+        group_ranks: list[list[int]],
+        local_rank: int,
+        parallel_mode: str,
+        **kwargs,
+    ) -> None:
+        self.group_ranks = group_ranks
+        self.parallel_mode = parallel_mode
+        self.device_group = object()
+        self.ulysses_group = kwargs.get("ulysses_group")
+        self.ring_group = kwargs.get("ring_group")
+        self.local_group = next(group for group in group_ranks if local_rank in group)
+        self.world_size = len(self.local_group)
+        self.rank_in_group = self.local_group.index(local_rank)
+
+    def destroy(self) -> None:
+        pass
+
+
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
+    """MoE+EP rank layout should map SP->PCP, CFG*DP->DP, and TP*SP*CFG*DP->EP."""
+    local_rank = 0
+    world_size = 32
+    created_groups: list[_FakeGroup] = []
+
+    def fake_init_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        parallel_mode=None,
+        group_name=None,
+        **kwargs,
+    ):
+        del backend, group_name
+        group = _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            parallel_mode or "",
+            **kwargs,
+        )
+        created_groups.append(group)
+        return group
+
+    fake_world_group = SimpleNamespace(
+        rank_in_group=local_rank,
+        local_rank=local_rank,
+        device_group=object(),
+    )
+    fake_forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(is_moe=True))
+    monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo")
+    monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
+    monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
+    monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
+
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+        monkeypatch.setattr(parallel_state, name, None)
+    for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+
+    parallel_state.initialize_model_parallel(
+        tensor_parallel_size=2,
+        sequence_parallel_size=2,
+        ulysses_degree=2,
+        ring_degree=1,
+        pipeline_parallel_size=2,
+        cfg_parallel_size=2,
+        data_parallel_size=2,
+        enable_expert_parallel=True,
+        backend="gloo",
+    )
+
+    assert parallel_state.vllm_parallel_state._PCP is parallel_state._SP
+    assert parallel_state.vllm_parallel_state._PCP.world_size == 2
+    assert parallel_state._DP.world_size == 2
+    assert parallel_state.vllm_parallel_state._DP is not parallel_state._DP
+    assert parallel_state.vllm_parallel_state._DP.world_size == 4
+    assert parallel_state.vllm_parallel_state._EP.world_size == 16
+    assert parallel_state.vllm_parallel_state._TP.world_size == 2
+    assert parallel_state._PP.world_size == 2
+
+    assert parallel_state.vllm_parallel_state._PCP.local_group == [0, 2]
+    assert parallel_state.vllm_parallel_state._DP.local_group == [0, 8, 16, 24]
+    assert parallel_state.vllm_parallel_state._EP.local_group == [
+        0,
+        1,
+        2,
+        3,
+        8,
+        9,
+        10,
+        11,
+        16,
+        17,
+        18,
+        19,
+        24,
+        25,
+        26,
+        27,
+    ]
+    assert parallel_state.get_expert_parallel_group_ranks() == [
+        [
+            0,
+            1,
+            2,
+            3,
+            8,
+            9,
+            10,
+            11,
+            16,
+            17,
+            18,
+            19,
+            24,
+            25,
+            26,
+            27,
+        ],
+        [
+            4,
+            5,
+            6,
+            7,
+            12,
+            13,
+            14,
+            15,
+            20,
+            21,
+            22,
+            23,
+            28,
+            29,
+            30,
+            31,
+        ],
+    ]
+
+    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "expert"]
+    assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
+
+
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
+    """vLLM DP should keep diffusion DP when expert parallelism is not enabled."""
+    local_rank = 0
+    world_size = 8
+
+    def fake_init_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        parallel_mode=None,
+        group_name=None,
+        **kwargs,
+    ):
+        del backend, group_name
+        return _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            parallel_mode or "",
+            **kwargs,
+        )
+
+    fake_world_group = SimpleNamespace(
+        rank_in_group=local_rank,
+        local_rank=local_rank,
+        device_group=object(),
+    )
+    monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo")
+    monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
+    monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
+
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+        monkeypatch.setattr(parallel_state, name, None)
+    for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+
+    parallel_state.initialize_model_parallel(
+        tensor_parallel_size=2,
+        sequence_parallel_size=1,
+        ulysses_degree=1,
+        ring_degree=1,
+        pipeline_parallel_size=1,
+        cfg_parallel_size=2,
+        data_parallel_size=2,
+        enable_expert_parallel=False,
+        backend="gloo",
+    )
+
+    assert parallel_state._DP.world_size == 2
+    assert parallel_state.vllm_parallel_state._DP is parallel_state._DP
+    assert parallel_state.vllm_parallel_state._DP.world_size == 2
+    assert parallel_state.vllm_parallel_state._DP.local_group == [0, 4]
+    assert parallel_state.vllm_parallel_state._PCP is None
+    assert parallel_state.vllm_parallel_state._EP is None
+    assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None
+
+
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
+    """Non-MoE diffusion configs should not create vLLM PCP/DP/EP remap state."""
+    local_rank = 0
+    world_size = 8
+
+    def fake_init_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        parallel_mode=None,
+        group_name=None,
+        **kwargs,
+    ):
+        del backend, group_name
+        return _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            parallel_mode or "",
+            **kwargs,
+        )
+
+    fake_world_group = SimpleNamespace(
+        rank_in_group=local_rank,
+        local_rank=local_rank,
+        device_group=object(),
+    )
+    fake_forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(is_moe=False))
+    monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo")
+    monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
+    monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
+    monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
+
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+        monkeypatch.setattr(parallel_state, name, None)
+    for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+
+    with pytest.raises(RuntimeError, match="Expert parallelism enabled for a non-MoE model"):
+        parallel_state.initialize_model_parallel(
+            tensor_parallel_size=2,
+            sequence_parallel_size=2,
+            ulysses_degree=2,
+            ring_degree=1,
+            pipeline_parallel_size=1,
+            cfg_parallel_size=2,
+            data_parallel_size=1,
+            enable_expert_parallel=True,
+            backend="gloo",
+        )
+
+    assert parallel_state.vllm_parallel_state._PCP is None
+    assert parallel_state.vllm_parallel_state._EP is None
+    assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None

vllm_omni/diffusion/distributed/parallel_state.py

@@ -35,7 +35,6 @@ from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion import envs
-from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.platforms import current_omni_platform
 
@@ -61,6 +60,11 @@ _DP: GroupCoordinator | None = None
 _FS: GroupCoordinator | None = None  # Fully Sharded (HSDP shard dimension)
 _DIT: GroupCoordinator | None = None
 
+# Rank-layout metadata for expert parallelism. This is not a process group;
+# it is reused by platform-specific runtimes that must build companion groups
+# with the same rank layout as vLLM EP.
+_EXPERT_PARALLEL_GROUP_RANKS: list[list[int]] | None = None
+
 
 def generate_masked_orthogonal_rank_groups(
     world_size: int, parallel_size: list[int], mask: list[bool]
@@ -190,7 +194,6 @@ class RankGenerator:
         self.fs = fs
         self.rank_offset = rank_offset
         self.world_size = tp * sp * pp * cfg * dp
-        self.ep = tp * sp * cfg * dp  # EP level exclude PP
 
         self.name_to_size = {
             "tp": self.tp,
@@ -253,15 +256,6 @@ class RankGenerator:
                 ranks.append(group)
             return ranks
 
-        if token == "ep":
-            ranks = []
-            num_pp_stages = self.pp
-            for i in range(num_pp_stages):
-                start = i * self.ep + self.rank_offset
-                end = start + self.ep
-                ranks.append(list(range(start, end)))
-            return ranks
-
         mask = self.get_mask(self.order, token)
         ranks = generate_masked_orthogonal_rank_groups(self.world_size, self.ordered_size, mask)
         if self.rank_offset > 0:
@@ -309,6 +303,12 @@ def get_ring_parallel_rank():
     return get_sp_group().ring_rank
 
 
+def get_expert_parallel_group_ranks() -> list[list[int]]:
+    assert vllm_parallel_state._EP is not None, "expert parallel group is not initialized"
+    assert _EXPERT_PARALLEL_GROUP_RANKS is not None, "expert parallel group ranks are not initialized"
+    return _EXPERT_PARALLEL_GROUP_RANKS
+
+
 # PP
 def get_pp_group() -> PipelineGroupCoordinator:
     assert _PP is not None, "pipeline model parallel group is not initialized"
@@ -782,6 +782,13 @@ def initialize_model_parallel(
         fs=fully_shard_degree,
         order="tp-sp-pp-cfg-dp",
     )
+    use_moe_parallel_mapping = False
+    if enable_expert_parallel:
+        od_config = get_forward_context().omni_diffusion_config
+        use_moe_parallel_mapping = bool(od_config and od_config.is_moe)
+        if not use_moe_parallel_mapping:
+            raise RuntimeError("Expert parallelism enabled for a non-MoE model")
+
     sp_group_ranks = rank_generator.get_ranks("sp")
     global _DP
     assert _DP is None, "data parallel group is already initialized"
@@ -828,6 +835,10 @@ def initialize_model_parallel(
         ulysses_group=ulysses_pg,
         ring_group=ring_pg,
     )
+    if use_moe_parallel_mapping:
+        # Diffusion normally uses its own SP group. Map it to vLLM PCP only for
+        # expert-parallel runtimes that rely on vLLM FusedMoE group semantics.
+        vllm_parallel_state._PCP = _SP
 
     assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized"
     vllm_parallel_state._TP = init_model_parallel_group(
@@ -836,8 +847,17 @@ def initialize_model_parallel(
         backend=backend,
         parallel_mode="tensor",
     )
+    if use_moe_parallel_mapping and cfg_parallel_size > 1:
+        # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP
+        # only when constructing the vLLM EP layout for expert-parallel paths.
+        vllm_parallel_state._DP = init_model_parallel_group(
+            group_ranks=rank_generator.get_ranks("cfg-dp"),
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            parallel_mode="data",
+        )
 
-    global _FS
+    global _FS, _EXPERT_PARALLEL_GROUP_RANKS
     assert _FS is None, "fully shard group is already initialized"
     _FS = init_model_parallel_group(
         group_ranks=rank_generator.get_ranks("fs", independent_ranks=True),
@@ -846,34 +866,40 @@ def initialize_model_parallel(
         parallel_mode="fully_shard",
     )
 
-    if enable_expert_parallel:
-        od_config: OmniDiffusionConfig | None = get_forward_context().omni_diffusion_config
-        if od_config and od_config.is_moe:
-            vllm_parallel_state._EP = init_model_parallel_group(
-                group_ranks=rank_generator.get_ranks("ep"),
-                local_rank=get_world_group().local_rank,
-                backend=backend,
-                parallel_mode="expert",
-            )
-        else:
-            raise RuntimeError("Expert parallelism enabled for a non-MoE model ")
+    _EXPERT_PARALLEL_GROUP_RANKS = None
+    if use_moe_parallel_mapping:
+        ep_group_ranks = rank_generator.get_ranks("tp-sp-cfg-dp")
+        vllm_parallel_state._EP = init_model_parallel_group(
+            group_ranks=ep_group_ranks,
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            parallel_mode="expert",
+        )
+        _EXPERT_PARALLEL_GROUP_RANKS = ep_group_ranks
 
     init_dit_group(dit_parallel_size, backend)
 
 
 def destroy_model_parallel():
     """Set the groups to none and destroy them."""
-    global _DP
+    global _DP, _CFG, _SP, _PP, _FS, _EXPERT_PARALLEL_GROUP_RANKS
+
+    if vllm_parallel_state._DP and vllm_parallel_state._DP is not _DP:
+        vllm_parallel_state._DP.destroy()
+    vllm_parallel_state._DP = None
+
     if _DP:
         _DP.destroy()
     _DP = None
 
-    global _CFG
     if _CFG:
         _CFG.destroy()
     _CFG = None
 
-    global _SP
+    if vllm_parallel_state._PCP and vllm_parallel_state._PCP is not _SP:
+        vllm_parallel_state._PCP.destroy()
+    vllm_parallel_state._PCP = None
+
     if _SP:
         _SP.destroy()
     _SP = None
@@ -885,13 +911,12 @@ def destroy_model_parallel():
     if vllm_parallel_state._EP:
         vllm_parallel_state._EP.destroy()
     vllm_parallel_state._EP = None
+    _EXPERT_PARALLEL_GROUP_RANKS = None
 
-    global _PP
     if _PP:
         _PP.destroy()
     _PP = None
 
-    global _FS
     if _FS:
         _FS.destroy()
     _FS = None

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -29,7 +29,7 @@ from transformers.modeling_outputs import (
     CausalLMOutputWithPast,
 )
 from transformers.modeling_utils import PreTrainedModel
-from vllm.config import CacheConfig
+from vllm.config import CacheConfig, get_current_vllm_config
 from vllm.distributed import (
     get_tensor_model_parallel_world_size,
 )
@@ -1624,6 +1624,7 @@ class HunYuanSparseMoeBlock(nn.Module):
         else:
             self.shared_mlp = None
 
+        enable_expert_parallel = get_current_vllm_config().parallel_config.enable_expert_parallel
         self.experts = HunyuanFusedMoE(
             shared_experts=self.shared_mlp,
             num_experts=self.n_routed_experts,
@@ -1635,7 +1636,7 @@ class HunYuanSparseMoeBlock(nn.Module):
             prefix=f"{prefix}.experts",
             enable_eplb=self.enable_eplb,
             num_redundant_experts=self.n_redundant_experts,
-            pcp_size=1,
+            pcp_size=None if enable_expert_parallel else 1,
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -246,9 +246,18 @@ class DiffusionWorker:
         # Create vllm_config for parallel configuration. Pass explicit device_config
         # so DeviceConfig does not rely on current_platform in worker subprocesses.
         vllm_config = _create_diffusion_worker_vllm_config(self.device, self.od_config)
-        vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
-        vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
-        vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel
+        parallel_config = self.od_config.parallel_config
+        vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size
+        vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size
+        if parallel_config.enable_expert_parallel and self.od_config.is_moe:
+            # Diffusion uses its own DP/CFG/SP groups normally. vLLM groups are
+            # only remapped for expert-parallel runtimes that consume vLLM's
+            # FusedMoE/EP semantics.
+            vllm_config.parallel_config.data_parallel_size = (
+                parallel_config.data_parallel_size * parallel_config.cfg_parallel_size
+            )
+            vllm_config.parallel_config.prefill_context_parallel_size = parallel_config.sequence_parallel_size
+        vllm_config.parallel_config.enable_expert_parallel = parallel_config.enable_expert_parallel
         vllm_config.profiler_config = self.od_config.profiler_config
         vllm_config.model_config = _make_diffusion_vllm_model_config(self.od_config)  # type: ignore[assignment]
         vllm_config.quant_config = self.od_config.quantization_config

vllm_omni/platforms/npu/models/hunyuan_fused_moe.py

@@ -7,7 +7,6 @@ import torch
 import vllm.forward_context as _vllm_fc
 from vllm.config import VllmConfig
 from vllm.distributed import get_ep_group
-from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
 from vllm.distributed.parallel_state import (
     init_model_parallel_group as vllm_init_model_parallel_group,
 )
@@ -17,7 +16,7 @@ from vllm_ascend.ops.fused_moe.moe_comm_method import _MoECommMethods
 from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type
 
 from vllm_omni.diffusion.distributed.parallel_state import (
-    get_data_parallel_world_size,
+    get_expert_parallel_group_ranks,
     get_world_group,
 )
 from vllm_omni.diffusion.forward_context import get_forward_context as omni_get_ctx
@@ -42,20 +41,14 @@ def _set_hunyuan_fused_moe_forward_context(num_tokens: int) -> None:
 
 
 def _init_mc2_group_for_diffusion(
-    world_size: int,
-    data_parallel_size: int,
-    tensor_parallel_size: int,
     backend: str,
     local_rank: int,
+    group_ranks: list[list[int]],
 ) -> None:
     import vllm_ascend.distributed.parallel_state as vllm_ascend_parallel_state
 
     if getattr(vllm_ascend_parallel_state, "_MC2", None) is not None:
         return
-    all_ranks = torch.arange(world_size).reshape(-1, data_parallel_size * tensor_parallel_size)
-    group_ranks = all_ranks.unbind(0)
-    group_ranks = [x.tolist() for x in group_ranks]
-
     vllm_ascend_parallel_state._MC2 = vllm_init_model_parallel_group(
         group_ranks,
         local_rank,
@@ -82,20 +75,17 @@ def _select_moe_comm_method(vllm_config: VllmConfig) -> MoECommType | None:
 
 
 def prepare_hunyuan_fused_moe_runtime() -> None:
-    world_size = torch.distributed.get_world_size()
-    data_parallel_size = get_data_parallel_world_size()
-    tensor_parallel_size = get_tensor_model_parallel_world_size()
-    backend = torch.distributed.get_backend(get_world_group().device_group)
-    local_rank = get_world_group().local_rank
-    _init_mc2_group_for_diffusion(
-        world_size=world_size,
-        data_parallel_size=data_parallel_size,
-        tensor_parallel_size=tensor_parallel_size,
-        backend=backend,
-        local_rank=local_rank,
-    )
+    vllm_config = omni_get_ctx().vllm_config
+    if vllm_config.parallel_config.enable_expert_parallel:
+        backend = torch.distributed.get_backend(get_world_group().device_group)
+        local_rank = get_world_group().local_rank
+        _init_mc2_group_for_diffusion(
+            backend=backend,
+            local_rank=local_rank,
+            group_ranks=get_expert_parallel_group_ranks(),
+        )
 
-    moe_comm_type = _select_moe_comm_method(vllm_config=omni_get_ctx().vllm_config)
+    moe_comm_type = _select_moe_comm_method(vllm_config=vllm_config)
     _ensure_forward_context_attr("num_tokens", int | None, None)
     _ensure_forward_context_attr("in_profile_run", bool, False)
     _ensure_forward_context_attr("moe_comm_type", MoECommType | None, moe_comm_type)