diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py index a97824c23..1aca62200 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py +++ b/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 = ( diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 121def4c6..bae5cbdd7 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/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: