From 113f6f080eb8fcb5924b1ae252eb66282743be15 Mon Sep 17 00:00:00 2001 From: zijiexia <37504505+zijiexia@users.noreply.github.com> Date: Sat, 19 Sep 2026 19:34:58 -0700 Subject: [PATCH] [PD] Enter the custom mem pool once when allocating DCP pack buffers (#40284) Co-authored-by: Claude Opus 5 Co-authored-by: kpham-sgl --- python/sglang/srt/disaggregation/mooncake/utils.py | 6 +++++- python/sglang/srt/disaggregation/nixl/conn.py | 13 ++++++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) 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):