diff --git a/python/sglang/srt/disaggregation/mooncake/utils.py b/python/sglang/srt/disaggregation/mooncake/utils.py index 279cf194d..b56c92a31 100644 --- a/python/sglang/srt/disaggregation/mooncake/utils.py +++ b/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}" ) diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 9a1094bba..753dcdb22 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/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):