[BugFix] bug fix for DeepSeek eagle3 in Attn-DP mode (#20492)

This commit is contained in:
khalilzhk
2026-03-19 14:48:46 +08:00
committed by GitHub
parent fd05532da1
commit 574572b21b

View File

@@ -44,7 +44,6 @@ 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
@@ -74,6 +73,7 @@ from sglang.srt.layers.communicator_nsa_cp import NSACPLayerCommunicator
from sglang.srt.layers.dp_attention import (
get_attention_cp_rank,
get_attention_cp_size,
get_attention_tp_group,
get_attention_tp_rank,
get_attention_tp_size,
is_dp_attention_enabled,
@@ -1998,7 +1998,7 @@ class DeepseekV2Model(nn.Module):
with ctx:
if i in self.layers_to_capture:
if self.enable_a2a_moe and i > self.first_k_dense_replace:
aux_hidden_state = tensor_model_parallel_all_gather(
aux_hidden_state = get_attention_tp_group().all_gather(
hidden_states + residual, dim=0
)
aux_hidden_states.append(aux_hidden_state)