From 187b85b7f38496653948a2aba546d53c09ada0f3 Mon Sep 17 00:00:00 2001
From: Shangming Cai <caishangming@linux.alibaba.com>
Date: Sat, 21 Jun 2025 00:50:39 +0800
Subject: [PATCH] [PD] Optimize custom mem pool usage and bump mooncake version
 (#7393)

Signed-off-by: Shangming Cai <caishangming@linux.alibaba.com>
---
 .../disaggregation/mooncake/memory_pool.py    | 47 -------------------
 python/sglang/srt/mem_cache/memory_pool.py    | 16 +++----
 scripts/ci_install_dependency.sh              |  2 +-
 3 files changed, 7 insertions(+), 58 deletions(-)
 delete mode 100644 python/sglang/srt/disaggregation/mooncake/memory_pool.py

diff --git a/python/sglang/srt/disaggregation/mooncake/memory_pool.py b/python/sglang/srt/disaggregation/mooncake/memory_pool.py
deleted file mode 100644
index 6e8edaf92..000000000
--- a/python/sglang/srt/disaggregation/mooncake/memory_pool.py
+++ /dev/null
@@ -1,47 +0,0 @@
-import os
-import threading
-from importlib import resources
-from typing import Dict, Final, Optional
-
-import torch
-from torch.cuda.memory import CUDAPluggableAllocator
-
-
-# TODO(shangming): move this class into mooncake's package for more general use cases
-class MooncakeNVLinkAllocator:
-    _instances: Dict[torch.device, CUDAPluggableAllocator] = {}
-    _lock: Final = threading.Lock()
-
-    @classmethod
-    def _get_so_path(cls) -> str:
-        """Dynamically locate hook.so in the mooncake package installation"""
-        try:
-            # Attempt to locate package resource
-            with resources.path("mooncake", "hook.so") as so_path:
-                if so_path.exists():
-                    return str(so_path)
-        except (ImportError, FileNotFoundError, TypeError):
-            pass
-
-        # Fallback strategy: check in package location via import metadata
-        try:
-            import mooncake
-
-            base_path = os.path.dirname(os.path.abspath(mooncake.__file__))
-            so_path = os.path.join(base_path, "hook.so")
-            if os.path.exists(so_path):
-                return so_path
-        except (ImportError, FileNotFoundError, TypeError):
-            raise ImportError(
-                "SGLANG_MOONCAKE_CUSTOM_MEM_POOL require mooncake-transfer-engine >= 0.3.3.post2."
-            )
-
-    @classmethod
-    def get_allocator(cls, device: torch.device) -> CUDAPluggableAllocator:
-        with cls._lock:
-            if device not in cls._instances:
-                so_path = cls._get_so_path()
-                cls._instances[device] = CUDAPluggableAllocator(
-                    so_path, "mc_nvlink_malloc", "mc_nvlink_free"
-                )
-            return cls._instances[device]
diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py
index c01807f1b..b5be2bb1b 100644
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/python/sglang/srt/mem_cache/memory_pool.py
@@ -270,12 +270,10 @@ class MHATokenToKVPool(KVCache):
             "SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false"
         )
         if self.enable_custom_mem_pool:
-            from sglang.srt.disaggregation.mooncake.memory_pool import (
-                MooncakeNVLinkAllocator,
-            )
-
             # TODO(shangming): abstract custom allocator class for more backends
-            allocator = MooncakeNVLinkAllocator.get_allocator(self.device)
+            from mooncake.allocator import NVLinkAllocator
+
+            allocator = NVLinkAllocator.get_allocator(self.device)
             self.custom_mem_pool = torch.cuda.MemPool(allocator.allocator())
         else:
             self.custom_mem_pool = None
@@ -602,12 +600,10 @@ class MLATokenToKVPool(KVCache):
             "SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false"
         )
         if self.enable_custom_mem_pool:
-            from sglang.srt.disaggregation.mooncake.memory_pool import (
-                MooncakeNVLinkAllocator,
-            )
-
             # TODO(shangming): abstract custom allocator class for more backends
-            allocator = MooncakeNVLinkAllocator.get_allocator(self.device)
+            from mooncake.allocator import NVLinkAllocator
+
+            allocator = NVLinkAllocator.get_allocator(self.device)
             self.custom_mem_pool = torch.cuda.MemPool(allocator.allocator())
         else:
             self.custom_mem_pool = None
diff --git a/scripts/ci_install_dependency.sh b/scripts/ci_install_dependency.sh
index 922c886c4..a1808019e 100755
--- a/scripts/ci_install_dependency.sh
+++ b/scripts/ci_install_dependency.sh
@@ -23,7 +23,7 @@ pip install -e "python[dev]"
 pip list
 
 # Install additional dependencies
-pip install mooncake-transfer-engine==0.3.2.post1 nvidia-cuda-nvrtc-cu12
+pip install mooncake-transfer-engine==0.3.4 nvidia-cuda-nvrtc-cu12
 
 # For lmms_evals evaluating MMMU
 git clone --branch v0.3.3 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git