From b8ac4fcb5191b7683172635f14dc425b956ed263 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Sun, 9 Nov 2025 18:51:34 +0800 Subject: [PATCH] [PD] feat: refactor custom mem pool and add barex pd support (#12332) --- .../srt/disaggregation/mooncake/conn.py | 16 ++- .../srt/disaggregation/mooncake/utils.py | 110 ++++++++++++++++++ python/sglang/srt/environ.py | 2 +- python/sglang/srt/mem_cache/memory_pool.py | 50 ++------ python/sglang/srt/mem_cache/utils.py | 33 ++++++ 5 files changed, 162 insertions(+), 49 deletions(-) create mode 100644 python/sglang/srt/disaggregation/mooncake/utils.py diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 2326a647a..0e1f844a5 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -28,14 +28,12 @@ from sglang.srt.disaggregation.common.utils import ( group_concurrent_contiguous, ) from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine +from sglang.srt.disaggregation.mooncake.utils import ( + check_mooncake_custom_mem_pool_enabled, +) from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import ( - format_tcp_address, - get_bool_env_var, - get_int_env_var, - is_valid_ipv6_address, -) +from sglang.srt.utils import format_tcp_address, get_int_env_var, is_valid_ipv6_address logger = logging.getLogger(__name__) @@ -201,8 +199,8 @@ class MooncakeKVManager(CommonKVManager): "SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT", 300 ) - self.enable_custom_mem_pool = get_bool_env_var( - "SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false" + self.enable_custom_mem_pool, self.custom_mem_pool_type = ( + check_mooncake_custom_mem_pool_enabled() ) elif self.disaggregation_mode == DisaggregationMode.DECODE: self.heartbeat_failures = {} @@ -549,7 +547,7 @@ class MooncakeKVManager(CommonKVManager): dst_aux_ptrs: list[int], ): # TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free - if self.enable_custom_mem_pool: + if self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK": return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs) transfer_blocks = [] diff --git a/python/sglang/srt/disaggregation/mooncake/utils.py b/python/sglang/srt/disaggregation/mooncake/utils.py new file mode 100644 index 000000000..767276139 --- /dev/null +++ b/python/sglang/srt/disaggregation/mooncake/utils.py @@ -0,0 +1,110 @@ +# Copyright 2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Mooncake-specific utilities for custom memory pool management.""" + +import logging +from typing import Any, Optional, Tuple + +import torch + +from sglang.srt.environ import envs + +logger = logging.getLogger(__name__) + +# Global constants for custom memory pool types +SUPPORTED_MOONCAKE_CUSTOM_MEM_POOL_TYPES = ["NVLINK", "BAREX"] + + +def init_mooncake_custom_mem_pool( + device: str, +) -> Tuple[bool, Optional[Any], Optional[str]]: + """ + Initialize custom memory pool based on environment variable. + + Args: + device: The device to allocate memory on + + Returns: + Tuple of (enable_custom_mem_pool, custom_mem_pool, custom_mem_pool_type) + """ + enable_custom_mem_pool, custom_mem_pool_type = ( + check_mooncake_custom_mem_pool_enabled() + ) + + custom_mem_pool = None + + if enable_custom_mem_pool: + try: + # TODO(shangming): abstract custom allocator class for more backends + if custom_mem_pool_type == "NVLINK": + from mooncake.allocator import NVLinkAllocator + + allocator = NVLinkAllocator.get_allocator(device) + elif custom_mem_pool_type == "BAREX": + from mooncake.allocator import BarexAllocator + + allocator = BarexAllocator.get_allocator(device) + else: + # This should not happen due to the enable_custom_mem_pool check above + raise ValueError( + f"Unsupported custom mem pool type: {custom_mem_pool_type}" + ) + + custom_mem_pool = torch.cuda.MemPool(allocator.allocator()) + logger.debug( + f"Initialized custom memory pool: {custom_mem_pool_type} on device {device}" + ) + except ImportError as e: + logger.warning( + f"Failed to import mooncake allocator for {custom_mem_pool_type}: {e}. " + f"Falling back to default memory pool." + ) + enable_custom_mem_pool = False + custom_mem_pool = None + custom_mem_pool_type = None + except Exception as e: + logger.error( + f"Failed to initialize custom memory pool {custom_mem_pool_type}: {e}. " + f"Falling back to default memory pool." + ) + enable_custom_mem_pool = False + custom_mem_pool = None + custom_mem_pool_type = None + else: + return False, None, None + + return enable_custom_mem_pool, custom_mem_pool, custom_mem_pool_type + + +def check_mooncake_custom_mem_pool_enabled() -> Tuple[bool, Optional[str]]: + """ + Check if custom memory pool is enabled without importing allocators. + + Returns: + Tuple of (enable_custom_mem_pool, custom_mem_pool_type) + """ + custom_mem_pool_type = envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() + + if custom_mem_pool_type is not None: + # Handle boolean True as NVLINK + if custom_mem_pool_type.lower() == "true": + custom_mem_pool_type = "NVLINK" + enable_custom_mem_pool = ( + custom_mem_pool_type in SUPPORTED_MOONCAKE_CUSTOM_MEM_POOL_TYPES + ) + else: + enable_custom_mem_pool = False + custom_mem_pool_type = None + + return enable_custom_mem_pool, custom_mem_pool_type diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 057e3b4f3..41ff73fc2 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -181,7 +181,7 @@ class Envs: SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None) # Mooncake KV Transfer - SGLANG_MOONCAKE_CUSTOM_MEM_POOL = EnvBool(False) + SGLANG_MOONCAKE_CUSTOM_MEM_POOL = EnvStr(None) ENABLE_ASCEND_TRANSFER_WITH_MOONCAKE = EnvBool(False) # AMD & ROCm diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 39e691145..244c0acd9 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -45,16 +45,11 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.utils import ( get_mla_kv_buffer_triton, + maybe_init_custom_mem_pool, set_mla_kv_buffer_triton, set_mla_kv_scale_buffer_triton, ) -from sglang.srt.utils import ( - get_bool_env_var, - is_cuda, - is_float4_e2m1fn_x2, - is_npu, - next_power_of_2, -) +from sglang.srt.utils import is_cuda, is_float4_e2m1fn_x2, is_npu, next_power_of_2 if TYPE_CHECKING: from sglang.srt.managers.cache_controller import LayerDoneCounter @@ -162,19 +157,13 @@ class MambaPool: ssm_dtype = cache_params.dtype.temporal num_mamba_layers = len(cache_params.layers) + self.size = size self.device = device - # for disagg with nvlink - self.enable_custom_mem_pool = get_bool_env_var( - "SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false" - ) - if self.enable_custom_mem_pool: - # TODO(shangming): abstract custom allocator class for more backends - 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 + # for disagg with nvlink + self.enable_custom_mem_pool, self.custom_mem_pool, _ = ( + maybe_init_custom_mem_pool(device=self.device) + ) self.is_kda_cache = isinstance(cache_params, KimiLinearCacheParams) with ( @@ -271,7 +260,6 @@ class MambaPool: f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, " f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB " ) - self.size = size self.free_slots = torch.arange( self.size, dtype=torch.int64, device=self.device ) @@ -484,17 +472,9 @@ class KVCache(abc.ABC): self.layer_transfer_counter = None # for disagg with nvlink - self.enable_custom_mem_pool = get_bool_env_var( - "SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false" + self.enable_custom_mem_pool, self.custom_mem_pool, _ = ( + maybe_init_custom_mem_pool(device=self.device) ) - if self.enable_custom_mem_pool: - # TODO(shangming): abstract custom allocator class for more backends - 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 def _finalize_allocation_log(self, num_tokens: int): """Common logging and mem_usage computation for KV cache allocation. @@ -1062,17 +1042,9 @@ class SWAKVPool(KVCache): assert not enable_kvcache_transpose # for disagg with nvlink - self.enable_custom_mem_pool = get_bool_env_var( - "SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false" + self.enable_custom_mem_pool, self.custom_mem_pool, _ = ( + maybe_init_custom_mem_pool(device=self.device) ) - if self.enable_custom_mem_pool: - # TODO(shangming): abstract custom allocator class for more backends - 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 self.swa_kv_pool = token_to_kv_pool_class( size=size_swa, diff --git a/python/sglang/srt/mem_cache/utils.py b/python/sglang/srt/mem_cache/utils.py index 03f5ec0a8..8f0e6d37d 100644 --- a/python/sglang/srt/mem_cache/utils.py +++ b/python/sglang/srt/mem_cache/utils.py @@ -13,10 +13,14 @@ # ============================================================================== """Common utilities.""" +from typing import Any, Optional, Tuple + import torch import triton import triton.language as tl +from sglang.srt.environ import envs + @triton.jit def set_mla_kv_buffer_kernel( @@ -208,3 +212,32 @@ def get_mla_kv_buffer_triton( nope_dim, rope_dim, ) + + +def maybe_init_custom_mem_pool( + device: str, +) -> Tuple[bool, Optional[Any], Optional[str]]: + """ + Initialize custom memory pool based on environment variable. + + This function can be modified to support more features that require a custom memory pool. + + Args: + device: The device to allocate memory on + + Returns: + Tuple of (enable_custom_mem_pool, custom_mem_pool, custom_mem_pool_type) + """ + enable_custom_mem_pool = ( + True if envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() is not None else False + ) + + if enable_custom_mem_pool: + # Currently, only mooncake requires a custom mem pool for MNNVL/Barex PD disaggregation + from sglang.srt.disaggregation.mooncake.utils import ( + init_mooncake_custom_mem_pool, + ) + + return init_mooncake_custom_mem_pool(device) + else: + return False, None, None