chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[PD] feat: refactor custom mem pool and add barex pd support (#12332)

Browse Source
This commit is contained in:
Teng Ma
2025-11-09 18:51:34 +08:00
committed by GitHub
parent 877cb52840
commit b8ac4fcb51
5 changed files with 162 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/srt/disaggregation/mooncake/conn.py
View File

@@ -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 = []

110
python/sglang/srt/disaggregation/mooncake/utils.py Normal file
View File

@@ -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

2
python/sglang/srt/environ.py
View File

@@ -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

50
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -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,

33
python/sglang/srt/mem_cache/utils.py
View File

@@ -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