diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index 9ddd167d6..b7d450753 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -18,6 +18,7 @@ from sglang.srt.mem_cache.memory_pool import ( MLATokenToKVPool, 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.mla import MLATokenToKVPoolHost from sglang.srt.server_args import ServerArgs @@ -54,6 +55,8 @@ class DecodeKVCacheOffloadManager: self.page_size, (env_stride // self.page_size) * self.page_size ) kv_cache = self.token_to_kv_pool_allocator.get_kvcache() + allocator_type = get_allocator_type(server_args) + if isinstance(kv_cache, MHATokenToKVPool): self.decode_host_mem_pool = get_mha_host_pool_cls(kv_cache)( kv_cache, @@ -61,6 +64,7 @@ class DecodeKVCacheOffloadManager: server_args.hicache_size, self.page_size, server_args.hicache_mem_layout, + allocator_type=allocator_type, ) elif isinstance(kv_cache, MLATokenToKVPool): self.decode_host_mem_pool = MLATokenToKVPoolHost( @@ -69,6 +73,7 @@ class DecodeKVCacheOffloadManager: server_args.hicache_size, self.page_size, server_args.hicache_mem_layout, + allocator_type=allocator_type, ) else: raise ValueError("Unsupported KV cache type for decode offload") diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 57166f23a..8551a3882 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -48,6 +48,7 @@ from sglang.srt.mem_cache.memory_pool import ( MiniMaxSparseKVPool, 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.mla import MLATokenToKVPoolHost from sglang.srt.mem_cache.radix_cache import ( @@ -80,6 +81,8 @@ class HiRadixCache(RadixCache): self.page_size = params.page_size self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache() + allocator_type = get_allocator_type(server_args) + if isinstance(self.kv_cache, MHATokenToKVPool): self.token_to_kv_pool_host = get_mha_host_pool_cls(self.kv_cache)( self.kv_cache, @@ -87,7 +90,7 @@ class HiRadixCache(RadixCache): server_args.hicache_size, self.page_size, server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, + allocator_type=allocator_type, ) elif isinstance(self.kv_cache, DSATokenToKVPool): # 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, self.page_size, server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, + allocator_type=allocator_type, ) else: raise ValueError("HiRadixCache only supports MHA, MLA, DSA, and MSA models") diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index a7af5e51a..22d500a26 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -21,6 +21,7 @@ from sglang.srt.mem_cache.memory_pool_host import ( MambaPoolHost, PoolEntry, ) +from sglang.srt.mem_cache.pool_host.common import get_allocator_type from sglang.srt.mem_cache.pool_host.mha import ( MHATokenToKOnlyPoolHost, get_mha_host_pool_cls, @@ -40,6 +41,10 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def _get_allocator_type(server_args: ServerArgs) -> str: + return get_allocator_type(server_args) + + def _make_layer_mapper( layer_mapping: dict[int, int], transfer_layer_num: int, @@ -72,7 +77,7 @@ def build_kv_host_pool( server_args.hicache_size, page_size, server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, + allocator_type=_get_allocator_type(server_args), **kwargs, ) @@ -353,7 +358,7 @@ def build_deepseek_v4_hicache_stack( num_host_pages=swa_num_host_pages, slot_page_size=kvcache.swa_page_size, 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 entries.append( @@ -379,7 +384,7 @@ def build_deepseek_v4_hicache_stack( num_host_pages=num_host_pages, slot_page_size=page_size, 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( pool_name=str(PoolName.DEEPSEEK_V4_C4_INDEXER), @@ -391,7 +396,7 @@ def build_deepseek_v4_hicache_stack( num_host_pages=num_host_pages, slot_page_size=page_size, layout=server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, + allocator_type=_get_allocator_type(server_args), ) entries.extend( [ @@ -422,7 +427,7 @@ def build_deepseek_v4_hicache_stack( num_host_pages=swa_num_host_pages, swa_page_size=kvcache.swa_page_size, 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( 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, swa_page_size=kvcache.swa_page_size, layout=server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, + allocator_type=_get_allocator_type(server_args), ) entries.extend( [ @@ -465,7 +470,7 @@ def build_deepseek_v4_hicache_stack( num_host_pages=num_host_pages, slot_page_size=page_size, 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. # page_size=256 % 128 == 0, so state pool is not consumed on load. @@ -533,7 +538,7 @@ def build_hybrid_mamba_stack( mamba_pool, server_args.hicache_ratio, server_args.hicache_size, - allocator_type=server_args.hicache_storage_backend, + allocator_type=_get_allocator_type(server_args), layout=server_args.hicache_mem_layout, ) entries = [ @@ -1107,7 +1112,7 @@ class _DsaStrategy(StackStrategy): full_kv_pool, kv_host_pool, server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, + allocator_type=_get_allocator_type(server_args), ), prefetch_threshold=prefetch_threshold, model_name=model_name, @@ -1545,7 +1550,7 @@ def attach_hybrid_dsa_pool_to_hiradix_cache( kv, kv_host_pool, 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, storage_backend_extra_config=extra_config, diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index 84d5e6bd5..7018bc579 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/python/sglang/srt/mem_cache/pool_host/base.py @@ -139,6 +139,7 @@ class HostKVCache(abc.ABC): ) 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. self.lock = threading.RLock() diff --git a/python/sglang/srt/mem_cache/pool_host/common.py b/python/sglang/srt/mem_cache/pool_host/common.py index 07a2cdfb9..770135654 100644 --- a/python/sglang/srt/mem_cache/pool_host/common.py +++ b/python/sglang/srt/mem_cache/pool_host/common.py @@ -1,11 +1,13 @@ from __future__ import annotations +import json import logging +import os from collections import defaultdict 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__) @@ -25,6 +27,43 @@ class HostTensorAllocator: 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): if allocator_type == "mooncake": try: @@ -54,10 +93,30 @@ def get_allocator_from_storage(allocator_type): exc, ) return HostTensorAllocator() + elif allocator_type == "shm": + return ShmHostTensorAllocator() else: 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: cudart = torch.cuda.cudart() n_bytes = buffer.numel() * buffer.element_size() diff --git a/python/sglang/srt/mem_cache/storage/backend_factory.py b/python/sglang/srt/mem_cache/storage/backend_factory.py index 093ac86f1..0fe83fdaf 100644 --- a/python/sglang/srt/mem_cache/storage/backend_factory.py +++ b/python/sglang/srt/mem_cache/storage/backend_factory.py @@ -187,6 +187,8 @@ class StorageBackendFactory: return backend_class(storage_config, mem_pool_host) elif backend_name == "mori": return backend_class(storage_config, mem_pool_host) + elif backend_name == "shm": + return backend_class(storage_config, mem_pool_host) else: raise ValueError(f"Unknown built-in backend: {backend_name}") @@ -237,3 +239,9 @@ StorageBackendFactory.register_backend( "sglang.srt.mem_cache.storage.umbp.umbp_store", "UMBPStore", ) + +StorageBackendFactory.register_backend( + "shm", + "sglang.srt.mem_cache.storage.shm", + "HiCacheShm", +) diff --git a/python/sglang/srt/mem_cache/storage/mmap/__init__.py b/python/sglang/srt/mem_cache/storage/mmap/__init__.py new file mode 100644 index 000000000..6ddd9552d --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/mmap/__init__.py @@ -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", +] diff --git a/python/sglang/srt/mem_cache/mmap_allocator.py b/python/sglang/srt/mem_cache/storage/mmap/mmap_allocator.py similarity index 62% rename from python/sglang/srt/mem_cache/mmap_allocator.py rename to python/sglang/srt/mem_cache/storage/mmap/mmap_allocator.py index f2aeeabfa..a3afd7ca3 100644 --- a/python/sglang/srt/mem_cache/mmap_allocator.py +++ b/python/sglang/srt/mem_cache/storage/mmap/mmap_allocator.py @@ -4,6 +4,7 @@ import logging import math import mmap import os +import uuid import weakref import torch @@ -39,6 +40,7 @@ _MAP_HUGETLB = 0x40000 _MAP_HUGE_2MB = 21 << 26 # 0x1400000 _MAP_HUGE_1GB = 30 << 26 # 0x78000000 _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: @@ -124,4 +126,76 @@ def alloc_mmap(dims: tuple, dtype: torch.dtype) -> torch.Tensor: flags=mmap.MAP_SHARED | mmap.MAP_ANONYMOUS | _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 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 diff --git a/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py b/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py index f1437c57e..1f3227073 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py +++ b/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py @@ -18,8 +18,8 @@ from sglang.srt.mem_cache.hicache_storage import ( PoolTransfer, 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.storage.mmap import alloc_mmap from sglang.srt.mem_cache.storage.nixl.nixl_cleaner import HiCacheL3Cleaner from .nixl_registry import NixlRegistry diff --git a/python/sglang/srt/mem_cache/storage/shm/__init__.py b/python/sglang/srt/mem_cache/storage/shm/__init__.py new file mode 100644 index 000000000..0f47f41f5 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/shm/__init__.py @@ -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"] diff --git a/python/sglang/srt/mem_cache/storage/shm/hicache_shm.py b/python/sglang/srt/mem_cache/storage/shm/hicache_shm.py new file mode 100644 index 000000000..5582bbbb3 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/shm/hicache_shm.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index bf74ffc25..a4ae73494 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2587,6 +2587,7 @@ class ServerArgs: "eic", "simm", "mori", + "shm", ], ), NS("memory"), diff --git a/test/registered/unit/mem_cache/test_mem_pool_host.py b/test/registered/unit/mem_cache/test_mem_pool_host.py index 510bf074c..ce94365d4 100644 --- a/test/registered/unit/mem_cache/test_mem_pool_host.py +++ b/test/registered/unit/mem_cache/test_mem_pool_host.py @@ -83,6 +83,24 @@ class TestHostKVCache(CustomTestCase): self.assertIn("Double-free", 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): self.assertEqual(self.host_pool.free(torch.empty(0, dtype=torch.int64)), 0) self.assertEqual(self.host_pool.num_release_slots, 0) diff --git a/test/registered/unit/mem_cache/test_mmap_allocator.py b/test/registered/unit/mem_cache/test_mmap_allocator.py new file mode 100644 index 000000000..0f1af9c4d --- /dev/null +++ b/test/registered/unit/mem_cache/test_mmap_allocator.py @@ -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()