Register allgather/reducescatter buffers with symm memory

Author: nvcastet

benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py

@@ -2,8 +2,8 @@
 import argparse
 import json
 import time
-from datetime import datetime
 from contextlib import nullcontext
+from datetime import datetime
 from typing import Any, Dict, List, Tuple, TypedDict
 
 import ray

python/sglang/srt/distributed/device_communicators/pynccl_allocator.py

@@ -5,7 +5,6 @@
 from torch.cuda.memory import CUDAPluggableAllocator
 
 from sglang.srt.distributed.parallel_state import GroupCoordinator
-from sglang.srt.managers.schedule_batch import global_server_args_dict
 
 nccl_allocator_source = """
 #include <nccl.h>
@@ -28,13 +27,21 @@
 _allocator = None
 _mem_pool = None
 _registered_base_addrs = set()
+_registered_tensor_addrs = set()
 _graph_pool_id = None
 
 
 def is_symmetric_memory_enabled():
+    # Import here to avoid circular import
+    from sglang.srt.managers.schedule_batch import global_server_args_dict
+
     return global_server_args_dict["enable_symm_mem"]
 
 
+def is_symmetric_memory_tensor(tensor: torch.Tensor):
+    return tensor.untyped_storage().data_ptr() in _registered_tensor_addrs
+
+
 def set_graph_pool_id(graph_pool_id):
     global _graph_pool_id
     _graph_pool_id = graph_pool_id
@@ -64,8 +71,17 @@ def get_nccl_mem_pool():
 
 
 class use_symmetric_memory:
-    def __init__(self, group_coordinator: GroupCoordinator):
-        if not is_symmetric_memory_enabled():
+    def __init__(
+        self,
+        group_coordinator: GroupCoordinator,
+        disabled: bool = False,
+    ):
+        self.disabled = (
+            disabled
+            or not is_symmetric_memory_enabled()
+            or group_coordinator.world_size == 1
+        )
+        if self.disabled:
             self.group_coordinator = None
             self._mem_pool_ctx = None
             self.is_graph_capture = None
@@ -79,7 +95,7 @@ def __init__(self, group_coordinator: GroupCoordinator):
             self.pre_2_8_0 = version.parse(torch.__version__) < version.parse("2.8.0")
 
     def __enter__(self):
-        if not is_symmetric_memory_enabled():
+        if self.disabled:
             return self
         assert (
             self.group_coordinator.pynccl_comm is not None
@@ -102,12 +118,13 @@ def __enter__(self):
         return self
 
     def tag(self, tensor: torch.Tensor):
-        if not is_symmetric_memory_enabled():
+        if self.disabled:
             return
-        tensor.symmetric_memory = True
+        global _registered_tensor_addrs
+        _registered_tensor_addrs.add(tensor.untyped_storage().data_ptr())
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if not is_symmetric_memory_enabled():
+        if self.disabled:
             return
         global _registered_base_addrs
         self._mem_pool_ctx.__exit__(exc_type, exc_val, exc_tb)

python/sglang/srt/distributed/parallel_state.py

@@ -270,7 +270,13 @@ def __init__(
         from sglang.srt.distributed.device_communicators.pynccl import (
             PyNcclCommunicator,
         )
+        from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+            is_symmetric_memory_tensor,
+            use_symmetric_memory,
+        )
 
+        self.is_symmetric_memory_tensor = is_symmetric_memory_tensor
+        self.use_symmetric_memory = use_symmetric_memory
         if is_hip():
             from sglang.srt.distributed.device_communicators.quick_all_reduce import (
                 QuickAllReduce,
@@ -499,11 +505,7 @@ def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
         if self.npu_communicator is not None and not self.npu_communicator.disabled:
             return self.npu_communicator.all_reduce(input_)
 
-        if (
-            self.pynccl_comm is not None
-            and hasattr(input_, "symmetric_memory")
-            and input_.symmetric_memory
-        ):
+        if self.pynccl_comm is not None and self.is_symmetric_memory_tensor(input_):
             with self.pynccl_comm.change_state(
                 enable=True, stream=torch.cuda.current_stream()
             ):
@@ -569,9 +571,23 @@ def reduce_scatter_tensor(
         self,
         output: torch.Tensor,
         input: torch.Tensor,
-    ) -> None:
-        # TODO(ch-wan): support other backends
-        torch.distributed.reduce_scatter_tensor(output, input, group=self.device_group)
+    ) -> torch.Tensor:
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and (
+            not pynccl_comm.disabled
+            or (
+                self.is_symmetric_memory_tensor(output)
+                and self.is_symmetric_memory_tensor(input)
+            )
+        ):
+            with pynccl_comm.change_state(
+                enable=True, stream=torch.cuda.current_stream()
+            ):
+                pynccl_comm.reduce_scatter(output, input)
+        else:
+            torch.distributed.reduce_scatter_tensor(
+                output, input, group=self.device_group
+            )
         return output
 
     def reduce_scatter(
@@ -618,8 +634,17 @@ def reduce_scatterv(
 
     def _all_gather_into_tensor(self, output: torch.Tensor, input: torch.Tensor):
         pynccl_comm = self.pynccl_comm
-        if pynccl_comm is not None and not pynccl_comm.disabled:
-            pynccl_comm.all_gather(output, input)
+        if pynccl_comm is not None and (
+            not pynccl_comm.disabled
+            or (
+                self.is_symmetric_memory_tensor(output)
+                and self.is_symmetric_memory_tensor(input)
+            )
+        ):
+            with pynccl_comm.change_state(
+                enable=True, stream=torch.cuda.current_stream()
+            ):
+                pynccl_comm.all_gather(output, input)
         else:
             torch.distributed.all_gather_into_tensor(
                 output, input, group=self.device_group
@@ -681,9 +706,11 @@ def all_gather(
         # torch.compile . see https://github.com/pytorch/pytorch/issues/138795
         output_size = (input_size[0] * world_size,) + input_size[1:]
         # Allocate output tensor.
-        output_tensor = torch.empty(
-            output_size, dtype=input_.dtype, device=input_.device
-        )
+        with self.use_symmetric_memory(self) as sm:
+            output_tensor = torch.empty(
+                output_size, dtype=input_.dtype, device=input_.device
+            )
+            sm.tag(output_tensor)
 
         # All-gather.
         if input_.is_cpu and is_shm_available(

python/sglang/srt/layers/communicator.py

@@ -21,8 +21,12 @@
 
 from sglang.srt.distributed import (
     get_tensor_model_parallel_world_size,
+    get_tp_group,
     tensor_model_parallel_all_reduce,
 )
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    use_symmetric_memory,
+)
 from sglang.srt.layers.dp_attention import (
     attn_tp_all_gather_into_tensor,
     attn_tp_reduce_scatter_tensor,
@@ -430,7 +434,13 @@ def _gather_hidden_states_and_residual(
             use_layer_norm_before_gather = context.attn_tp_size == 1
             if use_layer_norm_before_gather and hidden_states.shape[0] != 0:
                 residual = hidden_states
-                hidden_states = layernorm(hidden_states)
+                with use_symmetric_memory(
+                    get_tp_group(),
+                    disabled=not forward_batch.dp_padding_mode.is_max_len(),
+                ) as sm:
+                    hidden_states = layernorm(hidden_states)
+                    sm.tag(hidden_states)
+
             hidden_states, local_hidden_states = (
                 get_global_dp_buffer(),
                 hidden_states,

python/sglang/srt/layers/dp_attention.py

@@ -17,6 +17,9 @@
     get_tp_group,
     tensor_model_parallel_all_reduce,
 )
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    use_symmetric_memory,
+)
 
 if TYPE_CHECKING:
     from sglang.srt.configs.model_config import ModelConfig
@@ -72,6 +75,7 @@ class _DpGatheredBufferWrapper:
     _device: torch.device
     _global_dp_buffer_len: int
     _local_dp_buffer_len: int
+    _is_max_padding: bool
     _global_num_tokens: Optional[List[int]]
 
     @classmethod
@@ -85,27 +89,37 @@ def set_dp_buffer_len(
         cls,
         global_dp_buffer_len: int,
         local_dp_buffer_len: int,
+        is_max_padding: bool,
         global_num_tokens: Optional[List[int]] = None,
     ):
         cls._global_dp_buffer_len = global_dp_buffer_len
         cls._local_dp_buffer_len = local_dp_buffer_len
+        cls._is_max_padding = is_max_padding
         cls._global_num_tokens = global_num_tokens
 
     @classmethod
     def get_global_dp_buffer(cls) -> torch.Tensor:
-        return torch.empty(
-            (cls._global_dp_buffer_len, cls._hidden_size),
-            dtype=cls._dtype,
-            device=cls._device,
-        )
+        with use_symmetric_memory(get_tp_group()) as sm:
+            buffer = torch.empty(
+                (cls._global_dp_buffer_len, cls._hidden_size),
+                dtype=cls._dtype,
+                device=cls._device,
+            )
+            sm.tag(buffer)
+        return buffer
 
     @classmethod
     def get_local_dp_buffer(cls) -> torch.Tensor:
-        return torch.empty(
-            (cls._local_dp_buffer_len, cls._hidden_size),
-            dtype=cls._dtype,
-            device=cls._device,
-        )
+        with use_symmetric_memory(
+            get_tp_group(), disabled=not cls._is_max_padding
+        ) as sm:
+            buffer = torch.empty(
+                (cls._local_dp_buffer_len, cls._hidden_size),
+                dtype=cls._dtype,
+                device=cls._device,
+            )
+            sm.tag(buffer)
+        return buffer
 
     @classmethod
     def get_global_dp_buffer_len(cls) -> int:
@@ -119,14 +133,19 @@ def get_local_dp_buffer_len(cls) -> int:
     def get_dp_global_num_tokens(cls) -> List[int]:
         return cls._global_num_tokens
 
+    @classmethod
+    def is_max_padding(cls) -> bool:
+        return cls._is_max_padding
+
 
 def set_dp_buffer_len(
     global_dp_buffer_len: int,
     local_dp_buffer_len: int,
+    is_max_padding: bool,
     global_num_tokens: Optional[List[int]] = None,
 ):
     _DpGatheredBufferWrapper.set_dp_buffer_len(
-        global_dp_buffer_len, local_dp_buffer_len, global_num_tokens
+        global_dp_buffer_len, local_dp_buffer_len, is_max_padding, global_num_tokens
     )
 
 
@@ -150,6 +169,10 @@ def get_dp_global_num_tokens() -> List[int]:
     return _DpGatheredBufferWrapper.get_dp_global_num_tokens()
 
 
+def is_max_padding() -> bool:
+    return _DpGatheredBufferWrapper.is_max_padding()
+
+
 def compute_dp_attention_world_info(enable_dp_attention, tp_rank, tp_size, dp_size):
     if not enable_dp_attention:
         return tp_rank, tp_size, 0
@@ -408,7 +431,10 @@ def _dp_gather_via_all_gather(
     scattered_local_tokens = local_tokens.tensor_split(get_attention_tp_size())[
         get_attention_tp_rank()
     ]
-    get_attention_tp_group().reduce_scatter_tensor(scattered_local_tokens, local_tokens)
+    if get_attention_tp_size() > 1:
+        get_attention_tp_group().reduce_scatter_tensor(
+            scattered_local_tokens, local_tokens
+        )
     get_tp_group().all_gather_into_tensor(global_tokens, scattered_local_tokens)
 
 
@@ -467,7 +493,7 @@ def dp_scatter(
 
 
 def dp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
-    if get_tensor_model_parallel_world_size() == get_attention_dp_size():
+    if get_attention_tp_size() == 1:
         get_tp_group().reduce_scatter_tensor(output, input)
     else:
         scattered_local_tokens = input.tensor_split(

python/sglang/srt/layers/linear.py

@@ -1301,7 +1301,7 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
             # It does not support additional parameters.
             param.load_row_parallel_weight(loaded_weight)
 
-    def forward(self, input_, skip_all_reduce=False):
+    def forward(self, input_, skip_all_reduce=False, disable_symmetric_memory=True):
         if self.input_is_parallel:
             input_parallel = input_
         else:
@@ -1315,7 +1315,9 @@ def forward(self, input_, skip_all_reduce=False):
         # 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(parallel_state.get_tp_group()) as sm:
+        with use_symmetric_memory(
+            parallel_state.get_tp_group(), disabled=disable_symmetric_memory
+        ) as sm:
             output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
             sm.tag(output_parallel)
 

python/sglang/srt/layers/logits_processor.py

@@ -194,6 +194,7 @@ def compute_dp_attention_metadata(self):
         set_dp_buffer_len(
             self.global_dp_buffer_len,
             self.dp_local_num_tokens,
+            False,
             self.global_num_tokens_for_logprob_cpu,
         )
 

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -11,12 +11,8 @@
     get_moe_expert_parallel_world_size,
     get_moe_tensor_parallel_rank,
     get_moe_tensor_parallel_world_size,
-    get_tp_group,
     tensor_model_parallel_all_reduce,
 )
-from sglang.srt.distributed.device_communicators.pynccl_allocator import (
-    use_symmetric_memory,
-)
 from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
 from sglang.srt.layers.moe import (
     MoeRunnerConfig,
@@ -812,15 +808,12 @@ def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
                 raise NotImplementedError()
 
         # Matrix multiply.
-        with use_symmetric_memory(get_tp_group()) as sm:
-
-            final_hidden_states = self.quant_method.apply(
-                layer=self,
-                x=hidden_states,
-                topk_output=topk_output,
-                moe_runner_config=self.moe_runner_config,
-            )
-            sm.tag(final_hidden_states)
+        final_hidden_states = self.quant_method.apply(
+            layer=self,
+            x=hidden_states,
+            topk_output=topk_output,
+            moe_runner_config=self.moe_runner_config,
+        )
 
         final_hidden_states = final_hidden_states[
             ..., :origin_hidden_states_dim