diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index e8e14b9f3..447715230 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -1724,6 +1724,7 @@ def initialize_model_parallel( moe_data_model_parallel_size: int = 1, backend: Optional[str] = None, duplicate_tp_group: bool = False, + enable_symm_mem: bool = False, ) -> None: """ Initialize model parallel groups. @@ -1884,7 +1885,7 @@ def initialize_model_parallel( group_ranks, get_world_group().local_rank, backend, - use_pynccl=SYNC_TOKEN_IDS_ACROSS_TP, + use_pynccl=SYNC_TOKEN_IDS_ACROSS_TP or enable_symm_mem, use_mscclpp_allreduce=False, use_custom_allreduce=False, use_torch_symm_mem_allreduce=False, diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 7af9eb004..74d9a7271 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -1503,9 +1503,13 @@ class RowParallelLinear(LinearBase): # Only fuse bias add into GEMM for rank 0 (this ensures that # bias will not get added more than once in TP>1 case) bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias - with use_symmetric_memory( - get_tp_group(), disabled=not is_allocation_symmetric() - ): + if self.use_dp_attention_reduce: + symm_ctx = use_symmetric_memory(get_attention_tp_group()) + else: + symm_ctx = use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ) + with symm_ctx: output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_) if self.reduce_results and self.tp_size > 1 and not skip_all_reduce: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 76d9ded64..e4948bc47 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1027,6 +1027,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): attention_context_model_parallel_size=self.attn_cp_size, moe_data_model_parallel_size=self.moe_dp_size, duplicate_tp_group=self.server_args.enable_pdmux, + enable_symm_mem=self.server_args.enable_symm_mem, ) initialize_dp_attention( server_args=self.server_args, diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py index b8ad74015..447f57eaf 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -73,6 +73,7 @@ class LlamaMLP(nn.Module): reduce_results: bool = True, tp_rank: Optional[int] = None, tp_size: Optional[int] = None, + use_dp_attention_reduce: bool = False, ) -> None: super().__init__() self.gate_up_proj = MergedColumnParallelLinear( @@ -93,6 +94,7 @@ class LlamaMLP(nn.Module): reduce_results=reduce_results, tp_rank=tp_rank, tp_size=tp_size, + use_dp_attention_reduce=use_dp_attention_reduce, ) if hidden_act != "silu": raise ValueError(