[Fix] Run FlashInfer autotune on non-default stream for NCCL 2.29+ compatibility (#18987)

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,

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):