From e379befbac22eb786ea5022b25dd56921d9e212c Mon Sep 17 00:00:00 2001 From: Nicolas Castet <26874160+nvcastet@users.noreply.github.com> Date: Thu, 9 Apr 2026 00:34:58 -0500 Subject: [PATCH] Add symmetric debug mode to print stack trace of comm ops with unregistered tensors (#18569) --- docs/references/environment_variables.md | 1 + .../device_communicators/pynccl_allocator.py | 78 +++++++++++++++++++ .../sglang/srt/distributed/parallel_state.py | 9 +++ python/sglang/srt/environ.py | 1 + 4 files changed, 89 insertions(+) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index b7ac94a71..89854f01c 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -160,6 +160,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_TEST_RETRACT_NO_PREFILL_BS` | When SGLANG_TEST_RETRACT is enabled, no prefill is performed if the batch size exceeds SGLANG_TEST_RETRACT_NO_PREFILL_BS. | `2 ** 31` | | `SGLANG_RECORD_STEP_TIME` | Record step time for profiling | `false` | | `SGLANG_TEST_REQUEST_TIME_STATS` | Test request time statistics | `false` | +| `SGLANG_DEBUG_SYMM_MEM` | Enable debug checks that verify tensors passed to NCCL communication ops are allocated in the symmetric memory pool. Logs warnings (rank 0 only) with stack traces for any tensor not in the pool. | `false` | | `SGLANG_KERNEL_API_LOGLEVEL` | Controls crash-debug kernel API logging. `0` disables logging, `1` logs API names, `3` logs tensor metadata, `5` adds tensor statistics, and `10` also writes pre-call dump snapshots. | `0` | | `SGLANG_KERNEL_API_LOGDEST` | Destination for crash-debug kernel API logs. Use `stdout`, `stderr`, or a file path. `%i` is replaced with the process PID. | `stdout` | | `SGLANG_KERNEL_API_DUMP_DIR` | Output directory for level-10 kernel API input/output dumps. `%i` is replaced with the process PID. | `sglang_kernel_api_dumps` | diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py index 761b3c592..3c57e86ed 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py +++ b/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, + ) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index f0cba2189..e8e14b9f3 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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: diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index c13d58575..797502679 100644 --- a/python/sglang/srt/environ.py +++ b/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)