[DeepSeek V4] Fix fused MoE expert distribution (#57465)
Signed-off-by: Itay Alroy <ialroy@nvidia.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Itay Alroy committed
Sep 18, 2026 at 18:17 UTC
017dced6a6fd3cf430e4686a47b354da1cadfbf5
1 file changed
+5
-7
vllm/models/deepseek_v4/nvidia/model.py
+5
-7
@@ -993,20 +993,18 @@ class DeepseekV4MoE(nn.Module):
993
prefix: str,
994
) -> None:
995
parallel_config = vllm_config.parallel_config
996
- self.tp_rank = get_tensor_model_parallel_rank()
996
+ ep_group = get_ep_group()
997
+ ep_size = ep_group.world_size
998
+ ep_rank = ep_group.rank_in_group
999
1000
eplb_config = parallel_config.eplb_config
1001
self.n_redundant_experts = eplb_config.num_redundant_experts
1002
self.n_shared_experts = config.n_shared_experts or 0
1003
self.n_logical_experts = self.n_routed_experts
1004
self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
1003
- assert self.n_physical_experts % self.tp_size == 0, (
1004
- f"n_physical_experts={self.n_physical_experts} must be divisible by "
1005
- f"tp_size={self.tp_size}. Adjust num_redundant_experts."
1006
- )
1007
- self.n_local_physical_experts = self.n_physical_experts // self.tp_size
1005
+ self.n_local_physical_experts = self.n_physical_experts // ep_size
1006
self.n_local_experts = self.n_local_physical_experts
1009
- self.experts_start_idx = self.tp_rank * self.n_local_experts
1007
+ self.experts_start_idx = ep_rank * self.n_local_experts
1008
self.experts_end_idx = self.experts_start_idx + self.n_local_experts
1009
self.physical_expert_start = self.experts_start_idx
1010
self.physical_expert_end = self.experts_end_idx