diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 7de4337f9..30f068595 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -46,6 +46,7 @@ from sglang.srt.layers.dp_attention import ( get_attention_cp_rank, get_attention_cp_size, get_attention_dp_size, + get_attention_tp_group, get_attention_tp_rank, get_attention_tp_size, get_dp_global_num_tokens, @@ -855,7 +856,7 @@ class CommunicateSimpleFn: return tuple(gathered_hidden_states) hidden_states, local_hidden_states = ( - get_local_dp_buffer(), + get_local_dp_buffer(get_attention_tp_group()), hidden_states, ) attn_tp_all_gather_into_tensor( @@ -965,7 +966,7 @@ class CommunicateWithAllReduceAndLayerNormFn: if residual_input_mode == ScatterMode.SCATTERED and context.attn_tp_size > 1: residual, local_residual = ( - get_local_dp_buffer(), + get_local_dp_buffer(get_attention_tp_group()), residual, ) attn_tp_all_gather_into_tensor(residual, local_residual) @@ -982,7 +983,7 @@ class CommunicateWithAllReduceAndLayerNormFn: hidden_states += residual hidden_states, local_hidden_states = ( - get_global_dp_buffer(), + get_global_dp_buffer(get_tp_group()), hidden_states, ) dp_gather_partial(hidden_states, local_hidden_states, forward_batch) @@ -1210,8 +1211,12 @@ class CommunicateSummableTensorPairFn: context: CommunicateContext, allow_reduce_scatter: bool = False, ): + if get_tensor_model_parallel_world_size() == get_attention_dp_size(): + group = get_tp_group() + else: + group = get_attention_tp_group() hidden_states, global_hidden_states = ( - get_local_dp_buffer(), + get_local_dp_buffer(group), hidden_states, ) if should_use_dp_reduce_scatterv(): @@ -1237,7 +1242,7 @@ class CommunicateSummableTensorPairFn: hidden_states += residual residual = None hidden_states, local_hidden_states = ( - get_local_dp_buffer(), + get_local_dp_buffer(get_attention_tp_group()), hidden_states, ) attn_tp_all_gather_into_tensor( diff --git a/python/sglang/srt/layers/communicator_nsa_cp.py b/python/sglang/srt/layers/communicator_nsa_cp.py index 243b18e00..2508929f1 100644 --- a/python/sglang/srt/layers/communicator_nsa_cp.py +++ b/python/sglang/srt/layers/communicator_nsa_cp.py @@ -34,6 +34,7 @@ from sglang.srt.layers.communicator import ( from sglang.srt.layers.dp_attention import ( attn_cp_all_gather_into_tensor, attn_cp_reduce_scatter_tensor, + get_attention_cp_group, get_local_dp_buffer, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -154,7 +155,7 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn( if nsa_use_prefill_cp(forward_batch): assert context.attn_dp_size == 1 hidden_states, local_hidden_states = ( - get_local_dp_buffer(), + get_local_dp_buffer(get_attention_cp_group()), hidden_states, ) attn_cp_all_gather_into_tensor( diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index b8d761784..89d238928 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -126,8 +126,8 @@ class _DpGatheredBufferWrapper: cls._global_num_tokens = global_num_tokens @classmethod - def get_global_dp_buffer(cls) -> torch.Tensor: - with use_symmetric_memory(get_tp_group(), disabled=not cls._dp_max_padding): + def get_global_dp_buffer(cls, group: GroupCoordinator) -> torch.Tensor: + with use_symmetric_memory(group, disabled=not cls._dp_max_padding): buffer = torch.empty( (cls._global_dp_buffer_len, cls._hidden_size), dtype=cls._dtype, @@ -136,8 +136,8 @@ class _DpGatheredBufferWrapper: return buffer @classmethod - def get_local_dp_buffer(cls) -> torch.Tensor: - with use_symmetric_memory(get_tp_group(), disabled=not cls._dp_max_padding): + def get_local_dp_buffer(cls, group: GroupCoordinator) -> torch.Tensor: + with use_symmetric_memory(group, disabled=not cls._dp_max_padding): buffer = torch.empty( (cls._local_dp_buffer_len, cls._hidden_size), dtype=cls._dtype, @@ -193,12 +193,12 @@ def set_dp_buffer_len( ) -def get_global_dp_buffer() -> torch.Tensor: - return _DpGatheredBufferWrapper.get_global_dp_buffer() +def get_global_dp_buffer(group: GroupCoordinator) -> torch.Tensor: + return _DpGatheredBufferWrapper.get_global_dp_buffer(group=group) -def get_local_dp_buffer() -> torch.Tensor: - return _DpGatheredBufferWrapper.get_local_dp_buffer() +def get_local_dp_buffer(group: GroupCoordinator) -> torch.Tensor: + return _DpGatheredBufferWrapper.get_local_dp_buffer(group=group) def get_global_dp_buffer_len() -> int: diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py index 7658c28d4..caf99571d 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/standard.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/standard.py @@ -220,7 +220,10 @@ class StandardDispatcher(BaseDispatcher): def combine(self, combine_input: StandardCombineInput) -> torch.Tensor: (hidden_states,) = combine_input if should_use_flashinfer_cutlass_moe_fp4_allgather(): - 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, + ) get_tp_group().reduce_scatterv( global_hidden_states, output=hidden_states,