diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index c50fc327e7608..d0a116b97997a 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -287,7 +287,10 @@ class DeepseekV2MoE(nn.Module): ) self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() - if config.n_shared_experts is None or self.is_rocm_aiter_moe_enabled: + self.is_fusion_moe_shared_experts_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ) + if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled: self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts @@ -327,7 +330,7 @@ class DeepseekV2MoE(nn.Module): num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, n_shared_experts=config.n_shared_experts - if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + if self.is_fusion_moe_shared_experts_enabled else None, )