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

[PD] feat: support mooncake intra-node nvlink kv transfer (#17866)

Browse Source 
Co-authored-by: 百麒 <yaozhong.lyz@alibaba-inc.com>
Co-authored-by: Teng Ma <teng-ma@linux.alibaba.com>
This commit is contained in:
Lewis
2026-02-03 17:47:52 +08:00
committed by GitHub
parent cc69ac9e7a
commit a45647bce1
2 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/disaggregation/mooncake/utils.py
View File

@@ -23,7 +23,7 @@ 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"]
SUPPORTED_MOONCAKE_CUSTOM_MEM_POOL_TYPES = ["NVLINK", "BAREX", "INTRA_NVLINK"]
def init_mooncake_custom_mem_pool(
@@ -55,6 +55,8 @@ def init_mooncake_custom_mem_pool(
from mooncake.allocator import BarexAllocator
allocator = BarexAllocator.get_allocator(device)
elif custom_mem_pool_type == "INTRA_NODE_NVLINK":
return False, None, None
else:
# This should not happen due to the enable_custom_mem_pool check above
raise ValueError(

3
python/sglang/srt/disaggregation/utils.py
View File

@@ -11,6 +11,7 @@ import numpy as np
import torch
import torch.distributed as dist
from sglang.srt.environ import envs
from sglang.srt.utils import is_npu
if TYPE_CHECKING:
@@ -97,6 +98,8 @@ class MetadataBuffers:
elif self.custom_mem_pool:
# TODO(shangming): Fix me (use 'cuda') when nvlink_transport of Mooncake is bug-free
device = "cpu"
elif envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() == "INTRA_NODE_NVLINK":
device = "cuda"
with (
torch.cuda.use_mem_pool(self.custom_mem_pool)
if self.custom_mem_pool