diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 8c3f15719..33e820785 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -47,6 +47,7 @@ from sglang.srt.distributed import ( get_moe_expert_parallel_world_size, get_pp_group, get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce, ) from sglang.srt.environ import envs @@ -3127,6 +3128,10 @@ class DeepseekV2Model(nn.Module): ) ) self.layers_to_capture = [] + if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake(): + self.enable_a2a_moe = True + else: + self.enable_a2a_moe = False # llama_4_scaling: for supporting Mistral-Large-3 model self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None) @@ -3211,7 +3216,13 @@ class DeepseekV2Model(nn.Module): ) with ctx: if i in self.layers_to_capture: - aux_hidden_states.append(hidden_states + residual) + if self.enable_a2a_moe and i > self.first_k_dense_replace: + aux_hidden_state = tensor_model_parallel_all_gather( + hidden_states + residual, dim=0 + ) + aux_hidden_states.append(aux_hidden_state) + else: + aux_hidden_states.append(hidden_states + residual) layer = self.layers[i] hidden_states, residual = layer( positions,