[Celeste-jq]

Summary

Fixes NPU HunyuanImage3 fused MoE runtime failures when expert parallelism is enabled together with diffusion sequence parallelism. #3877. #4438 .

Root cause

HunyuanImage3 is a diffusion MoE model. vllm-omni has its own diffusion parallel dimensions:

TP  = tensor_parallel_size
SP  = sequence_parallel_size
CFG = cfg_parallel_size
DP  = data_parallel_size

vLLM FusedMoE, however, expects MoE expert placement to be expressed through its native dimensions:

MoE EP size = TP * PCP * DP

Before this patch, the NPU HunyuanImage3 fused MoE path did not consistently map diffusion SP/CFG into vLLM's MoE parallel semantics. This could make local expert placement and token dispatch disagree, causing runtime failures such as:

RuntimeError: shape '[8, 16]' is invalid for input of size 64

What this PR changes

This patch keeps the change in vllm-omni and does not modify vLLM or vllm-ascend.

For diffusion MoE models with enable_expert_parallel=True, vllm-omni maps:

diffusion TP       -> vLLM TP
diffusion SP       -> vLLM PCP
diffusion CFG * DP -> vLLM DP

So vLLM FusedMoE computes:

EP = TP * PCP * DP = TP * SP * CFG * DP

The mapping is gated on both:

od_config.is_moe
enable_expert_parallel

This avoids changing vLLM PCP/DP semantics for non-MoE models or for MoE models running with expert parallelism disabled.

Implementation details

• In diffusion_worker.py, set vLLM prefill_context_parallel_size and adjusted data_parallel_size only for MoE + EP.
• In parallel_state.py, expose diffusion _SP as vLLM _PCP, use cfg-dp as vLLM _DP when needed, and build vLLM _EP from tp-sp-cfg-dp.
• In hunyuan_image3_transformer.py, let HunyuanImage3 FusedMoE use vLLM PCP only when EP is enabled; otherwise pass pcp_size=1 to avoid requiring a PCP group.
• In platforms/npu/models/hunyuan_fused_moe.py, remove the previous effective-TP rewrite and keep only the NPU backend glue needed to sync vllm-ascend EP/MC2 state with vLLM's EP group.

Validation

Tested on 8 NPUs with HunyuanImage3 DiT-only offline serving.

TP4 + SP2 + EP enabled

Config highlights:

tensor_parallel_size: 4
sequence_parallel_size: 2
ulysses_degree: 2
ring_degree: 1
enable_expert_parallel: true
cfg_parallel_size: 1
data_parallel_size: 1

Expected MoE mapping:

vLLM TP  = 4
vLLM PCP = 2
vLLM DP  = 1
vLLM EP  = 8

Observed result:

e2e_requests                1
e2e_wall_time_ms            11,972.534
e2e_stage_0_wall_time_ms    11,970.978
denoise_step_latency_ms     239.397
image_pixels                1,048,576
output_unit_count           1
output_unit_type            image
stage_gen_time_ms           11,969.872

TP4 + SP2 + EP disabled

Config highlights:

tensor_parallel_size: 4
sequence_parallel_size: 2
ulysses_degree: 2
ring_degree: 1
enable_expert_parallel: false
cfg_parallel_size: 1
data_parallel_size: 1

Expected behavior:

No vLLM PCP/DP MoE mapping is enabled.
HunyuanImage3 FusedMoE uses pcp_size=1.

Observed result:

e2e_requests                1
e2e_wall_time_ms            9,407.734
e2e_stage_0_wall_time_ms    9,406.435
denoise_step_latency_ms     188.118
image_pixels                1,048,576
output_unit_count           1
output_unit_type            image
stage_gen_time_ms           9,405.898

Scope

The MoE parallel mapping is only enabled for diffusion MoE models when expert parallelism is enabled.

NPU-specific code only synchronizes vllm-ascend's MoE runtime state with the vLLM EP group and MC2 group. It does not redefine a separate NPU-only MoE formula.

GPU/XPU/default paths are not given a separate formula; the goal is to keep MoE TP/PCP/DP/EP semantics aligned across backends.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[Celeste-jq]

@zengchuang-hw @jiangmengyu18 @Gaohan123 @gcanlin @fan2956 PTAL

[Celeste-jq]

@hsliuustc0106 PTAL

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[Bounty-hunter]

we need to set pcp_size, dp_size to correct number?

[Celeste-jq]

For EP=true, pcp_size/dp_size are intentionally left as None so vLLM FusedMoE reads them from the already initialized vLLM PCP/DP groups and then computes the MoE EP layout from those sizes.
For EP=false, pcp_size=1 is required to avoid querying an uninitialized PCP group.

[jiangmengyu18]

Here, assigning a value to _EP in vllm_ascend.distributed.parallel_state is meaningless, because Ascend uses _EP from vllm.distributed.parallel_state, and vllm.distributed.parallel_state.EP has already been assigned a value in vllm_omni/diffusion/distributed/parallel_state.py.

[Celeste-jq]

you're right. done.

[Bounty-hunter]

LGTM

[Celeste-jq]

@hsliuustc0106 @gcanlin @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @yenuo26 @jiangmengyu18 @Bounty-hunter
This PR is now ready for merge. PTAL.

[yenuo26]

For further assurance, can we add NPU types to the hardware labels of the relevant test cases, or add E2E test cases for the corresponding scenarios, in subsequent PRs?

[jiangmengyu18]

For further assurance, can we add NPU types to the hardware labels of the relevant test cases, or add E2E test cases for the corresponding scenarios, in subsequent PRs?

Sure, we'll address that in a subsequent PR. This PR is intended to be a general change that works for both GPU and NPU.

[hsliuustc0106]

This dropped the old od_config.is_moe guard. Non-MoE diffusion configs with enable_expert_parallel=True now create vLLM PCP/DP/EP state instead of failing fast, even though this remap should be MoE-only. Please keep the remap gated on both EP and od_config.is_moe, and add a regression test for non-MoE + EP.

[Celeste-jq]

done

[Gaohan123]

LGTM. Thanks