diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 56d3fff29..591ce2d7f 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -1304,8 +1304,12 @@ class CommunicateSummableTensorPairFn: # DP scatter (if DP attention is enabled) if context.attn_dp_size > 1: + if get_tensor_model_parallel_world_size() == get_attention_dp_size(): + group = get_tp_group() + else: + group = get_attention_tp_group() hidden_states_output, global_hidden_states = ( - get_local_dp_buffer(), + get_local_dp_buffer(group), hidden_states, ) dp_scatter(hidden_states_output, global_hidden_states, forward_batch) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 1535a81db..4ab6d00e8 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -26,7 +26,11 @@ from sglang.jit_kernel.deepseek_v4 import ( fused_rope_inplace, ) from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config -from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size +from sglang.srt.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, + get_tp_group, +) from sglang.srt.environ import envs from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.layers.attention.dsv4.compressor import Compressor @@ -885,7 +889,10 @@ class DeepseekV4DecoderLayer(nn.Module): input_ids = input_ids[cp_rank::cp_size].contiguous() input_ids_global = input_ids elif _use_tp_moe_gather: - hidden_states, local_hidden_states = get_global_dp_buffer(), hidden_states + hidden_states, local_hidden_states = ( + get_global_dp_buffer(get_tp_group()), + hidden_states, + ) dp_gather_partial(hidden_states, local_hidden_states, forward_batch) _a2a_scatter_chunks: Optional[List[torch.Tensor]] = None if _use_tp_attn_a2a_scatter: @@ -901,7 +908,10 @@ class DeepseekV4DecoderLayer(nn.Module): input_ids_global=input_ids_global, ) if _use_tp_moe_gather: - hidden_states, global_hidden_states = get_local_dp_buffer(), hidden_states + hidden_states, global_hidden_states = ( + get_local_dp_buffer(get_tp_group()), + hidden_states, + ) dp_scatter(hidden_states, global_hidden_states, forward_batch) if _use_tp_attn_a2a_scatter: assert _a2a_scatter_chunks is not None