feat(hicache): Add shared memory allocator for host KV cache (#29326)
Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
co-authored by
Jimmy Shong
Zhangheng
parent
b83041c3cc
commit
ebcb74abd4
@@ -18,6 +18,7 @@ from sglang.srt.mem_cache.memory_pool import (
|
|||||||
MLATokenToKVPool,
|
MLATokenToKVPool,
|
||||||
ReqToTokenPool,
|
ReqToTokenPool,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.mem_cache.pool_host.common import get_allocator_type
|
||||||
from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
|
from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
|
||||||
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
|
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
|
||||||
from sglang.srt.server_args import ServerArgs
|
from sglang.srt.server_args import ServerArgs
|
||||||
@@ -54,6 +55,8 @@ class DecodeKVCacheOffloadManager:
|
|||||||
self.page_size, (env_stride // self.page_size) * self.page_size
|
self.page_size, (env_stride // self.page_size) * self.page_size
|
||||||
)
|
)
|
||||||
kv_cache = self.token_to_kv_pool_allocator.get_kvcache()
|
kv_cache = self.token_to_kv_pool_allocator.get_kvcache()
|
||||||
|
allocator_type = get_allocator_type(server_args)
|
||||||
|
|
||||||
if isinstance(kv_cache, MHATokenToKVPool):
|
if isinstance(kv_cache, MHATokenToKVPool):
|
||||||
self.decode_host_mem_pool = get_mha_host_pool_cls(kv_cache)(
|
self.decode_host_mem_pool = get_mha_host_pool_cls(kv_cache)(
|
||||||
kv_cache,
|
kv_cache,
|
||||||
@@ -61,6 +64,7 @@ class DecodeKVCacheOffloadManager:
|
|||||||
server_args.hicache_size,
|
server_args.hicache_size,
|
||||||
self.page_size,
|
self.page_size,
|
||||||
server_args.hicache_mem_layout,
|
server_args.hicache_mem_layout,
|
||||||
|
allocator_type=allocator_type,
|
||||||
)
|
)
|
||||||
elif isinstance(kv_cache, MLATokenToKVPool):
|
elif isinstance(kv_cache, MLATokenToKVPool):
|
||||||
self.decode_host_mem_pool = MLATokenToKVPoolHost(
|
self.decode_host_mem_pool = MLATokenToKVPoolHost(
|
||||||
@@ -69,6 +73,7 @@ class DecodeKVCacheOffloadManager:
|
|||||||
server_args.hicache_size,
|
server_args.hicache_size,
|
||||||
self.page_size,
|
self.page_size,
|
||||||
server_args.hicache_mem_layout,
|
server_args.hicache_mem_layout,
|
||||||
|
allocator_type=allocator_type,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
raise ValueError("Unsupported KV cache type for decode offload")
|
raise ValueError("Unsupported KV cache type for decode offload")
|
||||||
|
|||||||
@@ -48,6 +48,7 @@ from sglang.srt.mem_cache.memory_pool import (
|
|||||||
MiniMaxSparseKVPool,
|
MiniMaxSparseKVPool,
|
||||||
MLATokenToKVPool,
|
MLATokenToKVPool,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.mem_cache.pool_host.common import get_allocator_type
|
||||||
from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
|
from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
|
||||||
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
|
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
|
||||||
from sglang.srt.mem_cache.radix_cache import (
|
from sglang.srt.mem_cache.radix_cache import (
|
||||||
@@ -80,6 +81,8 @@ class HiRadixCache(RadixCache):
|
|||||||
self.page_size = params.page_size
|
self.page_size = params.page_size
|
||||||
self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
|
self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
|
||||||
|
|
||||||
|
allocator_type = get_allocator_type(server_args)
|
||||||
|
|
||||||
if isinstance(self.kv_cache, MHATokenToKVPool):
|
if isinstance(self.kv_cache, MHATokenToKVPool):
|
||||||
self.token_to_kv_pool_host = get_mha_host_pool_cls(self.kv_cache)(
|
self.token_to_kv_pool_host = get_mha_host_pool_cls(self.kv_cache)(
|
||||||
self.kv_cache,
|
self.kv_cache,
|
||||||
@@ -87,7 +90,7 @@ class HiRadixCache(RadixCache):
|
|||||||
server_args.hicache_size,
|
server_args.hicache_size,
|
||||||
self.page_size,
|
self.page_size,
|
||||||
server_args.hicache_mem_layout,
|
server_args.hicache_mem_layout,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=allocator_type,
|
||||||
)
|
)
|
||||||
elif isinstance(self.kv_cache, DSATokenToKVPool):
|
elif isinstance(self.kv_cache, DSATokenToKVPool):
|
||||||
# Filled by attach_hybrid_dsa_pool_to_hiradix_cache after storage extra_config is parsed.
|
# Filled by attach_hybrid_dsa_pool_to_hiradix_cache after storage extra_config is parsed.
|
||||||
@@ -102,7 +105,7 @@ class HiRadixCache(RadixCache):
|
|||||||
server_args.hicache_size,
|
server_args.hicache_size,
|
||||||
self.page_size,
|
self.page_size,
|
||||||
server_args.hicache_mem_layout,
|
server_args.hicache_mem_layout,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=allocator_type,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
raise ValueError("HiRadixCache only supports MHA, MLA, DSA, and MSA models")
|
raise ValueError("HiRadixCache only supports MHA, MLA, DSA, and MSA models")
|
||||||
|
|||||||
@@ -21,6 +21,7 @@ from sglang.srt.mem_cache.memory_pool_host import (
|
|||||||
MambaPoolHost,
|
MambaPoolHost,
|
||||||
PoolEntry,
|
PoolEntry,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.mem_cache.pool_host.common import get_allocator_type
|
||||||
from sglang.srt.mem_cache.pool_host.mha import (
|
from sglang.srt.mem_cache.pool_host.mha import (
|
||||||
MHATokenToKOnlyPoolHost,
|
MHATokenToKOnlyPoolHost,
|
||||||
get_mha_host_pool_cls,
|
get_mha_host_pool_cls,
|
||||||
@@ -40,6 +41,10 @@ if TYPE_CHECKING:
|
|||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _get_allocator_type(server_args: ServerArgs) -> str:
|
||||||
|
return get_allocator_type(server_args)
|
||||||
|
|
||||||
|
|
||||||
def _make_layer_mapper(
|
def _make_layer_mapper(
|
||||||
layer_mapping: dict[int, int],
|
layer_mapping: dict[int, int],
|
||||||
transfer_layer_num: int,
|
transfer_layer_num: int,
|
||||||
@@ -72,7 +77,7 @@ def build_kv_host_pool(
|
|||||||
server_args.hicache_size,
|
server_args.hicache_size,
|
||||||
page_size,
|
page_size,
|
||||||
server_args.hicache_mem_layout,
|
server_args.hicache_mem_layout,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=_get_allocator_type(server_args),
|
||||||
**kwargs,
|
**kwargs,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -353,7 +358,7 @@ def build_deepseek_v4_hicache_stack(
|
|||||||
num_host_pages=swa_num_host_pages,
|
num_host_pages=swa_num_host_pages,
|
||||||
slot_page_size=kvcache.swa_page_size,
|
slot_page_size=kvcache.swa_page_size,
|
||||||
layout=server_args.hicache_mem_layout,
|
layout=server_args.hicache_mem_layout,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=_get_allocator_type(server_args),
|
||||||
)
|
)
|
||||||
swa_attn_allocator = params.token_to_kv_pool_allocator.swa_attn_allocator
|
swa_attn_allocator = params.token_to_kv_pool_allocator.swa_attn_allocator
|
||||||
entries.append(
|
entries.append(
|
||||||
@@ -379,7 +384,7 @@ def build_deepseek_v4_hicache_stack(
|
|||||||
num_host_pages=num_host_pages,
|
num_host_pages=num_host_pages,
|
||||||
slot_page_size=page_size,
|
slot_page_size=page_size,
|
||||||
layout=server_args.hicache_mem_layout,
|
layout=server_args.hicache_mem_layout,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=_get_allocator_type(server_args),
|
||||||
)
|
)
|
||||||
c4_indexer_host_pool = DeepSeekV4PagedHostPool(
|
c4_indexer_host_pool = DeepSeekV4PagedHostPool(
|
||||||
pool_name=str(PoolName.DEEPSEEK_V4_C4_INDEXER),
|
pool_name=str(PoolName.DEEPSEEK_V4_C4_INDEXER),
|
||||||
@@ -391,7 +396,7 @@ def build_deepseek_v4_hicache_stack(
|
|||||||
num_host_pages=num_host_pages,
|
num_host_pages=num_host_pages,
|
||||||
slot_page_size=page_size,
|
slot_page_size=page_size,
|
||||||
layout=server_args.hicache_mem_layout,
|
layout=server_args.hicache_mem_layout,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=_get_allocator_type(server_args),
|
||||||
)
|
)
|
||||||
entries.extend(
|
entries.extend(
|
||||||
[
|
[
|
||||||
@@ -422,7 +427,7 @@ def build_deepseek_v4_hicache_stack(
|
|||||||
num_host_pages=swa_num_host_pages,
|
num_host_pages=swa_num_host_pages,
|
||||||
swa_page_size=kvcache.swa_page_size,
|
swa_page_size=kvcache.swa_page_size,
|
||||||
layout=server_args.hicache_mem_layout,
|
layout=server_args.hicache_mem_layout,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=_get_allocator_type(server_args),
|
||||||
)
|
)
|
||||||
c4_indexer_state_host_pool = DeepSeekV4StateHostPool(
|
c4_indexer_state_host_pool = DeepSeekV4StateHostPool(
|
||||||
pool_name=str(PoolName.DEEPSEEK_V4_C4_INDEXER_STATE),
|
pool_name=str(PoolName.DEEPSEEK_V4_C4_INDEXER_STATE),
|
||||||
@@ -433,7 +438,7 @@ def build_deepseek_v4_hicache_stack(
|
|||||||
num_host_pages=swa_num_host_pages,
|
num_host_pages=swa_num_host_pages,
|
||||||
swa_page_size=kvcache.swa_page_size,
|
swa_page_size=kvcache.swa_page_size,
|
||||||
layout=server_args.hicache_mem_layout,
|
layout=server_args.hicache_mem_layout,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=_get_allocator_type(server_args),
|
||||||
)
|
)
|
||||||
entries.extend(
|
entries.extend(
|
||||||
[
|
[
|
||||||
@@ -465,7 +470,7 @@ def build_deepseek_v4_hicache_stack(
|
|||||||
num_host_pages=num_host_pages,
|
num_host_pages=num_host_pages,
|
||||||
slot_page_size=page_size,
|
slot_page_size=page_size,
|
||||||
layout=server_args.hicache_mem_layout,
|
layout=server_args.hicache_mem_layout,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=_get_allocator_type(server_args),
|
||||||
)
|
)
|
||||||
# C128 state pool is intentionally not registered with hicache.
|
# C128 state pool is intentionally not registered with hicache.
|
||||||
# page_size=256 % 128 == 0, so state pool is not consumed on load.
|
# page_size=256 % 128 == 0, so state pool is not consumed on load.
|
||||||
@@ -533,7 +538,7 @@ def build_hybrid_mamba_stack(
|
|||||||
mamba_pool,
|
mamba_pool,
|
||||||
server_args.hicache_ratio,
|
server_args.hicache_ratio,
|
||||||
server_args.hicache_size,
|
server_args.hicache_size,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=_get_allocator_type(server_args),
|
||||||
layout=server_args.hicache_mem_layout,
|
layout=server_args.hicache_mem_layout,
|
||||||
)
|
)
|
||||||
entries = [
|
entries = [
|
||||||
@@ -1107,7 +1112,7 @@ class _DsaStrategy(StackStrategy):
|
|||||||
full_kv_pool,
|
full_kv_pool,
|
||||||
kv_host_pool,
|
kv_host_pool,
|
||||||
server_args.hicache_mem_layout,
|
server_args.hicache_mem_layout,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=_get_allocator_type(server_args),
|
||||||
),
|
),
|
||||||
prefetch_threshold=prefetch_threshold,
|
prefetch_threshold=prefetch_threshold,
|
||||||
model_name=model_name,
|
model_name=model_name,
|
||||||
@@ -1545,7 +1550,7 @@ def attach_hybrid_dsa_pool_to_hiradix_cache(
|
|||||||
kv,
|
kv,
|
||||||
kv_host_pool,
|
kv_host_pool,
|
||||||
server_args.hicache_mem_layout,
|
server_args.hicache_mem_layout,
|
||||||
allocator_type=server_args.hicache_storage_backend,
|
allocator_type=_get_allocator_type(server_args),
|
||||||
),
|
),
|
||||||
model_name=server_args.served_model_name,
|
model_name=server_args.served_model_name,
|
||||||
storage_backend_extra_config=extra_config,
|
storage_backend_extra_config=extra_config,
|
||||||
|
|||||||
@@ -139,6 +139,7 @@ class HostKVCache(abc.ABC):
|
|||||||
)
|
)
|
||||||
|
|
||||||
self.kv_buffer = self.init_kv_buffer()
|
self.kv_buffer = self.init_kv_buffer()
|
||||||
|
self.fd = getattr(self.allocator, "fd", None)
|
||||||
|
|
||||||
# A lock for synchronized operations on memory allocation and state transitions.
|
# A lock for synchronized operations on memory allocation and state transitions.
|
||||||
self.lock = threading.RLock()
|
self.lock = threading.RLock()
|
||||||
|
|||||||
@@ -1,11 +1,13 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
import logging
|
import logging
|
||||||
|
import os
|
||||||
from collections import defaultdict
|
from collections import defaultdict
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.srt.mem_cache.mmap_allocator import alloc_mmap
|
from sglang.srt.mem_cache.storage.mmap import alloc_mmap
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -25,6 +27,43 @@ class HostTensorAllocator:
|
|||||||
return alloc_mmap(dims, dtype)
|
return alloc_mmap(dims, dtype)
|
||||||
|
|
||||||
|
|
||||||
|
class ShmHostTensorAllocator(HostTensorAllocator):
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.fds = []
|
||||||
|
self.mms = []
|
||||||
|
|
||||||
|
@property
|
||||||
|
def fd(self):
|
||||||
|
return self.fds[0] if self.fds else None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def mm(self):
|
||||||
|
return self.mms[0] if self.mms else None
|
||||||
|
|
||||||
|
def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor:
|
||||||
|
assert (
|
||||||
|
device == "cpu"
|
||||||
|
), f"ShmHostTensorAllocator only supports CPU allocations; got device={device!r}"
|
||||||
|
self.dtype = dtype
|
||||||
|
self.dims = dims
|
||||||
|
from sglang.srt.mem_cache.storage.mmap import alloc_shm
|
||||||
|
|
||||||
|
tensor, fd, mm = alloc_shm(dims, dtype)
|
||||||
|
self.fds.append(fd)
|
||||||
|
self.mms.append(mm)
|
||||||
|
return tensor
|
||||||
|
|
||||||
|
def __del__(self):
|
||||||
|
for fd in getattr(self, "fds", []):
|
||||||
|
if fd is not None:
|
||||||
|
try:
|
||||||
|
os.close(fd)
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
self.fds = []
|
||||||
|
|
||||||
|
|
||||||
def get_allocator_from_storage(allocator_type):
|
def get_allocator_from_storage(allocator_type):
|
||||||
if allocator_type == "mooncake":
|
if allocator_type == "mooncake":
|
||||||
try:
|
try:
|
||||||
@@ -54,10 +93,30 @@ def get_allocator_from_storage(allocator_type):
|
|||||||
exc,
|
exc,
|
||||||
)
|
)
|
||||||
return HostTensorAllocator()
|
return HostTensorAllocator()
|
||||||
|
elif allocator_type == "shm":
|
||||||
|
return ShmHostTensorAllocator()
|
||||||
else:
|
else:
|
||||||
return HostTensorAllocator()
|
return HostTensorAllocator()
|
||||||
|
|
||||||
|
|
||||||
|
def get_allocator_type(server_args) -> str:
|
||||||
|
backend = getattr(server_args, "hicache_storage_backend", None)
|
||||||
|
if backend == "shm":
|
||||||
|
return "shm"
|
||||||
|
if backend == "dynamic":
|
||||||
|
extra_config_str = getattr(
|
||||||
|
server_args, "hicache_storage_backend_extra_config", None
|
||||||
|
)
|
||||||
|
if extra_config_str:
|
||||||
|
try:
|
||||||
|
config = json.loads(extra_config_str)
|
||||||
|
if config.get("allocator") == "shm":
|
||||||
|
return "shm"
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return backend or "default"
|
||||||
|
|
||||||
|
|
||||||
def _cuda_host_register(buffer: torch.Tensor) -> None:
|
def _cuda_host_register(buffer: torch.Tensor) -> None:
|
||||||
cudart = torch.cuda.cudart()
|
cudart = torch.cuda.cudart()
|
||||||
n_bytes = buffer.numel() * buffer.element_size()
|
n_bytes = buffer.numel() * buffer.element_size()
|
||||||
|
|||||||
@@ -187,6 +187,8 @@ class StorageBackendFactory:
|
|||||||
return backend_class(storage_config, mem_pool_host)
|
return backend_class(storage_config, mem_pool_host)
|
||||||
elif backend_name == "mori":
|
elif backend_name == "mori":
|
||||||
return backend_class(storage_config, mem_pool_host)
|
return backend_class(storage_config, mem_pool_host)
|
||||||
|
elif backend_name == "shm":
|
||||||
|
return backend_class(storage_config, mem_pool_host)
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"Unknown built-in backend: {backend_name}")
|
raise ValueError(f"Unknown built-in backend: {backend_name}")
|
||||||
|
|
||||||
@@ -237,3 +239,9 @@ StorageBackendFactory.register_backend(
|
|||||||
"sglang.srt.mem_cache.storage.umbp.umbp_store",
|
"sglang.srt.mem_cache.storage.umbp.umbp_store",
|
||||||
"UMBPStore",
|
"UMBPStore",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
StorageBackendFactory.register_backend(
|
||||||
|
"shm",
|
||||||
|
"sglang.srt.mem_cache.storage.shm",
|
||||||
|
"HiCacheShm",
|
||||||
|
)
|
||||||
|
|||||||
@@ -0,0 +1,11 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to SGLang project
|
||||||
|
|
||||||
|
"""Mmap allocator storage backend helpers for SGLang HiCache."""
|
||||||
|
|
||||||
|
from .mmap_allocator import alloc_mmap, alloc_shm
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"alloc_mmap",
|
||||||
|
"alloc_shm",
|
||||||
|
]
|
||||||
+74
@@ -4,6 +4,7 @@ import logging
|
|||||||
import math
|
import math
|
||||||
import mmap
|
import mmap
|
||||||
import os
|
import os
|
||||||
|
import uuid
|
||||||
import weakref
|
import weakref
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
@@ -39,6 +40,7 @@ _MAP_HUGETLB = 0x40000
|
|||||||
_MAP_HUGE_2MB = 21 << 26 # 0x1400000
|
_MAP_HUGE_2MB = 21 << 26 # 0x1400000
|
||||||
_MAP_HUGE_1GB = 30 << 26 # 0x78000000
|
_MAP_HUGE_1GB = 30 << 26 # 0x78000000
|
||||||
_MAP_FAILED = ctypes.c_void_p(-1).value
|
_MAP_FAILED = ctypes.c_void_p(-1).value
|
||||||
|
_MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23)
|
||||||
|
|
||||||
|
|
||||||
def _alloc_hugepage(n_bytes: int, alloc_bytes: int, extra_flags: int) -> ctypes.Array:
|
def _alloc_hugepage(n_bytes: int, alloc_bytes: int, extra_flags: int) -> ctypes.Array:
|
||||||
@@ -124,4 +126,76 @@ def alloc_mmap(dims: tuple, dtype: torch.dtype) -> torch.Tensor:
|
|||||||
flags=mmap.MAP_SHARED | mmap.MAP_ANONYMOUS | _MAP_POPULATE,
|
flags=mmap.MAP_SHARED | mmap.MAP_ANONYMOUS | _MAP_POPULATE,
|
||||||
prot=mmap.PROT_READ | mmap.PROT_WRITE,
|
prot=mmap.PROT_READ | mmap.PROT_WRITE,
|
||||||
)
|
)
|
||||||
|
try:
|
||||||
|
# MADV_POPULATE_WRITE guarantees pages are populated and writable,
|
||||||
|
# throwing an error on failure (e.g. out of memory).
|
||||||
|
mm.madvise(_MADV_POPULATE_WRITE)
|
||||||
|
except OSError:
|
||||||
|
# Fall back to MAP_POPULATE if MADV_POPULATE_WRITE is not supported (<5.14 kernel).
|
||||||
|
pass
|
||||||
return torch.frombuffer(mm, dtype=dtype, count=math.prod(dims)).reshape(dims)
|
return torch.frombuffer(mm, dtype=dtype, count=math.prod(dims)).reshape(dims)
|
||||||
|
|
||||||
|
|
||||||
|
def alloc_shm(dims: tuple, dtype: torch.dtype) -> tuple[torch.Tensor, int, mmap.mmap]:
|
||||||
|
"""Allocate a host tensor via shared memory (/dev/shm).
|
||||||
|
|
||||||
|
Returns a tuple of (tensor, fd, mm).
|
||||||
|
The caller is responsible for keeping the fd open if they need to share it,
|
||||||
|
and closing it when they are done.
|
||||||
|
"""
|
||||||
|
hugepage_size = (envs.SGLANG_HUGEPAGE_SIZE.get() or "").strip().upper()
|
||||||
|
n_bytes = math.prod(dims) * torch.empty([], dtype=dtype).element_size()
|
||||||
|
|
||||||
|
# Note: hugepages are not directly supported with /dev/shm mmap files
|
||||||
|
# without mounting hugetlbfs there, so we fall back to plain page size.
|
||||||
|
if hugepage_size != "":
|
||||||
|
logger.warning(
|
||||||
|
"Hugepages are not supported with SHM allocator. "
|
||||||
|
"Falling back to plain page-size mmap."
|
||||||
|
)
|
||||||
|
|
||||||
|
page_size = mmap.PAGESIZE
|
||||||
|
alloc_bytes = math.ceil(n_bytes / page_size) * page_size
|
||||||
|
|
||||||
|
# Create an anonymous shared memory file descriptor via memfd_create
|
||||||
|
fd = None
|
||||||
|
try:
|
||||||
|
# MFD_CLOEXEC is standard on Linux 3.17+
|
||||||
|
fd = os.memfd_create(
|
||||||
|
f"sglang_host_pool_{uuid.uuid4().hex}",
|
||||||
|
flags=getattr(os, "MFD_CLOEXEC", 1),
|
||||||
|
)
|
||||||
|
except (AttributeError, OSError):
|
||||||
|
# Fallback to creating a file in /dev/shm if memfd_create is not supported
|
||||||
|
shm_path = f"/dev/shm/sglang_host_pool_{uuid.uuid4().hex}.mmap"
|
||||||
|
try:
|
||||||
|
fd = os.open(shm_path, os.O_CREAT | os.O_RDWR | os.O_TRUNC, 0o600)
|
||||||
|
try:
|
||||||
|
os.unlink(shm_path)
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
except Exception as e:
|
||||||
|
raise OSError(f"Failed to create shm file: {e}")
|
||||||
|
|
||||||
|
try:
|
||||||
|
os.ftruncate(fd, alloc_bytes)
|
||||||
|
mm = mmap.mmap(
|
||||||
|
fd,
|
||||||
|
alloc_bytes,
|
||||||
|
flags=mmap.MAP_SHARED | _MAP_POPULATE,
|
||||||
|
prot=mmap.PROT_READ | mmap.PROT_WRITE,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
# MADV_POPULATE_WRITE guarantees pages are populated and writable,
|
||||||
|
# throwing an error on failure (e.g. out of memory).
|
||||||
|
mm.madvise(_MADV_POPULATE_WRITE)
|
||||||
|
except OSError:
|
||||||
|
# Fall back to MAP_POPULATE if MADV_POPULATE_WRITE is not supported (<5.14 kernel).
|
||||||
|
pass
|
||||||
|
except Exception as e:
|
||||||
|
if fd is not None:
|
||||||
|
os.close(fd)
|
||||||
|
raise e
|
||||||
|
|
||||||
|
tensor = torch.frombuffer(mm, dtype=dtype, count=math.prod(dims)).reshape(dims)
|
||||||
|
return tensor, fd, mm
|
||||||
@@ -18,8 +18,8 @@ from sglang.srt.mem_cache.hicache_storage import (
|
|||||||
PoolTransfer,
|
PoolTransfer,
|
||||||
PoolTransferResult,
|
PoolTransferResult,
|
||||||
)
|
)
|
||||||
from sglang.srt.mem_cache.mmap_allocator import alloc_mmap
|
|
||||||
from sglang.srt.mem_cache.pool_host import HostKVCache
|
from sglang.srt.mem_cache.pool_host import HostKVCache
|
||||||
|
from sglang.srt.mem_cache.storage.mmap import alloc_mmap
|
||||||
from sglang.srt.mem_cache.storage.nixl.nixl_cleaner import HiCacheL3Cleaner
|
from sglang.srt.mem_cache.storage.nixl.nixl_cleaner import HiCacheL3Cleaner
|
||||||
|
|
||||||
from .nixl_registry import NixlRegistry
|
from .nixl_registry import NixlRegistry
|
||||||
|
|||||||
@@ -0,0 +1,8 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to SGLang project
|
||||||
|
|
||||||
|
"""Shared memory storage backend for SGLang HiCache."""
|
||||||
|
|
||||||
|
from .hicache_shm import HiCacheShm
|
||||||
|
|
||||||
|
__all__ = ["HiCacheShm"]
|
||||||
@@ -0,0 +1,96 @@
|
|||||||
|
import logging
|
||||||
|
from typing import Any, List, Optional
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.mem_cache.hicache_storage import (
|
||||||
|
HiCacheStorage,
|
||||||
|
HiCacheStorageConfig,
|
||||||
|
HiCacheStorageExtraInfo,
|
||||||
|
PoolTransfer,
|
||||||
|
PoolTransferResult,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class HiCacheShm(HiCacheStorage):
|
||||||
|
"""
|
||||||
|
Dummy storage backend for shared memory allocator.
|
||||||
|
Since shm is a local allocator, there's no actual storage transfer needed.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self, storage_config: HiCacheStorageConfig, mem_pool_host: Optional[Any] = None
|
||||||
|
):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def get(
|
||||||
|
self,
|
||||||
|
key: str,
|
||||||
|
target_location: Optional[Any] = None,
|
||||||
|
target_sizes: Optional[Any] = None,
|
||||||
|
) -> torch.Tensor | None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
def batch_get(
|
||||||
|
self,
|
||||||
|
keys: List[str],
|
||||||
|
target_locations: Optional[Any] = None,
|
||||||
|
target_sizes: Optional[Any] = None,
|
||||||
|
) -> List[torch.Tensor | None]:
|
||||||
|
return [None] * len(keys)
|
||||||
|
|
||||||
|
def set(
|
||||||
|
self,
|
||||||
|
key: str,
|
||||||
|
value: Optional[Any] = None,
|
||||||
|
target_location: Optional[Any] = None,
|
||||||
|
target_sizes: Optional[Any] = None,
|
||||||
|
) -> bool:
|
||||||
|
return True
|
||||||
|
|
||||||
|
def batch_set(
|
||||||
|
self,
|
||||||
|
keys: List[str],
|
||||||
|
values: Optional[Any] = None,
|
||||||
|
target_locations: Optional[Any] = None,
|
||||||
|
target_sizes: Optional[Any] = None,
|
||||||
|
) -> bool:
|
||||||
|
return True
|
||||||
|
|
||||||
|
def exists(self, key: str) -> bool:
|
||||||
|
return False
|
||||||
|
|
||||||
|
def batch_exists_v2(
|
||||||
|
self,
|
||||||
|
keys: List[str],
|
||||||
|
pool_transfers: Optional[List[PoolTransfer]] = None,
|
||||||
|
extra_info: Optional[HiCacheStorageExtraInfo] = None,
|
||||||
|
) -> PoolTransferResult:
|
||||||
|
return PoolTransferResult(0, {})
|
||||||
|
|
||||||
|
def batch_get_v2(
|
||||||
|
self,
|
||||||
|
transfers: List[PoolTransfer],
|
||||||
|
extra_info: Optional[HiCacheStorageExtraInfo] = None,
|
||||||
|
) -> dict[str, List[bool]]:
|
||||||
|
results = {}
|
||||||
|
for transfer in transfers:
|
||||||
|
keys = transfer.keys or []
|
||||||
|
results[transfer.name] = [False] * len(keys)
|
||||||
|
return results
|
||||||
|
|
||||||
|
def batch_set_v2(
|
||||||
|
self,
|
||||||
|
transfers: List[PoolTransfer],
|
||||||
|
extra_info: Optional[HiCacheStorageExtraInfo] = None,
|
||||||
|
) -> dict[str, List[bool]]:
|
||||||
|
results = {}
|
||||||
|
for transfer in transfers:
|
||||||
|
keys = transfer.keys or []
|
||||||
|
results[transfer.name] = [True] * len(keys)
|
||||||
|
return results
|
||||||
|
|
||||||
|
def clear(self) -> bool:
|
||||||
|
return True
|
||||||
@@ -2587,6 +2587,7 @@ class ServerArgs:
|
|||||||
"eic",
|
"eic",
|
||||||
"simm",
|
"simm",
|
||||||
"mori",
|
"mori",
|
||||||
|
"shm",
|
||||||
],
|
],
|
||||||
),
|
),
|
||||||
NS("memory"),
|
NS("memory"),
|
||||||
|
|||||||
@@ -83,6 +83,24 @@ class TestHostKVCache(CustomTestCase):
|
|||||||
self.assertIn("Double-free", msg)
|
self.assertIn("Double-free", msg)
|
||||||
self.assertIn(str(indices.tolist()), msg)
|
self.assertIn(str(indices.tolist()), msg)
|
||||||
|
|
||||||
|
def test_shm_allocator(self):
|
||||||
|
shm_host_pool = MHATokenToKVPoolHost(
|
||||||
|
device_pool=self.device_pool,
|
||||||
|
host_to_device_ratio=2.0,
|
||||||
|
host_size=0,
|
||||||
|
page_size=self.page_size,
|
||||||
|
layout="layer_first",
|
||||||
|
pin_memory=False,
|
||||||
|
device="cpu",
|
||||||
|
allocator_type="shm",
|
||||||
|
)
|
||||||
|
self.assertIsNotNone(shm_host_pool.fd)
|
||||||
|
self.assertGreaterEqual(shm_host_pool.fd, 0)
|
||||||
|
|
||||||
|
indices = shm_host_pool.alloc(4)
|
||||||
|
self.assertEqual(len(indices), 4)
|
||||||
|
shm_host_pool.free(indices)
|
||||||
|
|
||||||
def test_empty_free_keeps_release_list_empty(self):
|
def test_empty_free_keeps_release_list_empty(self):
|
||||||
self.assertEqual(self.host_pool.free(torch.empty(0, dtype=torch.int64)), 0)
|
self.assertEqual(self.host_pool.free(torch.empty(0, dtype=torch.int64)), 0)
|
||||||
self.assertEqual(self.host_pool.num_release_slots, 0)
|
self.assertEqual(self.host_pool.num_release_slots, 0)
|
||||||
|
|||||||
@@ -0,0 +1,119 @@
|
|||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
import sys
|
||||||
|
|
||||||
|
sys.modules["libtpu"] = None
|
||||||
|
import mmap
|
||||||
|
import os
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.mem_cache.pool_host.common import ShmHostTensorAllocator
|
||||||
|
from sglang.srt.mem_cache.storage.mmap import alloc_mmap, alloc_shm
|
||||||
|
|
||||||
|
|
||||||
|
class TestMmapAllocator(unittest.TestCase):
|
||||||
|
def test_alloc_mmap(self):
|
||||||
|
dims = (10, 1024)
|
||||||
|
dtype = torch.float32
|
||||||
|
tensor = alloc_mmap(dims, dtype)
|
||||||
|
self.assertEqual(tensor.shape, dims)
|
||||||
|
self.assertEqual(tensor.dtype, dtype)
|
||||||
|
# Verify it has mapped memory address
|
||||||
|
self.assertGreater(tensor.data_ptr(), 0)
|
||||||
|
|
||||||
|
def test_alloc_shm(self):
|
||||||
|
dims = (10, 1024)
|
||||||
|
dtype = torch.float32
|
||||||
|
tensor, fd, mm = alloc_shm(dims, dtype)
|
||||||
|
|
||||||
|
self.assertEqual(tensor.shape, dims)
|
||||||
|
self.assertEqual(tensor.dtype, dtype)
|
||||||
|
self.assertGreater(tensor.data_ptr(), 0)
|
||||||
|
self.assertGreaterEqual(fd, 0)
|
||||||
|
self.assertIsInstance(mm, mmap.mmap)
|
||||||
|
|
||||||
|
# Check that we can write to the tensor
|
||||||
|
tensor[0, 0] = 42.0
|
||||||
|
self.assertEqual(tensor[0, 0].item(), 42.0)
|
||||||
|
|
||||||
|
# Check that the FD is open and valid
|
||||||
|
try:
|
||||||
|
os.lseek(fd, 0, os.SEEK_SET)
|
||||||
|
except OSError:
|
||||||
|
self.fail("FD is not valid or closed")
|
||||||
|
|
||||||
|
# Cleanup
|
||||||
|
mm.close()
|
||||||
|
os.close(fd)
|
||||||
|
|
||||||
|
def test_shm_host_tensor_allocator(self):
|
||||||
|
allocator = ShmHostTensorAllocator()
|
||||||
|
dims = (2, 512)
|
||||||
|
dtype = torch.int32
|
||||||
|
|
||||||
|
tensor = allocator.allocate(dims, dtype, "cpu")
|
||||||
|
self.assertEqual(tensor.shape, dims)
|
||||||
|
self.assertEqual(tensor.dtype, dtype)
|
||||||
|
self.assertIsNotNone(allocator.fd)
|
||||||
|
self.assertGreaterEqual(allocator.fd, 0)
|
||||||
|
|
||||||
|
# Write data and check
|
||||||
|
tensor[1, 1] = 99
|
||||||
|
self.assertEqual(tensor[1, 1].item(), 99)
|
||||||
|
|
||||||
|
# Test destructor cleans up fd
|
||||||
|
fd = allocator.fd
|
||||||
|
# Trigger GC / deletion
|
||||||
|
del allocator
|
||||||
|
|
||||||
|
# Verify fd is closed
|
||||||
|
with self.assertRaises(OSError):
|
||||||
|
os.fstat(fd)
|
||||||
|
|
||||||
|
def test_alloc_shm_unlinked(self):
|
||||||
|
dims = (4, 256)
|
||||||
|
dtype = torch.float32
|
||||||
|
tensor, fd, mm = alloc_shm(dims, dtype)
|
||||||
|
|
||||||
|
# On Linux, the path of an unlinked fd shows up in /proc/self/fd/
|
||||||
|
# with a ' (deleted)' suffix.
|
||||||
|
fd_path = f"/proc/self/fd/{fd}"
|
||||||
|
try:
|
||||||
|
resolved_path = os.readlink(fd_path)
|
||||||
|
self.assertIn("sglang_host_pool_", resolved_path)
|
||||||
|
self.assertTrue(resolved_path.endswith(" (deleted)"))
|
||||||
|
except OSError:
|
||||||
|
# If procfs is not available or readlink fails, fallback to direct path existence check
|
||||||
|
self.assertFalse(os.path.exists(f"/dev/shm/sglang_host_pool_"))
|
||||||
|
|
||||||
|
# Cleanup
|
||||||
|
mm.close()
|
||||||
|
os.close(fd)
|
||||||
|
|
||||||
|
def test_alloc_shm_hugepage_warning(self):
|
||||||
|
from sglang.srt.environ import envs
|
||||||
|
|
||||||
|
envs.SGLANG_HUGEPAGE_SIZE.override("2MB")
|
||||||
|
try:
|
||||||
|
# Should succeed by falling back to plain page size mapping
|
||||||
|
dims = (2, 2)
|
||||||
|
tensor, fd, mm = alloc_shm(dims, torch.float32)
|
||||||
|
self.assertEqual(tensor.shape, dims)
|
||||||
|
mm.close()
|
||||||
|
os.close(fd)
|
||||||
|
finally:
|
||||||
|
envs.SGLANG_HUGEPAGE_SIZE.override(None)
|
||||||
|
|
||||||
|
def test_shm_host_tensor_allocator_invalid_device(self):
|
||||||
|
allocator = ShmHostTensorAllocator()
|
||||||
|
with self.assertRaises(AssertionError) as ctx:
|
||||||
|
allocator.allocate((2, 2), torch.float32, device="cuda")
|
||||||
|
self.assertIn("only supports CPU allocations", str(ctx.exception))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user