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

[mem_cache][5/N] refactor: extract host KV cache base layer into pool_host package (#27273)

Browse Source
This commit is contained in:
shuwenn
2026-06-20 20:44:08 +08:00
committed by GitHub
parent 1109acc24b
commit ff1fc1fbdf
16 changed files with 380 additions and 332 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/cache_controller.py
+1 -1
View File
@@ -31,7 +31,7 @@ from sglang.srt.mem_cache.hicache_storage import (
if TYPE_CHECKING:
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool_host import HostKVCache
from sglang.srt.mem_cache.pool_host import HostKVCache
from sglang.srt.distributed import (
get_pipeline_model_parallel_rank,
python/sglang/srt/mem_cache/hicache_storage.py
+1 -1
View File
@@ -14,7 +14,7 @@ import torch
from sglang.srt.environ import envs
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool_host import HostKVCache
from sglang.srt.mem_cache.pool_host import HostKVCache
logger = logging.getLogger(__name__)
python/sglang/srt/mem_cache/memory_pool_host.py
+11 -316
View File
@@ -1,11 +1,8 @@
from __future__ import annotations
import abc
import logging
import threading
from collections import defaultdict
from dataclasses import dataclass
from functools import wraps
from typing import TYPE_CHECKING, Any, Callable, Optional
if TYPE_CHECKING:
@@ -40,12 +37,10 @@ from sglang.jit_kernel.hicache import (
from sglang.jit_kernel.hisparse import transfer_cache_dsv4_mla
from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool,
KVCache,
MambaPool,
MHATokenToKVPool,
MLATokenToKVPool,
)
from sglang.srt.mem_cache.mmap_allocator import alloc_mmap
from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
_is_cuda = is_cuda()
@@ -74,321 +69,21 @@ if _is_npu:
logger = logging.getLogger(__name__)
# Host RAM to leave free when sizing HiCache pools (OS, other processes).
HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3)
from sglang.srt.mem_cache.pool_host import HostKVCache
from sglang.srt.mem_cache.pool_host.base import (
HICACHE_HOST_MEMORY_RESERVE_BYTES,
synchronized,
)
from sglang.srt.mem_cache.pool_host.common import (
ALLOC_MEMORY_FUNCS,
get_allocator_from_storage,
)
from sglang.srt.mem_cache.pool_host.hisparse import HiSparseHostPoolMixin
_WRITE_BACK_STAGING_PAGE_CHUNK = 64
def synchronized(func):
@wraps(func)
def wrapper(self, *args, **kwargs):
with self.lock:
return func(self, *args, **kwargs)
return wrapper
class HostTensorAllocator:
def __init__(self):
"""Initialize the HostTensorAllocator."""
self.dtype = None
self.dims = None
def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor:
assert (
device == "cpu"
), f"HostTensorAllocator only supports CPU allocations; got device={device!r}"
self.dtype = dtype
self.dims = dims
return alloc_mmap(dims, dtype)
class HiSparseHostPoolMixin:
def _round_up_to_page_size(self, size: int) -> int:
return (size + self.page_size - 1) // self.page_size * self.page_size
def alloc_page(self, num_pages: int) -> Optional[torch.Tensor]:
return self.alloc(num_pages * self.page_size)
def alloc_paged_token_slots(
self,
req_to_host_pool: torch.Tensor,
req_to_host_pool_allocated_len: torch.Tensor,
req_pool_idx: int,
start_pos: int,
num_tokens: int,
) -> torch.Tensor:
"""Allocate request host slots by page and return token-granular slots."""
device = req_to_host_pool.device
if num_tokens <= 0:
return torch.empty((0,), dtype=torch.int64, device=device)
allocated_len = int(req_to_host_pool_allocated_len[req_pool_idx])
end_pos = start_pos + num_tokens
page_end = self._round_up_to_page_size(end_pos)
assert start_pos <= allocated_len
if page_end > allocated_len:
num_new_pages = (page_end - allocated_len) // self.page_size
host_locs = self.alloc_page(num_new_pages)
if host_locs is None:
logger.error(
"HiSparse: host mem pool alloc failed for %d host pages "
"(req_pool_idx=%d, start_pos=%d, num_tokens=%d)",
num_new_pages,
req_pool_idx,
start_pos,
num_tokens,
)
raise RuntimeError(
f"HiSparse host mem pool alloc failed for {num_new_pages} pages"
)
req_to_host_pool[req_pool_idx, allocated_len:page_end] = host_locs.to(
device=device, non_blocking=True
)
req_to_host_pool_allocated_len[req_pool_idx] = page_end
return req_to_host_pool[req_pool_idx, start_pos:end_pos]
def allocated_host_indices(
self,
req_to_host_pool: torch.Tensor,
req_pool_idx: int,
allocated_len: int,
) -> torch.Tensor:
allocated_len = int(allocated_len)
host_len = min(
self._round_up_to_page_size(allocated_len),
req_to_host_pool.shape[1],
)
host_indices = req_to_host_pool[req_pool_idx, :host_len]
return host_indices[host_indices >= 0]
def get_allocator_from_storage(allocator_type):
if allocator_type == "mooncake":
try:
from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import (
MooncakeHostTensorAllocator,
)
return MooncakeHostTensorAllocator()
except ImportError:
logger.warning(
"Mooncake's tensor allocator requires mooncake >= 0.3.8.post1. "
"Please upgrade Mooncake by 'pip install mooncake-transfer-engine --upgrade'. "
"Fallback to use default allocator."
)
return HostTensorAllocator()
else:
return HostTensorAllocator()
def _cuda_host_register(buffer: torch.Tensor) -> None:
cudart = torch.cuda.cudart()
n_bytes = buffer.numel() * buffer.element_size()
rc = cudart.cudaHostRegister(buffer.data_ptr(), n_bytes, 0)
if int(rc) != 0:
raise RuntimeError(
f"cudaHostRegister failed (rc={int(rc)}, "
f"{cudart.cudaGetErrorString(rc)}) for ptr={buffer.data_ptr():#x} "
f"size={n_bytes}; host buffer is not pinned and device transfers "
f"may silently return stale data."
)
def alloc_with_host_register(
dims: tuple,
dtype: torch.dtype,
device: str,
pin_memory: bool,
allocator: HostTensorAllocator,
) -> torch.Tensor:
"""
Allocate tensor and register host memory with cudaHostRegister.
CudaHostRegister only applies when pin_memory=True.
"""
buffer = allocator.allocate(dims, dtype=dtype, device=device)
if pin_memory:
_cuda_host_register(buffer)
return buffer
def alloc_with_pin_memory(
dims: tuple,
dtype: torch.dtype,
device: str,
pin_memory: bool,
allocator: None,
) -> torch.Tensor:
"""
Allocate tensor using PyTorch's built-in pin_memory flag.
"""
buffer = torch.empty(dims, dtype=dtype, device=device, pin_memory=pin_memory)
return buffer
ALLOC_MEMORY_FUNCS = defaultdict(
lambda: alloc_with_host_register,
{
"npu": alloc_with_pin_memory,
"musa": alloc_with_pin_memory,
},
)
class HostKVCache(abc.ABC):
def __init__(
self,
device_pool: KVCache,
host_to_device_ratio: float,
host_size: int,
page_size: int,
layout: str,
pin_memory: bool,
device: str,
allocator_type: str = "default",
):
self.device_pool = device_pool
self.page_size = page_size
self.layout = layout
self.pin_memory = pin_memory
self.device = device
self.allocator = get_allocator_from_storage(allocator_type)
self.can_use_write_back_jit = False
self.dtype = device_pool.store_dtype
self.size_per_token = self.get_size_per_token()
if host_size > 0:
self.size = int(host_size * 1e9 // self.size_per_token)
else:
self.size = int(device_pool.size * host_to_device_ratio)
# Align up the host memory pool size to the page size
self.page_num = self.size // self.page_size + 1
self.size = self.page_num * self.page_size
self.start_layer = device_pool.start_layer
self.end_layer = device_pool.end_layer
assert (
self.size > device_pool.size
), "The host memory should be larger than the device memory with the current protocol"
# Verify there is enough available host memory.
host_mem = psutil.virtual_memory()
requested_bytes = self.size * self.size_per_token
available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES
if requested_bytes > available_bytes:
raise ValueError(
f"Not enough host memory available. Requesting "
f"{requested_bytes / 1e9:.2f} GB but only have "
f"{available_bytes / 1e9:.2f} GB free. Please reduce the "
f"size of the hierarchical cache."
)
else:
logger.info(
f"Allocating {requested_bytes / 1e9:.2f} GB host memory for hierarchical KV cache."
)
self.kv_buffer = self.init_kv_buffer()
# A lock for synchronized operations on memory allocation and state transitions.
self.lock = threading.RLock()
self.clear()
@abc.abstractmethod
def get_size_per_token(self):
raise NotImplementedError()
@abc.abstractmethod
def init_kv_buffer(self):
raise NotImplementedError()
@abc.abstractmethod
def load_to_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
) -> None:
"""
Load KV data from the host memory pool to the device memory pool for a specific layer.
"""
raise NotImplementedError()
@abc.abstractmethod
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend
) -> None:
"""
Backup KV data from the device memory pool to the host memory pool for all layers.
"""
raise NotImplementedError()
@abc.abstractmethod
def get_data_page(self, index, flat: bool = True) -> torch.Tensor:
"""
Get a flat data page from the host memory pool.
"""
raise NotImplementedError()
@abc.abstractmethod
def get_dummy_flat_data_page(self) -> torch.Tensor:
"""
Get a dummy flat data page from the host memory pool.
This is used for prefetching or initializing empty pages.
"""
raise NotImplementedError()
@abc.abstractmethod
def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
"""
Set a flat data page to the host memory pool.
"""
raise NotImplementedError()
def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
"""Return True if per-page strides are multiples of *page_size_bytes*.
Subclasses should override this with a layout-specific stride formula.
This base implementation logs a warning and returns False (safe default).
"""
logger.warning(
"%s does not implement is_stride_page_aligned(); assuming not aligned. "
"O_DIRECT with a file-based NIXL backend will fall back to copy mode for this pool.",
type(self).__name__,
)
return False
@synchronized
def clear(self):
# Initialize memory states and tracking structures.
self.mem_state = torch.zeros(
(self.size,), dtype=torch.uint8, device=self.device
)
self.free_slots = torch.arange(self.size, dtype=torch.int64)
def available_size(self):
return len(self.free_slots)
@synchronized
def alloc(self, need_size: int) -> Optional[torch.Tensor]:
assert (
need_size % self.page_size == 0
), "The requested size should be a multiple of the page size."
if need_size > self.available_size():
return None
select_index = self.free_slots[:need_size]
self.free_slots = self.free_slots[need_size:]
return select_index
@synchronized
def free(self, indices: torch.Tensor) -> int:
self.free_slots = torch.cat([self.free_slots, indices.cpu()])
return len(indices)
class MHATokenToKVPoolHost(HostKVCache):
device_pool: MHATokenToKVPool
python/sglang/srt/mem_cache/pool_host/__init__.py
+7
View File
@@ -0,0 +1,7 @@
from sglang.srt.mem_cache.pool_host.base import HostKVCache
from sglang.srt.mem_cache.pool_host.common import HostTensorAllocator
__all__ = [
"HostKVCache",
"HostTensorAllocator",
]
python/sglang/srt/mem_cache/pool_host/base.py
+177
View File
@@ -0,0 +1,177 @@
from __future__ import annotations
import abc
import logging
import threading
from functools import wraps
from typing import Optional
import psutil
import torch
from sglang.srt.mem_cache.memory_pool import KVCache
from sglang.srt.mem_cache.pool_host.common import get_allocator_from_storage
logger = logging.getLogger(__name__)
# Host RAM to leave free when sizing HiCache pools (OS, other processes).
HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3)
def synchronized(func):
@wraps(func)
def wrapper(self, *args, **kwargs):
with self.lock:
return func(self, *args, **kwargs)
return wrapper
class HostKVCache(abc.ABC):
def __init__(
self,
device_pool: KVCache,
host_to_device_ratio: float,
host_size: int,
page_size: int,
layout: str,
pin_memory: bool,
device: str,
allocator_type: str = "default",
):
self.device_pool = device_pool
self.page_size = page_size
self.layout = layout
self.pin_memory = pin_memory
self.device = device
self.allocator = get_allocator_from_storage(allocator_type)
self.can_use_write_back_jit = False
self.dtype = device_pool.store_dtype
self.size_per_token = self.get_size_per_token()
if host_size > 0:
self.size = int(host_size * 1e9 // self.size_per_token)
else:
self.size = int(device_pool.size * host_to_device_ratio)
# Align up the host memory pool size to the page size
self.page_num = self.size // self.page_size + 1
self.size = self.page_num * self.page_size
self.start_layer = device_pool.start_layer
self.end_layer = device_pool.end_layer
assert (
self.size > device_pool.size
), "The host memory should be larger than the device memory with the current protocol"
# Verify there is enough available host memory.
host_mem = psutil.virtual_memory()
requested_bytes = self.size * self.size_per_token
available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES
if requested_bytes > available_bytes:
raise ValueError(
f"Not enough host memory available. Requesting "
f"{requested_bytes / 1e9:.2f} GB but only have "
f"{available_bytes / 1e9:.2f} GB free. Please reduce the "
f"size of the hierarchical cache."
)
else:
logger.info(
f"Allocating {requested_bytes / 1e9:.2f} GB host memory for hierarchical KV cache."
)
self.kv_buffer = self.init_kv_buffer()
# A lock for synchronized operations on memory allocation and state transitions.
self.lock = threading.RLock()
self.clear()
@abc.abstractmethod
def get_size_per_token(self):
raise NotImplementedError()
@abc.abstractmethod
def init_kv_buffer(self):
raise NotImplementedError()
@abc.abstractmethod
def load_to_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
) -> None:
"""
Load KV data from the host memory pool to the device memory pool for a specific layer.
"""
raise NotImplementedError()
@abc.abstractmethod
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend
) -> None:
"""
Backup KV data from the device memory pool to the host memory pool for all layers.
"""
raise NotImplementedError()
@abc.abstractmethod
def get_data_page(self, index, flat: bool = True) -> torch.Tensor:
"""
Get a flat data page from the host memory pool.
"""
raise NotImplementedError()
@abc.abstractmethod
def get_dummy_flat_data_page(self) -> torch.Tensor:
"""
Get a dummy flat data page from the host memory pool.
This is used for prefetching or initializing empty pages.
"""
raise NotImplementedError()
@abc.abstractmethod
def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
"""
Set a flat data page to the host memory pool.
"""
raise NotImplementedError()
def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
"""Return True if per-page strides are multiples of *page_size_bytes*.
Subclasses should override this with a layout-specific stride formula.
This base implementation logs a warning and returns False (safe default).
"""
logger.warning(
"%s does not implement is_stride_page_aligned(); assuming not aligned. "
"O_DIRECT with a file-based NIXL backend will fall back to copy mode for this pool.",
type(self).__name__,
)
return False
@synchronized
def clear(self):
# Initialize memory states and tracking structures.
self.mem_state = torch.zeros(
(self.size,), dtype=torch.uint8, device=self.device
)
self.free_slots = torch.arange(self.size, dtype=torch.int64)
def available_size(self):
return len(self.free_slots)
@synchronized
def alloc(self, need_size: int) -> Optional[torch.Tensor]:
assert (
need_size % self.page_size == 0
), "The requested size should be a multiple of the page size."
if need_size > self.available_size():
return None
select_index = self.free_slots[:need_size]
self.free_slots = self.free_slots[need_size:]
return select_index
@synchronized
def free(self, indices: torch.Tensor) -> int:
self.free_slots = torch.cat([self.free_slots, indices.cpu()])
return len(indices)
python/sglang/srt/mem_cache/pool_host/common.py
+97
View File
@@ -0,0 +1,97 @@
from __future__ import annotations
import logging
from collections import defaultdict
import torch
from sglang.srt.mem_cache.mmap_allocator import alloc_mmap
logger = logging.getLogger(__name__)
class HostTensorAllocator:
def __init__(self):
"""Initialize the HostTensorAllocator."""
self.dtype = None
self.dims = None
def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor:
assert (
device == "cpu"
), f"HostTensorAllocator only supports CPU allocations; got device={device!r}"
self.dtype = dtype
self.dims = dims
return alloc_mmap(dims, dtype)
def get_allocator_from_storage(allocator_type):
if allocator_type == "mooncake":
try:
from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import (
MooncakeHostTensorAllocator,
)
return MooncakeHostTensorAllocator()
except ImportError:
logger.warning(
"Mooncake's tensor allocator requires mooncake >= 0.3.8.post1. "
"Please upgrade Mooncake by 'pip install mooncake-transfer-engine --upgrade'. "
"Fallback to use default allocator."
)
return HostTensorAllocator()
else:
return HostTensorAllocator()
def _cuda_host_register(buffer: torch.Tensor) -> None:
cudart = torch.cuda.cudart()
n_bytes = buffer.numel() * buffer.element_size()
rc = cudart.cudaHostRegister(buffer.data_ptr(), n_bytes, 0)
if int(rc) != 0:
raise RuntimeError(
f"cudaHostRegister failed (rc={int(rc)}, "
f"{cudart.cudaGetErrorString(rc)}) for ptr={buffer.data_ptr():#x} "
f"size={n_bytes}; host buffer is not pinned and device transfers "
f"may silently return stale data."
)
def alloc_with_host_register(
dims: tuple,
dtype: torch.dtype,
device: str,
pin_memory: bool,
allocator: HostTensorAllocator,
) -> torch.Tensor:
"""
Allocate tensor and register host memory with cudaHostRegister.
CudaHostRegister only applies when pin_memory=True.
"""
buffer = allocator.allocate(dims, dtype=dtype, device=device)
if pin_memory:
_cuda_host_register(buffer)
return buffer
def alloc_with_pin_memory(
dims: tuple,
dtype: torch.dtype,
device: str,
pin_memory: bool,
allocator: None,
) -> torch.Tensor:
"""
Allocate tensor using PyTorch's built-in pin_memory flag.
"""
buffer = torch.empty(dims, dtype=dtype, device=device, pin_memory=pin_memory)
return buffer
ALLOC_MEMORY_FUNCS = defaultdict(
lambda: alloc_with_host_register,
{
"npu": alloc_with_pin_memory,
"musa": alloc_with_pin_memory,
},
)
python/sglang/srt/mem_cache/pool_host/hisparse.py
+71
View File
@@ -0,0 +1,71 @@
from __future__ import annotations
import logging
from typing import Optional
import torch
logger = logging.getLogger(__name__)
class HiSparseHostPoolMixin:
def _round_up_to_page_size(self, size: int) -> int:
return (size + self.page_size - 1) // self.page_size * self.page_size
def alloc_page(self, num_pages: int) -> Optional[torch.Tensor]:
return self.alloc(num_pages * self.page_size)
def alloc_paged_token_slots(
self,
req_to_host_pool: torch.Tensor,
req_to_host_pool_allocated_len: torch.Tensor,
req_pool_idx: int,
start_pos: int,
num_tokens: int,
) -> torch.Tensor:
"""Allocate request host slots by page and return token-granular slots."""
device = req_to_host_pool.device
if num_tokens <= 0:
return torch.empty((0,), dtype=torch.int64, device=device)
allocated_len = int(req_to_host_pool_allocated_len[req_pool_idx])
end_pos = start_pos + num_tokens
page_end = self._round_up_to_page_size(end_pos)
assert start_pos <= allocated_len
if page_end > allocated_len:
num_new_pages = (page_end - allocated_len) // self.page_size
host_locs = self.alloc_page(num_new_pages)
if host_locs is None:
logger.error(
"HiSparse: host mem pool alloc failed for %d host pages "
"(req_pool_idx=%d, start_pos=%d, num_tokens=%d)",
num_new_pages,
req_pool_idx,
start_pos,
num_tokens,
)
raise RuntimeError(
f"HiSparse host mem pool alloc failed for {num_new_pages} pages"
)
req_to_host_pool[req_pool_idx, allocated_len:page_end] = host_locs.to(
device=device, non_blocking=True
)
req_to_host_pool_allocated_len[req_pool_idx] = page_end
return req_to_host_pool[req_pool_idx, start_pos:end_pos]
def allocated_host_indices(
self,
req_to_host_pool: torch.Tensor,
req_pool_idx: int,
allocated_len: int,
) -> torch.Tensor:
allocated_len = int(allocated_len)
host_len = min(
self._round_up_to_page_size(allocated_len),
req_to_host_pool.shape[1],
)
host_indices = req_to_host_pool[req_pool_idx, :host_len]
return host_indices[host_indices >= 0]
python/sglang/srt/mem_cache/storage/aibrix_kvcache/aibrix_kvcache_storage.py
+1 -1
View File
@@ -18,7 +18,7 @@ from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorageConfig,
HiCacheStorageExtraInfo,
)
from sglang.srt.mem_cache.memory_pool_host import HostKVCache
from sglang.srt.mem_cache.pool_host import HostKVCache
logger = logging.getLogger(__name__)
python/sglang/srt/mem_cache/storage/eic/eic_storage.py
+1 -1
View File
@@ -13,7 +13,7 @@ from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorageConfig,
HiCacheStorageExtraInfo,
)
from sglang.srt.mem_cache.memory_pool_host import HostKVCache
from sglang.srt.mem_cache.pool_host import HostKVCache
logger = logging.getLogger(__name__)
python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py
+1 -1
View File
@@ -22,7 +22,7 @@ from sglang.srt.mem_cache.hicache_storage import (
PoolTransfer,
PoolTransferResult,
)
from sglang.srt.mem_cache.memory_pool_host import HostKVCache
from sglang.srt.mem_cache.pool_host import HostKVCache
from sglang.srt.mem_cache.storage.hf3fs.hf3fs_client import Hf3fsClient
from sglang.srt.observability.metrics_collector import StorageMetrics
python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
+2 -5
View File
@@ -21,11 +21,8 @@ from sglang.srt.mem_cache.hicache_storage import (
PoolTransfer,
PoolTransferResult,
)
from sglang.srt.mem_cache.memory_pool_host import (
HostKVCache,
HostTensorAllocator,
MLATokenToKVPoolHost,
)
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
from sglang.srt.mem_cache.pool_host import HostKVCache, HostTensorAllocator
from sglang.srt.observability.metrics_collector import StorageMetrics
DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024 # 16 MB
python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py
+1 -1
View File
@@ -13,8 +13,8 @@ from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorageConfig,
HiCacheStorageExtraInfo,
)
from sglang.srt.mem_cache.memory_pool_host import HostKVCache
from sglang.srt.mem_cache.mmap_allocator import alloc_mmap
from sglang.srt.mem_cache.pool_host import HostKVCache
from .nixl_registry import NixlRegistry
from .nixl_utils import NixlBackendConfig, NixlBackendSelection, NixlFileManager
python/sglang/srt/mem_cache/storage/simm/hicache_simm.py
+1 -1
View File
@@ -16,7 +16,7 @@ from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorageConfig,
HiCacheStorageExtraInfo,
)
from sglang.srt.mem_cache.memory_pool_host import HostKVCache
from sglang.srt.mem_cache.pool_host import HostKVCache
# Third Party
try: