From ef4b3c0e967d84b96b54e9ac4146050252bec8df Mon Sep 17 00:00:00 2001 From: EkiRui <374389385@qq.com> Date: Mon, 29 Dec 2025 08:28:50 +0800 Subject: [PATCH] Add host tensor allocator for memory_pool_host and support Mooncake standalone storage (#14873) Signed-off-by: Xingrui Yi Co-authored-by: Teng Ma --- python/sglang/srt/environ.py | 2 + python/sglang/srt/mem_cache/hiradix_cache.py | 2 + .../sglang/srt/mem_cache/memory_pool_host.py | 57 +++++++- .../srt/mem_cache/storage/backend_factory.py | 2 +- .../storage/mooncake_store/mooncake_store.py | 128 ++++++++++++++---- 5 files changed, 164 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index eb2868bb1..4fb5d8f88 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -251,6 +251,7 @@ class Envs: # Mooncake Store SGLANG_HICACHE_MOONCAKE_CONFIG_PATH = EnvStr(None) MOONCAKE_MASTER = EnvStr(None) + MOONCAKE_CLIENT = EnvStr(None) MOONCAKE_LOCAL_HOSTNAME = EnvStr("localhost") MOONCAKE_TE_META_DATA_SERVER = EnvStr("P2PHANDSHAKE") MOONCAKE_GLOBAL_SEGMENT_SIZE = EnvStr("4gb") @@ -258,6 +259,7 @@ class Envs: MOONCAKE_DEVICE = EnvStr("") MOONCAKE_MASTER_METRICS_PORT = EnvInt(9003) MOONCAKE_CHECK_SERVER = EnvBool(False) + MOONCAKE_STANDALONE_STORAGE = EnvBool(False) # AMD & ROCm SGLANG_USE_AITER = EnvBool(False) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index dd6ddc290..a38f1ac5f 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -52,6 +52,7 @@ class HiRadixCache(RadixCache): server_args.hicache_size, self.page_size, server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, ) elif isinstance(self.kv_cache, MLATokenToKVPool): self.token_to_kv_pool_host = MLATokenToKVPoolHost( @@ -60,6 +61,7 @@ class HiRadixCache(RadixCache): server_args.hicache_size, self.page_size, server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, ) else: raise ValueError(f"HiRadixCache only supports MHA and MLA yet") diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 08ff838ee..46394158f 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -52,17 +52,51 @@ def synchronized(func): return wrapper +class HostTensorAllocator(abc.ABC): + def __init__(self): + """Initialize the HostTensorAllocator.""" + self.dtype = None + self.dims = None + + def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor: + """Allocate a tensor of given dims and dtype on the memory.""" + self.dtype = dtype + self.dims = dims + tensor = torch.empty(dims, dtype=dtype, device=device) + return tensor + + +def get_allocator_from_storage(allocator_type): + if allocator_type == "mooncake": + try: + from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import ( + MooncakeHostTensorAllocator, + ) + + return MooncakeHostTensorAllocator() + except ImportError: + logger.warning( + "Mooncake's tensor allocator requires mooncake >= 0.3.8. " + "Please upgrade Mooncake by 'pip install mooncake --upgrade'. " + "Fallback to use default allocator." + ) + return HostTensorAllocator() + else: + return HostTensorAllocator() + + def alloc_with_host_register( dims, dtype: torch.dtype, device: str, pin_memory: bool, + allocator: HostTensorAllocator, ) -> torch.Tensor: """ Allocate tensor and register host memory with cudaHostRegister. CudaHostRegister only applies when pin_memory=True. """ - buffer = torch.empty(dims, dtype=dtype, device=device) + buffer = allocator.allocate(dims, dtype=dtype, device=device) if pin_memory: torch.cuda.cudart().cudaHostRegister( buffer.data_ptr(), buffer.numel() * buffer.element_size(), 0 @@ -75,6 +109,7 @@ def alloc_with_pin_memory( dtype: torch.dtype, device: str, pin_memory: bool, + allocator: None, ) -> torch.Tensor: """ Allocate tensor using PyTorch's built-in pin_memory flag. @@ -102,12 +137,14 @@ class HostKVCache(abc.ABC): layout: str, pin_memory: bool, device: str, + allocator_type: str = "default", ): self.device_pool = device_pool self.page_size = page_size self.layout = layout self.pin_memory = pin_memory self.device = device + self.allocator = get_allocator_from_storage(allocator_type) self.dtype = device_pool.store_dtype self.size_per_token = self.get_size_per_token() @@ -239,6 +276,7 @@ class MHATokenToKVPoolHost(HostKVCache): layout: str, pin_memory: bool = True, device: str = "cpu", + allocator_type: str = "default", ): super().__init__( device_pool, @@ -248,6 +286,7 @@ class MHATokenToKVPoolHost(HostKVCache): layout, pin_memory, device, + allocator_type, ) self.element_dim = self.device_pool.head_num * self.device_pool.head_dim self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel( @@ -307,7 +346,11 @@ class MHATokenToKVPoolHost(HostKVCache): alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] buffer = alloc_func( - dims, dtype=self.dtype, device=self.device, pin_memory=self.pin_memory + dims, + dtype=self.dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, ) return buffer @@ -645,6 +688,7 @@ class MLATokenToKVPoolHost(HostKVCache): layout: str, pin_memory: bool = True, device: str = "cpu", + allocator_type: str = "default", ): super().__init__( device_pool, @@ -654,6 +698,7 @@ class MLATokenToKVPoolHost(HostKVCache): layout, pin_memory, device, + allocator_type, ) self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] self.data_ptrs = torch.tensor( @@ -715,12 +760,14 @@ class MLATokenToKVPoolHost(HostKVCache): dtype=self.dtype, device=self.device, pin_memory=self.pin_memory, + allocator=self.allocator, ) self.v_buffer = alloc_func( (*base_dims, self.qk_rope_head_dim), dtype=self.dtype, device=self.device, pin_memory=self.pin_memory, + allocator=self.allocator, ) # Return k_buffer to preserve original kv_buffer and data_refs init logic, # though Ascend doesn't use these parameters. @@ -734,7 +781,11 @@ class MLATokenToKVPoolHost(HostKVCache): alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] buffer = alloc_func( - dims, dtype=self.dtype, device=self.device, pin_memory=self.pin_memory + dims, + dtype=self.dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, ) return buffer diff --git a/python/sglang/srt/mem_cache/storage/backend_factory.py b/python/sglang/srt/mem_cache/storage/backend_factory.py index 4b195c863..eaa5a3e18 100644 --- a/python/sglang/srt/mem_cache/storage/backend_factory.py +++ b/python/sglang/srt/mem_cache/storage/backend_factory.py @@ -163,7 +163,7 @@ class StorageBackendFactory: elif backend_name == "nixl": return backend_class(storage_config) elif backend_name == "mooncake": - backend = backend_class(storage_config) + backend = backend_class(storage_config, mem_pool_host) return backend elif backend_name == "aibrix": backend = backend_class(storage_config, mem_pool_host) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index acc472369..b015db46c 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -1,3 +1,4 @@ +import ctypes import json import logging import os @@ -15,7 +16,7 @@ from sglang.srt.mem_cache.hicache_storage import ( HiCacheStorageConfig, HiCacheStorageExtraInfo, ) -from sglang.srt.mem_cache.memory_pool_host import HostKVCache +from sglang.srt.mem_cache.memory_pool_host import HostKVCache, HostTensorAllocator DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024 # 16 MB SETUP_TIMEOUT = 600 # 10min @@ -23,6 +24,41 @@ SETUP_TIMEOUT = 600 # 10min logger = logging.getLogger(__name__) +class MooncakeHostTensorAllocator(HostTensorAllocator): + def __init__(self): + super().__init__() + from mooncake.store import MooncakeHostMemAllocator + + self.allocator = MooncakeHostMemAllocator() + self.ptr = None + + def allocate( + self, dims: tuple, dtype: torch.dtype, device: str = "cpu" + ) -> torch.Tensor: + """ + Allocates memory using MooncakeHostMemAllocator and wraps it in a PyTorch tensor. + """ + self.dims = dims + self.dtype = dtype + size = 1 + for d in dims: + size *= d + size *= torch.tensor([], dtype=self.dtype).element_size() + ptr_int = self.allocator.alloc(size) + self.ptr = ptr_int + c_type = ctypes.c_byte * size + c_array = c_type.from_address(ptr_int) + + tensor = torch.frombuffer(c_array, dtype=torch.uint8, count=size) + + if dtype != torch.uint8: + element_size = torch.tensor([], dtype=dtype).element_size() + assert size % element_size == 0, "Size must be divisible by element size" + tensor = tensor.view(dtype) + + return tensor.view(dims) + + def _parse_global_segment_size(value) -> int: if isinstance(value, int): return value @@ -49,6 +85,8 @@ class MooncakeStoreConfig: master_server_address: str master_metrics_port: int check_server: bool + standalone_storage: bool + client_server_address: str @staticmethod def from_file() -> "MooncakeStoreConfig": @@ -64,8 +102,13 @@ class MooncakeStoreConfig: except Exception as e: raise RuntimeError(f"Failed to load config from {file_path}: {str(e)}") - if "master_server_address" not in config: - raise ValueError("master_server_address is required in config file") + if ( + "master_server_address" not in config + and "client_server_address" not in config + ): + raise ValueError( + "Either master_server_address or client_server_address is required in config file" + ) return MooncakeStoreConfig( local_hostname=config.get( @@ -81,11 +124,19 @@ class MooncakeStoreConfig: ), protocol=config.get("protocol", envs.MOONCAKE_PROTOCOL.default), device_name=config.get("device_name", envs.MOONCAKE_DEVICE.default), - master_server_address=config.get("master_server_address"), + master_server_address=config.get( + "master_server_address", envs.MOONCAKE_MASTER.default + ), master_metrics_port=config.get( "master_metrics_port", envs.MOONCAKE_MASTER_METRICS_PORT.default ), check_server=config.get("check_server", envs.MOONCAKE_CHECK_SERVER.default), + standalone_storage=config.get( + "standalone_storage", envs.MOONCAKE_STANDALONE_STORAGE.default + ), + client_server_address=config.get( + "client_server_address", envs.MOONCAKE_CLIENT.default + ), ) @staticmethod @@ -97,8 +148,10 @@ class MooncakeStoreConfig: export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE" """ # other required environment variables... - if not envs.MOONCAKE_MASTER.is_set(): - raise ValueError("The environment variable 'MOONCAKE_MASTER' is not set.") + if not envs.MOONCAKE_MASTER.is_set() and not envs.MOONCAKE_CLIENT.is_set(): + raise ValueError( + "Either the environment variable 'MOONCAKE_MASTER' or 'MOONCAKE_CLIENT' is not set." + ) # Special handling for local_hostname: try MOONCAKE_LOCAL_HOSTNAME first, # then fall back to LOCAL_HOSTNAME if not set. @@ -121,13 +174,20 @@ class MooncakeStoreConfig: master_server_address=envs.MOONCAKE_MASTER.get(), master_metrics_port=envs.MOONCAKE_MASTER_METRICS_PORT.get(), check_server=envs.MOONCAKE_CHECK_SERVER.get(), + standalone_storage=envs.MOONCAKE_STANDALONE_STORAGE.get(), + client_server_address=envs.MOONCAKE_CLIENT.get(), ) @staticmethod def load_from_extra_config(extra_config: dict) -> "MooncakeStoreConfig": """Load config from extra_config dictionary.""" - if "master_server_address" not in extra_config: - raise ValueError("master_server_address is required in extra_config") + if ( + "master_server_address" not in extra_config + and "client_server_address" not in extra_config + ): + raise ValueError( + "Either master_server_address or client_server_address is required in extra_config" + ) return MooncakeStoreConfig( local_hostname=extra_config.get( @@ -143,19 +203,29 @@ class MooncakeStoreConfig: ), protocol=extra_config.get("protocol", envs.MOONCAKE_PROTOCOL.default), device_name=extra_config.get("device_name", envs.MOONCAKE_DEVICE.default), - master_server_address=extra_config["master_server_address"], + master_server_address=extra_config.get( + "master_server_address", envs.MOONCAKE_MASTER.default + ), master_metrics_port=extra_config.get( "master_metrics_port", envs.MOONCAKE_MASTER_METRICS_PORT.default ), check_server=extra_config.get( "check_server", envs.MOONCAKE_CHECK_SERVER.default ), + standalone_storage=extra_config.get( + "standalone_storage", envs.MOONCAKE_STANDALONE_STORAGE.default + ), + client_server_address=extra_config.get( + "client_server_address", envs.MOONCAKE_CLIENT.default + ), ) class MooncakeStore(HiCacheStorage): - def __init__(self, storage_config: HiCacheStorageConfig = None): + def __init__( + self, storage_config: HiCacheStorageConfig = None, mem_pool: HostKVCache = None + ): try: from mooncake.store import MooncakeDistributedStore except ImportError as e: @@ -174,9 +244,9 @@ class MooncakeStore(HiCacheStorage): else None ) # Load configuration with master_server_address prioritized from extra_config if available - if ( - extra_config is not None - and extra_config.get("master_server_address") is not None + if extra_config is not None and ( + extra_config.get("master_server_address") is not None + or extra_config.get("client_server_address") is not None ): # Load from extra_config self.config = MooncakeStoreConfig.load_from_extra_config(extra_config) @@ -226,16 +296,28 @@ class MooncakeStore(HiCacheStorage): f"Failed to parse device_name as JSON: {device_name}" ) device_name = "" - - ret_code = self.store.setup( - self.config.local_hostname, - self.config.metadata_server, - per_tp_global_segment_size, - DEFAULT_LOCAL_BUFFER_SIZE, # Zero copy interface does not need local buffer - self.config.protocol, - device_name, - self.config.master_server_address, - ) + if self.config.standalone_storage: + if not isinstance(mem_pool.allocator, MooncakeHostTensorAllocator): + raise RuntimeError( + "MooncakeStore with standalone_storage=True requires MooncakeHostTensorAllocator. " + "Please set standalone_storage=False " + "or upgrade Mooncake by 'pip install mooncake --upgrade'." + ) + ret_code = self.store.setup_dummy( + mem_pool.size * mem_pool.size_per_token, + DEFAULT_LOCAL_BUFFER_SIZE, # Zero copy interface does not need local buffer + self.config.client_server_address, + ) + else: + ret_code = self.store.setup( + self.config.local_hostname, + self.config.metadata_server, + per_tp_global_segment_size, + DEFAULT_LOCAL_BUFFER_SIZE, # Zero copy interface does not need local buffer + self.config.protocol, + device_name, + self.config.master_server_address, + ) if ret_code: raise RuntimeError( f"Failed to setup Mooncake store, error code: {ret_code}"