[PD] Enter the custom mem pool once when allocating DCP pack buffers (#40284)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: kpham-sgl <khoa.pham@radixark.ai>

python/sglang/srt/disaggregation/mooncake/utils.py

@@ -63,7 +63,11 @@ def init_mooncake_custom_mem_pool(
                     f"Unsupported custom mem pool type: {custom_mem_pool_type}"
                 )
 
-            custom_mem_pool = torch.cuda.MemPool(allocator.allocator())
+            # MemPool binds to the current device; on a non-main thread (e.g.
+            # PD bootstrap) that is device 0, so ranks on other GPUs hit the
+            # CUDACachingAllocator use_count assert in use_mem_pool().
+            with torch.cuda.device(device):
+                custom_mem_pool = torch.cuda.MemPool(allocator.allocator())
             logger.debug(
                 f"Initialized custom memory pool: {custom_mem_pool_type} on device {device}"
             )

python/sglang/srt/disaggregation/nixl/conn.py

@@ -3077,7 +3077,18 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager):
                     logger.debug(f"{room=} is bootstrapped")
                     self.update_status(room, KVPoll.WaitingForInput)
 
-        threading.Thread(target=bootstrap_thread).start()
+        def bootstrap_thread_guarded():
+            try:
+                bootstrap_thread()
+            except Exception:
+                logger.exception(
+                    "prefill bootstrap_thread died on engine_rank=%s; requests to "
+                    "this rank will time out in KVPoll.Bootstrapping",
+                    self.kv_args.engine_rank,
+                )
+                raise
+
+        threading.Thread(target=bootstrap_thread_guarded).start()
 
 
 class NixlKVSender(CommonKVSender):