Add symmetric debug mode to print stack trace of comm ops with unregistered tensors (#18569)

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

@@ -1,11 +1,14 @@
+import logging
 import os
 import tempfile
+import traceback
 from contextlib import nullcontext
 
 import torch
 from torch.cuda.memory import CUDAPluggableAllocator
 
 from sglang.srt.distributed.parallel_state import GroupCoordinator
+from sglang.srt.environ import envs
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils.common import torch_release
 
@@ -211,3 +214,78 @@ def use_symmetric_memory(group_coordinator: GroupCoordinator, disabled: bool = F
         or group_coordinator.world_size == 1
     )
     return SymmetricMemoryContext(group_coordinator) if not disabled else nullcontext()
+
+
+# --- Debug mode for symmetric memory validation ---
+
+_symm_mem_logger = logging.getLogger(__name__)
+_debug_seen_traces: set = set()
+
+
+def is_tensor_in_symmetric_mempool(tensor: torch.Tensor) -> bool:
+    """Check if a tensor's storage is allocated in the NCCL symmetric memory pool."""
+
+    if _mem_pool is None:
+        return False  # Pool not initialized
+
+    data_ptr = tensor.untyped_storage().data_ptr()
+
+    for segment in _mem_pool.snapshot():
+        for block in segment["blocks"]:
+            if block["address"] == data_ptr:
+                return True
+    return False
+
+
+def debug_check_symmetric_mempool(
+    group_coordinator: GroupCoordinator,
+    tensors: dict,
+    op_name: str,
+) -> None:
+    """
+    Debug check: verify that tensors passed to communication ops are allocated
+    in the NCCL symmetric memory pool.
+
+    Enabled by setting SGLANG_DEBUG_SYMM_MEM=1.
+    Only prints warnings on rank 0 and deduplicates identical stack traces.
+
+    Args:
+        tensors: dict mapping argument name to tensor
+                 (e.g. {"input": t1, "output": t2})
+        op_name: name of the communication operation being checked
+    """
+    if not envs.SGLANG_DEBUG_SYMM_MEM.get() or not is_symmetric_memory_enabled():
+        return
+
+    # Only print on rank 0
+    if not group_coordinator.is_first_rank:
+        return
+
+    bad_names = []
+    bad_details = []
+    for name, tensor in tensors.items():
+        if not is_tensor_in_symmetric_mempool(tensor):
+            bad_names.append(name)
+            bad_details.append(
+                f"  - '{name}' (data_ptr=0x{tensor.storage().data_ptr():x}, "
+                f"shape={list(tensor.shape)}, dtype={tensor.dtype})"
+            )
+
+    if bad_names:
+        traces = traceback.format_stack()
+        # Skip autotune stack traces
+        if any("_flashinfer_autotune" in trace for trace in traces):
+            return
+        stack = "".join(traces[:-1])
+        trace_key = f"{op_name}:{','.join(bad_names)}:{stack}"
+        if trace_key not in _debug_seen_traces:
+            _debug_seen_traces.add(trace_key)
+            _symm_mem_logger.warning(
+                "[SymmMem Debug] %s: %d tensor(s) are NOT in the "
+                "NCCL symmetric memory pool:\n%s\n"
+                "Stack trace:\n%s",
+                op_name,
+                len(bad_names),
+                "\n".join(bad_details),
+                stack,
+            )

python/sglang/srt/distributed/parallel_state.py

@@ -334,6 +334,7 @@ class GroupCoordinator:
             PyNcclCommunicator,
         )
         from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+            debug_check_symmetric_mempool,
             is_symmetric_memory_enabled,
             use_symmetric_memory,
         )
@@ -345,6 +346,7 @@ class GroupCoordinator:
         self.is_symmetric_memory_enabled = is_symmetric_memory_enabled
         self.use_symmetric_memory = use_symmetric_memory
         self.is_allocation_symmetric = is_allocation_symmetric
+        self.debug_check_symmetric_mempool = debug_check_symmetric_mempool
         if is_hip():
             from sglang.srt.distributed.device_communicators.quick_all_reduce import (
                 QuickAllReduce,
@@ -577,6 +579,7 @@ class GroupCoordinator:
             return self.npu_communicator.all_reduce(input_)
 
         if self.pynccl_comm is not None and self.is_symmetric_memory_enabled():
+            self.debug_check_symmetric_mempool(self, {"input": input_}, "all_reduce")
             with self.pynccl_comm.change_state(enable=True):
                 self.pynccl_comm.all_reduce(input_)
                 return input_
@@ -718,6 +721,9 @@ class GroupCoordinator:
         if pynccl_comm is not None and (
             not pynccl_comm.disabled or self.is_symmetric_memory_enabled()
         ):
+            self.debug_check_symmetric_mempool(
+                self, {"output": output, "input": input}, "reduce_scatter_tensor"
+            )
             with pynccl_comm.change_state(enable=True):
                 pynccl_comm.reduce_scatter(output, input)
         else:
@@ -779,6 +785,9 @@ class GroupCoordinator:
         if pynccl_comm is not None and (
             not pynccl_comm.disabled or self.is_symmetric_memory_enabled()
         ):
+            self.debug_check_symmetric_mempool(
+                self, {"output": output}, "all_gather_into_tensor"
+            )
             with pynccl_comm.change_state(enable=True):
                 pynccl_comm.all_gather(output, input)
         else:

python/sglang/srt/environ.py

@@ -524,6 +524,7 @@ class Envs:
 
     # Symmetric Memory
     SGLANG_SYMM_MEM_PREALLOC_GB_SIZE = EnvInt(-1)
+    SGLANG_DEBUG_SYMM_MEM = EnvBool(False)
 
     # Aiter
     SGLANG_USE_AITER_FP8_PER_TOKEN = EnvBool(False)