Expert Parallel¶

EP_SIZE = TP_SIZE × SP_SIZE × CFG_SIZE × DP_SIZE

ep_size = 8  # Expert Parallel size (typically equals TP size)
num_experts = 64
num_local_experts = num_experts // ep_size  # 8 experts per card

# Check divisibility
assert num_experts % ep_size == 0, "Experts must be divisible by EP size"

from vllm.model_executor.layers.linear import ReplicatedLinear
class HunYuanSparseMoeBlock(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        layer_id: int = -1,
        prefix: str = "",
    ):
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
        self.n_routed_experts = config.num_experts  # 64

        # Calculate local experts per rank (key for EP)
        if self.tp_size > self.n_routed_experts:
            raise ValueError(f"TP size {self.tp_size} > experts {self.n_routed_experts}")

        # Routing gate (replicated on all ranks, computes scores for all tokens to all experts)
        self.gate = ReplicatedLinear(
            config.hidden_size,
            config.num_experts,
            bias=False,
            quant_config=None,
            prefix=f"{prefix}.gate",
        )

        # EP expert layer (factory loads platform-specific implementation)
        self.experts = HunyuanFusedMoE(...)

from vllm.utils.import_utils import resolve_obj_by_qualname
# Call during __init__ or model loading
op_name = "hunyuan_fused_moe"

# Prepare EP runtime: establish communication groups, assign local expert indices, init _expert_map
current_omni_platform.prepare_diffusion_op_runtime(op_name)

# Factory automatically resolves platform implementation (GPU: FusedMoE / NPU: AscendFusedMoE)
impl = resolve_obj_by_qualname(
    current_omni_platform.get_diffusion_model_impl_qualname(op_name)
)

# Get expert parameter mapping (different per rank)
expert_mapping = HunyuanFusedMoE.make_expert_params_mapping(
    model=self,
    ckpt_gate_proj_name="gate_proj",
    ckpt_down_proj_name="down_proj",
    ckpt_up_proj_name="up_proj",
    num_experts=64,
    num_redundant_experts=0,
)
# Returns: [(param_name, weight_name, expert_id, shard_id), ...]
# Note: Each rank only contains mappings for its local expert_ids

# Filter non-local experts during loading
for name, loaded_weight in weights:
    if "mlp.experts" in name:
        # Parse expert_id from weight name (implementation needed)
        expert_id = parse_expert_id_from_name(name)
        local_expert_start = (ep_rank) * num_local_experts
        local_expert_end = (ep_rank + 1) * num_local_experts

        if not (local_expert_start <= expert_id < local_expert_end):
            continue  # Skip non-local expert weights

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    orig_shape = hidden_states.shape
    hidden_states = hidden_states.view(-1, hidden_states.shape[-1])

    # 1. Global routing computation (all tokens, all expert scores)
    # hidden_states: [num_tokens, hidden_dim] (full tensor)
    router_logits, _ = self.gate(hidden_states)  # [num_tokens, num_experts]

    # 2. EP dispatch and compute (HunyuanFusedMoE handles all_to_all internally)
    # - Dispatch: Send tokens to target ranks based on router_logits
    # - Local Compute: Each rank processes only its num_local_experts
    # - Combine: Results returned to original token positions
    final_hidden_states = self.experts(
        hidden_states=hidden_states,
        router_logits=router_logits,
    )

    # 3. Add shared expert output (not EP, computed on all ranks)
    if self.shared_mlp is not None:
        shared_out = self.shared_mlp(hidden_states)
        final_hidden_states = final_hidden_states + shared_out

    # 4. Tensor Parallel All-Reduce (synchronize across TP group)
    if self.tp_size > 1:
        final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
            final_hidden_states
        )

    return final_hidden_states.view(orig_shape)