Fix hang in deepgemm compilation with symmetric memory enabled (#12715)

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

@@ -62,6 +62,7 @@ _allocator = None
 _mem_pool = None
 _graph_pool_id = None
 _cur_device = None
+_active_symmetric_memory_context = None
 
 
 def is_symmetric_memory_enabled():
@@ -73,6 +74,19 @@ def set_graph_pool_id(graph_pool_id):
     _graph_pool_id = graph_pool_id
 
 
+def disable_symmetric_memory_context():
+    if _active_symmetric_memory_context is None:
+        return None
+    saved_context = _active_symmetric_memory_context
+    saved_context.__exit__(None, None, None)
+    return saved_context
+
+
+def restore_symmetric_memory_context(saved_context):
+    if saved_context is not None:
+        saved_context.__enter__()
+
+
 def get_nccl_mem_pool():
     global _allocator, _mem_pool, _cur_device
     if _mem_pool is None:
@@ -114,6 +128,7 @@ class SymmetricMemoryContext:
         self.group_coordinator = group_coordinator
         self._mem_pool_ctx = torch.cuda.use_mem_pool(get_nccl_mem_pool())
         self.is_graph_capture = torch.cuda.is_current_stream_capturing()
+        self.exited = False
 
     def __enter__(self):
         assert (
@@ -132,12 +147,20 @@ class SymmetricMemoryContext:
                     _cur_device, _graph_pool_id
                 )
 
+        if self.exited:
+            # mempool ctx (@contextlib.contextmanager) is not re-entrant
+            self._mem_pool_ctx = torch.cuda.use_mem_pool(get_nccl_mem_pool())
+            self.exited = False
         self._mem_pool_ctx.__enter__()
 
         # Set the env var to pass this argument to the C functions.
         os.environ["SGLANG_TMP_NCCL_COMM_VALUE"] = str(
             self.group_coordinator.pynccl_comm.comm.value
         )
+
+        global _active_symmetric_memory_context
+        _active_symmetric_memory_context = self
+
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -151,6 +174,11 @@ class SymmetricMemoryContext:
             else:
                 torch._C._cuda_beginAllocateToPool(_cur_device, _graph_pool_id)
 
+        global _active_symmetric_memory_context
+        _active_symmetric_memory_context = None
+
+        self.exited = True
+
 
 def use_symmetric_memory(group_coordinator: GroupCoordinator, disabled: bool = False):
     disabled = (

python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py

@@ -7,6 +7,10 @@ from typing import Dict, List, Tuple
 import torch
 from tqdm import tqdm
 
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    disable_symmetric_memory_context,
+    restore_symmetric_memory_context,
+)
 from sglang.srt.environ import envs
 from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM
 from sglang.srt.server_args import ServerArgs
@@ -120,25 +124,32 @@ def _compile_deep_gemm_one_type_all(
     num_groups: int,
     m_list: List[int],
 ) -> None:
-    if kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG:
-        m_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout()
-        m_list = sorted(list(set(m for m in m_list if m % m_alignment == 0)))
+    # Symmetric memory allocation performs a collective operation across all the GPUs.
+    # Temporary disable symmetric memory during compilation since it only runs on the first rank.
+    saved_context = disable_symmetric_memory_context()
+    try:
+        if kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG:
+            m_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout()
+            m_list = sorted(list(set(m for m in m_list if m % m_alignment == 0)))
 
-    executor = _BaseWarmupExecutor.create(
-        kernel_type, max_m=max(m_list), n=n, k=k, num_groups=num_groups
-    )
+        executor = _BaseWarmupExecutor.create(
+            kernel_type, max_m=max(m_list), n=n, k=k, num_groups=num_groups
+        )
 
-    old_compile_mode = deep_gemm.get_compile_mode()
-    deep_gemm.set_compile_mode(1)
-    # TODO can use multi thread
-    for m in tqdm(m_list, desc=f"DeepGEMM warmup"):
-        executor.execute(m=m)
-    deep_gemm.set_compile_mode(old_compile_mode)
+        old_compile_mode = deep_gemm.get_compile_mode()
+        deep_gemm.set_compile_mode(1)
+        # TODO can use multi thread
+        for m in tqdm(m_list, desc=f"DeepGEMM warmup"):
+            executor.execute(m=m)
+        deep_gemm.set_compile_mode(old_compile_mode)
 
-    # clean up input buffers
-    torch.cuda.current_stream().synchronize()
-    del executor
-    torch.cuda.empty_cache()
+        # clean up input buffers
+        torch.cuda.current_stream().synchronize()
+        del executor
+        torch.cuda.empty_cache()
+    finally:
+        # Restore symmetric memory context
+        restore_symmetric_memory_context(saved_context)
 
 
 class _BaseWarmupExecutor: