diff --git a/python/sglang/srt/disaggregation/mooncake/utils.py b/python/sglang/srt/disaggregation/mooncake/utils.py index 767276139..63e7bd29d 100644 --- a/python/sglang/srt/disaggregation/mooncake/utils.py +++ b/python/sglang/srt/disaggregation/mooncake/utils.py @@ -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( diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 816acc871..b3786f881 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -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