diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py index 7dfc08c29..761b3c592 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py @@ -104,7 +104,17 @@ def get_nccl_mem_pool(): if _mem_pool is None: import torch.utils.cpp_extension - out_dir = tempfile.gettempdir() + out_dir = os.path.join(tempfile.gettempdir(), "symm_allocator") + os.makedirs(out_dir, exist_ok=True) + # Make sure to clean up leftover pytorch lock files + # from previous runs and synchronize across processes + # right after + try: + os.remove(os.path.join(out_dir, "lock")) + except FileNotFoundError: + pass + torch.distributed.barrier() + nccl_allocator_libname = "nccl_allocator" torch.utils.cpp_extension.load_inline( name=nccl_allocator_libname, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index d451c614a..275775a73 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1862,11 +1862,15 @@ class ModelRunner(ModelRunnerKVCacheMixin): logger.info("Running FlashInfer autotune...") - self._dummy_run( - batch_size=self.req_to_token_pool.size, - run_ctx=autotune(), - ) - + # Run warmup on the non-default stream to avoid NCCL 2.29+ cudaMemcpyBatchAsync + # calls on default stream (unsupported by CUDA) when --enable-symm-mem is used. + self.forward_stream.wait_stream(torch.cuda.current_stream()) + with torch.get_device_module(self.device).stream(self.forward_stream): + with torch.inference_mode(), autotune(): + self._dummy_run( + batch_size=self.req_to_token_pool.size, run_ctx=autotune() + ) + torch.cuda.current_stream().wait_stream(self.forward_stream) logger.info("FlashInfer autotune completed.") def _dummy_run(self, batch_size: int, run_ctx=None):