chore: import upstream snapshot with attribution
PR Test (NPU) / check-changes (push) Has been cancelled
PR Test (NPU) / pr-gate (push) Has been cancelled
PR Test (NPU) / set-image-config (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-4-npu-a3 (push) Has been cancelled
PR Test (NPU) / stage-b-test-16-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-1-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-2-npu-a3 (push) Has been cancelled
PR Test (Arm64) / pr-gate (push) Has been cancelled
PR Test (Arm64) / check-changes (push) Has been cancelled
PR Test (Arm64) / build-test (push) Has been cancelled
PR Test (sgl-router) / gate (push) Has been cancelled
PR Test (sgl-router) / tier-1 — lint (push) Has been cancelled
PR Test (sgl-router) / tier-2 — build + test (push) Has been cancelled
PR Test (sgl-router) / tier-3 — docker (placeholder) (push) Has been cancelled
PR Test (sgl-router) / tier-3 — k8s integration (push) Has been cancelled
PR Test (sgl-router) / tier-3 — e2e (push) Has been cancelled
PR Test (sgl-router) / finish (push) Has been cancelled
PR Test (NPU) / single-node-poc (map[name:qwen3_6_27b_w8a8_1p_in64k_out1k_50ms runner:linux-aarch64-a3-2 test_case:test/registered/ascend/performance/qwen3_6_27b/test_npu_qwen3_6_27b_w8a8_1p_in64k_out1k_50ms.py test_type:perf]) (push) Has been cancelled
PR Test (NPU) / pr-test-npu-finish (push) Has been cancelled
PR Test (Xeon) / pr-gate (push) Has been cancelled
PR Test (Xeon) / check-changes (push) Has been cancelled
PR Test (Xeon) / build-test (, xeon-gnr, base-b-test-cpu) (push) Has been cancelled
PR Test (XPU) / check-changes (push) Has been cancelled
PR Test (XPU) / pr-gate (push) Has been cancelled
PR Test (XPU) / stage-a-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / wait-for-stage-a (push) Has been cancelled
PR Test (XPU) / stage-b-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / finish (push) Has been cancelled
CI Model Inventory / build-inventory (push) Has been cancelled
Lint / lint (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Compilation Check (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Manual Policy (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Request Processing (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Summary (push) Has been cancelled
PR Test (SMG) / build-wheel (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on windows (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (x86_64 - auto) (push) Has been cancelled
PR Test (SMG) / python-unit-tests (push) Has been cancelled
PR Test (SMG) / unit-tests (push) Has been cancelled
PR Test (SMG) / benchmarks (push) Has been cancelled
PR Test (SMG) / chat-completions (push) Has been cancelled
PR Test (SMG) / chat-completions-4gpu (push) Has been cancelled
PR Test (SMG) / e2e (push) Has been cancelled
PR Test (SMG) / docker-build-test (push) Has been cancelled
PR Test (SMG) / k8s-integration (push) Has been cancelled
PR Test (SMG) / finish (push) Has been cancelled
PR Test (SMG) / summarize-benchmarks (push) Has been cancelled
Release SGLang Model Gateway Docker Image / publish (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Build SDist (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Upload to PyPI (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (aarch64, 12.9, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (x86_64, 12.9, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu129 (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (aarch64, 13.0, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (x86_64, 13.0, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu130 (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 700) (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 720) (push) Has been cancelled
Release SGLang Kernels / release-rocm700 (push) Has been cancelled
Release SGLang Kernels / release-rocm720 (push) Has been cancelled
Release SGLang Kernels / build-musa43 (43, 3.10) (push) Has been cancelled
Release SGLang Kernels / release-musa43 (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 12:38:16 +08:00
commit 94057c3d3e
7152 changed files with 2120455 additions and 0 deletions
@@ -0,0 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to SGLang project
"""Storage backend module for SGLang HiCache."""
from .backend_factory import StorageBackendFactory
__all__ = [
"StorageBackendFactory",
]
@@ -0,0 +1,37 @@
# AIBrix KVCache as L3 KV Cache
This document provides brief instructions for setting up a AIBrixKVCache storage backend + AIBrixKVCache + SGLang runtime environment from scratch, describing how to utilize AIBrixKVCache as the L3 KV cache for SGLang.
The process consists of three main steps:
## Step1:Install AIbrix KVCache
Refer to the [AIBrix KVCache documentation](https://github.com/vllm-project/aibrix/blob/main/python/aibrix_kvcache/README.md) to install AIBrix KVCache.
## Step2: Deploy AIBrix Distributed KVCache Storage
AIBrix KVCache currently supports multiple distributed KVCache backends, including ByteDance's open-source Infinistore and the not-yet-open source PrisKV incubated by ByteDance's PrisDB & IAAS & DMI team.
For the Infinistore installation process, please refer to [this link](https://github.com/bytedance/InfiniStore).
PrisKV for AIBrix KVCache is currently in the open-source preparation stage, and no public documentation is available yet.
## Step3: Deploy Model Serving
For information on configuring a distributed KVCache backend for AIBrixKVCache, please refer to [this link](https://aibrix.readthedocs.io/latest/designs/aibrix-kvcache-offloading-framework.html)
Using PrisKV as an example, the startup command is as follows:
```bash
export AIBRIX_KV_CACHE_OL_L1_CACHE_ENABLED="0"
export AIBRIX_KV_CACHE_OL_L2_CACHE_BACKEND="PRIS"
export AIBRIX_KV_CACHE_OL_PRIS_REMOTE_ADDR="127.0.0.1"
export AIBRIX_KV_CACHE_OL_PRIS_REMOTE_PORT="6379"
export AIBRIX_KV_CACHE_OL_PRIS_PASSWORD="kvcache-redis"
MODEL_LENGTH=32768&&NCCL_MIN_NCHANNELS=24&&NCCL_IB_QPS_PER_CONNECTION=8&&NCCL_DEBUG=INFO \
python3 -m sglang.launch_server \
--model-path /code/models/Qwen3-32B \
--host 0.0.0.0 --port 8080 \
--enable-hierarchical-cache \
--hicache-storage-backend aibrix \
--page-size 16 \
--hicache-write-policy write_back \
--enable-metrics --hicache-ratio=2
```
@@ -0,0 +1,157 @@
import logging
from typing import Any, List, Optional
import torch
from aibrix_kvcache import (
BaseKVCacheManager,
BlockHashes,
KVCacheBlockLayout,
KVCacheBlockSpec,
KVCacheConfig,
KVCacheTensorSpec,
ModelSpec,
)
from aibrix_kvcache.common.absl_logging import log_every_n_seconds
from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorage,
HiCacheStorageConfig,
HiCacheStorageExtraInfo,
)
from sglang.srt.mem_cache.pool_host import HostKVCache
logger = logging.getLogger(__name__)
class AibrixKVCacheStorage(HiCacheStorage):
def __init__(self, storage_config: HiCacheStorageConfig, mem_pool: HostKVCache):
if storage_config is not None:
self.is_mla_backend = storage_config.is_mla_model
self.local_rank = storage_config.tp_rank
else:
self.is_mla_backend = False
self.local_rank = 0
kv_cache = mem_pool.device_pool
self.page_size = mem_pool.page_size
self.kv_cache_dtype = kv_cache.dtype
self.layer_num = kv_cache.layer_num
self.kv_head_ids = [
self.local_rank * kv_cache.head_num + i for i in range(kv_cache.head_num)
]
if not self.is_mla_backend:
self.layer_ids = range(
kv_cache.start_layer, kv_cache.end_layer
) # for pipeline parallel
self.block_spec = KVCacheBlockSpec(
block_ntokens=self.page_size,
block_dtype=self.kv_cache_dtype,
block_layout=KVCacheBlockLayout(KVCacheBlockLayout.NCLD),
tensor_spec=KVCacheTensorSpec(
heads=self.kv_head_ids,
layers=self.layer_ids,
head_size=kv_cache.head_dim,
),
)
logger.info(self.block_spec)
config = KVCacheConfig(
block_spec=self.block_spec, model_spec=ModelSpec(102400)
)
self.kv_cache_manager = BaseKVCacheManager(config)
else:
raise NotImplementedError(
"MLA is not supported by AibrixKVCacheStorage yet."
)
def _aibrix_kvcache_metrics_report(self):
self.kv_cache_manager.metrics.summary()
self.kv_cache_manager.metrics.reset()
def batch_get(
self,
keys: List[str],
target_locations: List[torch.Tensor],
target_sizes: Optional[Any] = None,
) -> List[torch.Tensor | None]:
block_hash = BlockHashes(keys, self.page_size)
status = self.kv_cache_manager.acquire(None, block_hash)
log_every_n_seconds(
logger, logging.INFO, self._aibrix_kvcache_metrics_report(), 1
)
if status.is_ok():
num_fetched_tokens, handle = status.value
kv_blocks = handle.to_tensors()
assert len(kv_blocks) == len(target_locations)
for i in range(len(kv_blocks)):
assert (
target_locations[i].nbytes == kv_blocks[i].nbytes
), f"{target_locations[i].nbytes}, {kv_blocks[i].nbytes}"
target_locations[i].copy_(kv_blocks[i].flatten())
handle.release()
return target_locations
return [None] * len(keys)
def get(
self,
key: str,
target_location: Optional[Any] = None,
target_size: Optional[Any] = None,
) -> torch.Tensor | None:
return self.batch_get([key], [target_location], [target_size])[0]
def batch_set(
self,
keys: List[str],
values: Optional[Any] = None,
target_locations: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> bool:
block_hash = BlockHashes(keys, self.page_size)
status = self.kv_cache_manager.allocate_for(None, block_hash)
if not status.is_ok():
logger.warning(
f"aibrix_kvcache set allocate failed, error_code {status.error_code}"
)
return False
handle = status.value
tensors = handle.to_tensors()
if len(tensors) != len(values):
logger.warning("aibrix_kvcache set allocate not enough")
return False
for i in range(len(tensors)):
assert (
tensors[i].nbytes == values[i].nbytes
), f"{tensors[i].nbytes}, {values[i].nbytes}"
tensors[i].reshape(values[i].shape).copy_(values[i]).reshape(
tensors[i].shape
)
status = self.kv_cache_manager.put(None, block_hash, handle)
if not status.is_ok():
logger.info(
f"AIBrix KVCache Storage set failed, error_code {status.error_code}"
)
return False
completed = status.value
return completed == len(keys) * self.page_size
def set(
self,
key: str,
value: Optional[Any] = None,
target_location: Optional[Any] = None,
target_size: Optional[Any] = None,
) -> bool:
return self.batch_set([key], [value], [target_location], [target_size])
def batch_exists(
self, keys: List[str], extra_info: Optional[HiCacheStorageExtraInfo] = None
) -> int:
block_hash = BlockHashes(keys, self.page_size)
status = self.kv_cache_manager.exists(None, block_hash)
if status.is_ok():
return status.value // self.page_size
return 0
def exists(self, key: str) -> bool | dict:
return self.batch_exists([key]) > 0
@@ -0,0 +1,97 @@
import logging
import os
import torch
import torch.distributed
from aibrix_kvcache.common.absl_logging import log_every_n_seconds
from aibrix_kvcache_storage import AibrixKVCacheStorage
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
def setup():
os.environ["RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = "63886"
class AIBrixKVCacheStorageTest:
def test_with_page_size(self):
config = HiCacheStorageConfig(
tp_rank=0,
tp_size=1,
is_mla_model=False,
is_page_first_layout=True,
model_name="test",
)
for page_size in range(1, 3):
logger.info(f"page_size: {page_size}")
batch_size = 2
head_num = 1
layer_num = 64
head_dim = 128
kv_cache = MHATokenToKVPool(
1024,
page_size,
torch.float16,
head_num,
head_dim,
layer_num,
"cpu",
False,
0,
layer_num,
)
mem_pool = MHATokenToKVPoolHost(kv_cache, 2, 0, page_size, "layer_first")
query_length = batch_size * 2
partial = batch_size
self.aibrix_kvcache = AibrixKVCacheStorage(config, mem_pool)
target_shape = (2, layer_num, page_size, head_num, head_dim)
rand_tensor = [
torch.rand(target_shape, dtype=torch.float16)
for _ in range(query_length)
]
keys = ["hash" + str(i) for i in range(query_length)]
partial_keys = keys[batch_size:query_length]
assert self.aibrix_kvcache.batch_exists(keys) == 0
assert self.aibrix_kvcache.batch_set(keys, rand_tensor)
get_tensor = [
torch.rand(target_shape, dtype=torch.float16).flatten()
for _ in range(query_length)
]
self.aibrix_kvcache.batch_get(keys, get_tensor)
for i in range(query_length):
assert torch.equal(get_tensor[i], rand_tensor[i].flatten())
ret = self.aibrix_kvcache.batch_exists(keys)
assert self.aibrix_kvcache.batch_exists(keys) == query_length
assert self.aibrix_kvcache.batch_exists(partial_keys) == partial
partial_get_tensor = [
torch.rand(target_shape, dtype=torch.float16).flatten()
for _ in range(partial)
]
self.aibrix_kvcache.batch_get(partial_keys, partial_get_tensor)
for i in range(partial):
assert torch.equal(
partial_get_tensor[i], rand_tensor[i + partial].flatten()
)
log_every_n_seconds(
logger,
logging.INFO,
self.aibrix_kvcache.kv_cache_manager.metrics.summary(),
1,
)
if __name__ == "__main__":
setup()
test = AIBrixKVCacheStorageTest()
test.test_with_page_size()
@@ -0,0 +1,239 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to SGLang project
import importlib
import logging
from typing import TYPE_CHECKING, Any, Dict
from sglang.srt.mem_cache.hicache_storage import HiCacheStorage, HiCacheStorageConfig
if TYPE_CHECKING:
pass
logger = logging.getLogger(__name__)
class StorageBackendFactory:
"""Factory for creating storage backend instances with support for dynamic loading."""
_registry: Dict[str, Dict[str, Any]] = {}
@staticmethod
def _load_backend_class(
module_path: str, class_name: str, backend_name: str
) -> type[HiCacheStorage]:
"""Load and validate a backend class from module path."""
try:
module = importlib.import_module(module_path)
backend_class = getattr(module, class_name)
if not issubclass(backend_class, HiCacheStorage):
raise TypeError(
f"Backend class {class_name} must inherit from HiCacheStorage"
)
return backend_class
except ImportError as e:
raise ImportError(
f"Failed to import backend '{backend_name}' from '{module_path}': {e}"
) from e
except AttributeError as e:
raise AttributeError(
f"Class '{class_name}' not found in module '{module_path}': {e}"
) from e
@classmethod
def register_backend(cls, name: str, module_path: str, class_name: str) -> None:
"""Register a storage backend with lazy loading.
Args:
name: Backend identifier
module_path: Python module path containing the backend class
class_name: Name of the backend class
"""
if name in cls._registry:
logger.warning(f"Backend '{name}' is already registered, overwriting")
def loader() -> type[HiCacheStorage]:
"""Lazy loader function to import the backend class."""
return cls._load_backend_class(module_path, class_name, name)
cls._registry[name] = {
"loader": loader,
"module_path": module_path,
"class_name": class_name,
}
@classmethod
def create_backend(
cls,
backend_name: str,
storage_config: HiCacheStorageConfig,
mem_pool_host: Any,
**kwargs,
) -> HiCacheStorage:
"""Create a storage backend instance.
Args:
backend_name: Name of the backend to create
storage_config: Storage configuration
mem_pool_host: Memory pool host object
**kwargs: Additional arguments passed to external backends
Returns:
Initialized storage backend instance
Raises:
ValueError: If backend is not registered and cannot be dynamically loaded
ImportError: If backend module cannot be imported
Exception: If backend initialization fails
"""
# First check if backend is already registered
if backend_name in cls._registry:
registry_entry = cls._registry[backend_name]
backend_class = registry_entry["loader"]()
logger.info(
f"Creating storage backend '{backend_name}' "
f"({registry_entry['module_path']}.{registry_entry['class_name']})"
)
return cls._create_builtin_backend(
backend_name, backend_class, storage_config, mem_pool_host
)
# Try to dynamically load backend from extra_config
if backend_name == "dynamic" and storage_config.extra_config is not None:
backend_config = storage_config.extra_config
return cls._create_dynamic_backend(
backend_config, storage_config, mem_pool_host, **kwargs
)
# Backend not found
available_backends = list(cls._registry.keys())
raise ValueError(
f"Unknown storage backend '{backend_name}'. "
f"Registered backends: {available_backends}. "
)
@classmethod
def _create_dynamic_backend(
cls,
backend_config: Dict[str, Any],
storage_config: HiCacheStorageConfig,
mem_pool_host: Any,
**kwargs,
) -> HiCacheStorage:
"""Create a backend dynamically from configuration."""
required_fields = ["backend_name", "module_path", "class_name"]
for field in required_fields:
if field not in backend_config:
raise ValueError(
f"Missing required field '{field}' in backend config for 'dynamic' backend"
)
backend_name = backend_config["backend_name"]
module_path = backend_config["module_path"]
class_name = backend_config["class_name"]
try:
# Import the backend class
backend_class = cls._load_backend_class(
module_path, class_name, backend_name
)
logger.info(
f"Creating dynamic storage backend '{backend_name}' "
f"({module_path}.{class_name})"
)
# Create the backend instance with storage_config
return backend_class(storage_config, kwargs)
except Exception as e:
logger.error(
f"Failed to create dynamic storage backend '{backend_name}': {e}"
)
raise
@classmethod
def _create_builtin_backend(
cls,
backend_name: str,
backend_class: type[HiCacheStorage],
storage_config: HiCacheStorageConfig,
mem_pool_host: Any,
) -> HiCacheStorage:
"""Create built-in backend with original initialization logic."""
if backend_name == "file":
return backend_class(storage_config)
elif backend_name == "nixl":
return backend_class(storage_config)
elif backend_name == "mooncake":
backend = backend_class(storage_config, mem_pool_host)
return backend
elif backend_name == "aibrix":
backend = backend_class(storage_config, mem_pool_host)
return backend
elif backend_name == "hf3fs":
# Calculate bytes_per_page based on memory pool layout
if mem_pool_host.layout in ["page_first", "page_first_direct"]:
bytes_per_page = (
mem_pool_host.get_ksize_per_token() * mem_pool_host.page_size
)
elif mem_pool_host.layout == "layer_first":
bytes_per_page = (
mem_pool_host.get_size_per_token() * mem_pool_host.page_size
)
dtype = mem_pool_host.dtype
return backend_class.from_env_config(bytes_per_page, dtype, storage_config)
elif backend_name == "eic":
return backend_class(storage_config, mem_pool_host)
elif backend_name == "simm":
return backend_class(storage_config, mem_pool_host)
elif backend_name == "mori":
return backend_class(storage_config, mem_pool_host)
else:
raise ValueError(f"Unknown built-in backend: {backend_name}")
# Register built-in storage backends
StorageBackendFactory.register_backend(
"file", "sglang.srt.mem_cache.hicache_storage", "HiCacheFile"
)
StorageBackendFactory.register_backend(
"nixl",
"sglang.srt.mem_cache.storage.nixl.hicache_nixl",
"HiCacheNixl",
)
StorageBackendFactory.register_backend(
"mooncake",
"sglang.srt.mem_cache.storage.mooncake_store.mooncake_store",
"MooncakeStore",
)
StorageBackendFactory.register_backend(
"hf3fs",
"sglang.srt.mem_cache.storage.hf3fs.storage_hf3fs",
"HiCacheHF3FS",
)
StorageBackendFactory.register_backend(
"aibrix",
"sglang.srt.mem_cache.storage.aibrix_kvcache.aibrix_kvcache_storage",
"AibrixKVCacheStorage",
)
StorageBackendFactory.register_backend(
"eic",
"sglang.srt.mem_cache.storage.eic.eic_storage",
"EICStorage",
)
StorageBackendFactory.register_backend(
"simm",
"sglang.srt.mem_cache.storage.simm.hicache_simm",
"HiCacheSiMM",
)
StorageBackendFactory.register_backend(
"mori",
"sglang.srt.mem_cache.storage.umbp.umbp_store",
"UMBPStore",
)
@@ -0,0 +1,24 @@
# EIC as sglang HiCache Storage
EIC(Elastic Instant Cache) is a distributed database designed for LLM KV Cache. It supports RDMA, GDR and has the capabilities of distributed disaster tolerance and expansion.
You can understand the principles and architecture of EIC through these articles: https://mp.weixin.qq.com/s/tasDqXf0Gxr3o_WCJ2IJUQ https://mp.weixin.qq.com/s/b_4YhTa96Zeklh23lv8qBw
## Deploy EIC
You can visit the official link https://console.volcengine.com/eic and deploy EIC KVCache on your compute cluster with web UI.In addition, we provide particular image in volcano engine, which integrates various optimizations based on the official image.
You may use test_unit.py to detect the connectivity of EIC.
## Deploy Model With EIC
You can enable EIC KVCache offload with the official interface, such as
```bash
python -m sglang.launch_server \
--model-path [model_path]
--enable-hierarchical-cache \
--hicache-storage-backend eic \
--hicache-write-policy 'write_through' \
--hicache-mem-layout 'page_first' \
```
For more details, you can see https://www.volcengine.com/docs/85848/1749188 .
@@ -0,0 +1,778 @@
import json
import logging
import os
import time
from typing import Any, List, Optional, Tuple
import eic
import torch
import yaml
from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorage,
HiCacheStorageConfig,
HiCacheStorageExtraInfo,
)
from sglang.srt.mem_cache.pool_host import HostKVCache
logger = logging.getLogger(__name__)
TensorPoolSize = 2048
REMOTE_EIC_YAML_ENV_VAR = "REMOTE_EIC_YAML"
# gpu direct rdma for kv set
G_EnableKVSetGPUDirect = False
# gpu direct rdma for kv get
G_EnableKVGetGPUDirect = False
# gpu nic affinity
G_EnableGPUNicAffinity = False
# default H20 gpu nic affinity
GPUNicAffinity = {
"cuda:0": "eth1",
"cuda:1": "eth1",
"cuda:2": "eth2",
"cuda:3": "eth2",
"cuda:4": "eth3",
"cuda:5": "eth3",
"cuda:6": "eth4",
"cuda:7": "eth4",
}
# default H20 cpu nic affinity
CPUNicAffinity = {
"cuda:0": "cpu",
"cuda:1": "cpu",
"cuda:2": "cpu",
"cuda:3": "cpu",
"cuda:4": "cpu",
"cuda:5": "cpu",
"cuda:6": "cpu",
"cuda:7": "cpu",
}
def get_eic_config_file_path():
if os.environ.get(REMOTE_EIC_YAML_ENV_VAR) is not None:
logger.info(f"eic init with env var {REMOTE_EIC_YAML_ENV_VAR}")
config_file = os.environ.get(REMOTE_EIC_YAML_ENV_VAR)
else:
config_file = "/sgl-workspace/config/remote-eic.yaml"
logger.info(f"eic init with default config, config_file {config_file}")
return config_file
class FlexibleKVCacheMemoryPool:
def __init__(self, conn, kvcache_shape, kvcache_dtype, device):
self.connection = conn
if device.startswith("cpu") and G_EnableGPUNicAffinity:
gpu_id = torch.cuda.current_device()
self.device = CPUNicAffinity["cuda:" + str(gpu_id)]
# current memory pool size is 5 times of CPU TensorPoolSize
mempool_size = TensorPoolSize * 5
else:
self.device = device
mempool_size = TensorPoolSize
self.kvcache_shape = kvcache_shape
self.kvcache_dtype = kvcache_dtype
self.kv_cache_numel = 1
for i in self.kvcache_shape:
self.kv_cache_numel *= i
self.free_data_addr = set()
self.data_ptr_to_index = dict()
if self.device.startswith("cpu"):
self.kvcache_mempool = torch.zeros(
(mempool_size,) + kvcache_shape,
dtype=kvcache_dtype,
device=self.device,
pin_memory=True,
)
else:
self.kvcache_mempool = torch.zeros(
(mempool_size,) + kvcache_shape, dtype=kvcache_dtype, device=self.device
)
for i in range(mempool_size):
self.free_data_addr.add(i)
self.data_ptr_to_index[self.kvcache_mempool[i].data_ptr()] = i
meminfo = eic.MemoryInfo()
meminfo.type = eic.MemoryType.MEMORY_CUDA
meminfo.cuda_id = 0
vals = eic.IOBuffers()
vals.append(
self.kvcache_mempool.data_ptr(),
self.kvcache_mempool.numel() * self.kvcache_mempool.element_size(),
True,
)
self.connection.register_memory(vals, meminfo)
logger.info(
f"allocate memory pool, size {self.kvcache_mempool.numel() * self.kvcache_mempool.element_size()}, device {self.device}"
)
def try_allocate_kv_cache(self, shape, dtype, count=1):
if len(self.free_data_addr) < count:
return None
numel = 1
for i in shape:
numel *= i
if numel != self.kv_cache_numel or dtype != self.kvcache_dtype:
logger.error(
f"allocate from mempool failed, self.kvcache_shape {self.kvcache_shape}, dtype {self.kvcache_dtype}, require shape {shape}, dtype {dtype}"
)
return None
ret = []
for _ in range(count):
free_index = self.free_data_addr.pop()
ret.append(self.kvcache_mempool[free_index])
return ret
def free_to_mempool(self, data_ptr):
if data_ptr not in self.data_ptr_to_index:
logger.error(
f"free_to_mempool failed, data_ptr {data_ptr} not in allocated_data_addr"
)
return
self.free_data_addr.add(self.data_ptr_to_index[data_ptr])
def check_data_ptr_allocated(self, data_ptr):
return data_ptr in self.data_ptr_to_index
def left_count(self):
return len(self.free_data_addr)
class EICStorage(HiCacheStorage):
def __init__(
self, hicache_config: HiCacheStorageConfig, memory_pool_host: HostKVCache
):
global G_EnableKVSetGPUDirect, G_EnableKVGetGPUDirect
global GPUNicAffinity, CPUNicAffinity, G_EnableGPUNicAffinity
config_file = get_eic_config_file_path()
if os.path.exists(config_file) is False:
logger.error(f"config file {config_file} not exists")
raise RuntimeError(f"eic config file {config_file} not exists")
with open(config_file, "r") as fin:
config = yaml.safe_load(fin)
remote_url = config.get("remote_url", None)
if remote_url is None:
AssertionError("remote_url is None")
endpoint = remote_url[len("eic://") :]
logger.info(f"eic remote_url:" + remote_url + " endpoint: " + endpoint)
eic_instance_id = config.get("eic_instance_id", None)
logger.info(f"eic instance_id: {eic_instance_id}")
eic_thread_num = config.get("eic_thread_num", 1)
logger.info(f"eic thread_num: {eic_thread_num}")
eic_log_dir = config.get("eic_log_dir", None)
logger.info(f"eic log_dir: {eic_log_dir}")
eic_log_level = config.get("eic_log_level", 2)
logger.info(f"eic log_level: {eic_log_level}")
eic_trans_type = config.get("eic_trans_type", 3)
logger.info(f"eic trans_type: {eic_trans_type}")
eic_flag_file = config.get("eic_flag_file", None)
logger.info(f"eic flag_file: {eic_flag_file}")
# GDR now is not used
G_EnableKVSetGPUDirect = (
config.get("enable_kvset_gpu_direct", False) and torch.cuda.is_available()
)
logger.debug(f"eic enable_kvset_gpu_direct: {G_EnableKVSetGPUDirect}")
G_EnableKVGetGPUDirect = (
config.get("enable_kvget_gpu_direct", False) and torch.cuda.is_available()
)
logger.debug(f"eic enable_kvget_gpu_direct: {G_EnableKVGetGPUDirect}")
self.model_name = hicache_config.model_name
# rdma
enable_kv_set_direct = config.get("enable_kvset_direct", True)
logger.info(f"eic enable_kv_set_direct: {enable_kv_set_direct}")
self.enable_kv_set_direct = enable_kv_set_direct
enable_kv_get_direct = config.get("enable_kvget_direct", True)
logger.info(f"eic enable_kv_get_direct: {enable_kv_get_direct}")
self.enable_kv_get_direct = enable_kv_get_direct
# gpu nic affinity
G_EnableGPUNicAffinity = config.get("enable_gpu_nic_affinity", False)
logger.info(f"eic enable_gpu_nic_affinity: {G_EnableGPUNicAffinity}")
self.enable_gpu_nic_affinity = G_EnableGPUNicAffinity
if G_EnableGPUNicAffinity:
if "gpu_nic_affinity_config" in config:
GPUNicAffinity = json.loads(config["gpu_nic_affinity_config"])
if "cpu_nic_affinity_config" in config:
CPUNicAffinity = json.loads(config["cpu_nic_affinity_config"])
logger.info(f"eic gpu nic affinity {GPUNicAffinity}")
logger.info(f"eic cpu nic affinity {CPUNicAffinity}")
eic_namespace = config.get("eic_namespace", "")
logger.info(f"eic namespace: {eic_namespace}")
self.eic_namespace = eic_namespace
if not os.path.exists(eic_log_dir) and not os.path.isdir(eic_log_dir):
os.makedirs(eic_log_dir, exist_ok=True)
self.connection = eic.Client()
init_option = eic.InitOption()
init_option.log_dir = eic_log_dir
init_option.log_level = eic.LogLevel(eic_log_level)
init_option.transport_type = eic.TransportType(eic_trans_type)
init_option.flag_file = eic_flag_file
if G_EnableGPUNicAffinity:
gpu_id = torch.cuda.current_device()
init_option.multi_net_local_interface_names = GPUNicAffinity[
"cuda:" + str(gpu_id)
]
logger.info(
f"gpu {gpu_id} set gpu nic affinity to {init_option.multi_net_local_interface_names}"
)
ret = self.connection.init(eic_instance_id, endpoint, init_option)
if ret != 0:
logger.error(f"fail to init eic client, ret: {ret}")
raise RuntimeError("EIC Client Init Failed.")
self.warmup()
self.memory_pool_host = memory_pool_host
self.host_kvcache_layout = self.memory_pool_host.layout
self.trans_type = eic.TransportType(eic_trans_type)
self.kv_cache_dtype = self.memory_pool_host.dtype
self.is_mla_model = hicache_config.is_mla_model
self.rank = hicache_config.tp_rank
self.world_size = hicache_config.tp_size
self.page_size = self.memory_pool_host.page_size
self.use_zero_copy = self.memory_pool_host.layout == "page_first"
self.mha_zero_copy = self.use_zero_copy and not self.is_mla_model
if not self.use_zero_copy:
self.kv_cache_shape = self.memory_pool_host.get_data_page(
0, flat=True
).shape
if self.enable_kv_set_direct:
self.kv_cache_write_mem_pool = FlexibleKVCacheMemoryPool(
self.connection, self.kv_cache_shape, self.kv_cache_dtype, "cpu"
)
if self.enable_kv_get_direct:
self.kv_cache_get_mem_pool = FlexibleKVCacheMemoryPool(
self.connection, self.kv_cache_shape, self.kv_cache_dtype, "cpu"
)
self._init_eic_prefix()
def warmup(self):
logger.info("begin warm up eic client")
start_time = time.perf_counter()
num_warmup = 1024
preheat_keys = ["warmup_key_" + str(i) for i in range(num_warmup)]
batch_size = 32
for i in range(0, num_warmup, batch_size):
keys_vec = eic.StringVector()
for key in preheat_keys[i : i + batch_size]:
keys_vec.append(key)
exist_option = eic.ExistOption()
_, _ = self.connection.mexist(keys_vec, exist_option)
logger.info(
f"finish eic client warm up, warm up cost {time.perf_counter() - start_time:.2f} seconds"
)
def register_mem_pool_host(self, memory_pool_host: HostKVCache) -> None:
# no need judge meminfo type, cuda_id, etc.
meminfo = eic.MemoryInfo()
meminfo.type = eic.MemoryType.MEMORY_CUDA
meminfo.cuda_id = 0
vals = eic.IOBuffers()
buffer = memory_pool_host.kv_buffer
vals.append(
buffer.data_ptr(),
buffer.numel() * buffer.element_size(),
True,
)
self.connection.register_memory(vals, meminfo)
def _init_eic_prefix(self):
if self.is_mla_model:
self.eic_prefix = (
f"{self.model_name}_mla_att_{self.host_kvcache_layout}@sglang"
)
else:
self.eic_prefix = f"{self.model_name}_mha_attn_{self.host_kvcache_layout}_{self.rank}_{self.world_size}_@sglang"
def _get_eic_key(self, keys: List[str]) -> str:
return [f"{self.eic_prefix}_{key}" for key in keys]
def set(
self,
key: str,
value: Optional[Any] = None,
target_location: Optional[Any] = None,
target_size: Optional[Any] = None,
) -> bool:
# now is not used
if self.use_zero_copy:
return self.zero_copy_batch_set([key], [target_location])
else:
return self.generic_batch_set([key], [value])
# target_locations and target_sizes are not used for now
def batch_set(
self,
keys: List[str],
values: Optional[Any] = None,
target_locations: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> bool:
if len(keys) == 0:
return True
if self.use_zero_copy:
return self.zero_copy_batch_set(keys, values)
else:
return self.generic_batch_set(keys, values)
def get(
self,
key,
target_location: Optional[Any] = None,
target_size: Optional[Any] = None,
) -> torch.Tensor | None:
# now is not used
if self.use_zero_copy:
return self.zero_copy_batch_get([key], [target_location])
else:
return self.generic_batch_get([key], [target_location])
# use for v1 interface, and shound not be called directly
def batch_get(
self,
keys: List[str],
target_locations: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> List[torch.Tensor | None]:
assert len(keys) == len(target_locations)
if len(keys) == 0:
return None
if self.use_zero_copy:
return self.zero_copy_batch_get(keys, target_locations)
else:
return self.generic_batch_get(keys, target_locations)
def _batch_exists_impl(self, keys) -> List[bool]:
if len(keys) == 0:
return 0
eic_keys = self._get_eic_key(keys)
logger.debug(f"eic exists {len(keys)}")
result = []
exist_bs = 1024
for i in range(0, len(eic_keys), exist_bs):
batch_keys = eic_keys[i : i + exist_bs]
keys_vec = eic.StringVector()
for key in batch_keys:
keys_vec.append(key)
exist_option = eic.ExistOption()
exist_option.ns = self.eic_namespace
status_code, exist_outcome = self.connection.mexist(keys_vec, exist_option)
if status_code != eic.StatusCode.SUCCESS:
logger.error(
f"eic exists {len(keys)} failed, status_code {status_code}"
)
result.extend([False] * len(batch_keys))
for err_code in exist_outcome.status_codes:
result.append(err_code == eic.StatusCode.SUCCESS)
return result
def exists(self, key) -> bool:
exist_num = self.batch_exists([key])
return exist_num == 1
def batch_exists(
self, keys, extra_info: Optional[HiCacheStorageExtraInfo] = None
) -> int:
if len(keys) == 0:
return 0
if self.mha_zero_copy:
keys = self._get_mha_zero_copy_keys(keys)
exist_mask = self._batch_exists_impl(keys)
prefix_success = 0
for exist in exist_mask:
if exist:
prefix_success += 1
else:
break
if self.mha_zero_copy:
prefix_success = prefix_success // 2
return prefix_success
def delete(self, key) -> None:
eic_keys = self._get_eic_key([key])
keys_vec = eic.StringVector()
for eic_key in eic_keys:
keys_vec.append(eic_key)
del_option = eic.DelOption()
self.connection.mdel(keys_vec, del_option)
def clear(self) -> None:
return
# Not used for now
def _filter_kv_cache(self, total_len) -> Tuple[int, int]:
mean_len = total_len // self.world_size
remainder = total_len % self.world_size
tp_keys_len = mean_len + (1 if self.rank < remainder else 0)
start = self.rank * mean_len + min(self.rank, remainder)
end = start + tp_keys_len
logger.debug(f"start: {start}, end: {end}, tp_keys_len: {tp_keys_len}")
return start, end
def zero_copy_batch_set(self, keys: List[str], values: List[torch.Tensor]) -> bool:
logger.debug(f"eic zero copy set {len(keys)} keys")
if len(keys) == 0:
return True
eic_keys = self._get_eic_key(keys)
keys_vec = eic.StringVector()
vals_vec = eic.IOBuffers()
# set data key & value
for i, key in enumerate(eic_keys):
# set data key & value
keys_vec.append(key)
vals_vec.append(
values[i].data_ptr(),
values[i].element_size() * values[i].numel(),
True,
)
# set options
set_option = eic.SetOption()
set_option.ns = self.eic_namespace
set_option.ttl_second = -1
status_code, set_outcome = self.connection.mset(keys_vec, vals_vec, set_option)
if status_code != eic.StatusCode.SUCCESS:
logger.error(f"eic mset {len(keys)} failed, status_code {status_code}")
return [False] * len(keys)
else:
logger.debug(f"eic zero copy mset {len(keys)} success")
return [True] * len(keys)
def zero_copy_batch_get(
self, keys: List[str], values: List[torch.Tensor]
) -> List[bool]:
logger.debug(f"eic zero copy get {len(keys)} keys")
# Get Data: generate data keys and vals
get_data_start_time = time.perf_counter()
eic_keys = self._get_eic_key(keys)
data_keys = eic.StringVector()
data_vals = eic.IOBuffers()
success_mask = [True] * len(keys)
count = len(keys)
for i, key in enumerate(eic_keys):
data_keys.append(key)
data_vals.append(
values[i].data_ptr(),
values[i].element_size() * values[i].numel(),
True,
)
# Get data: recv data buffer tensor
get_option = eic.GetOption()
get_option.ns = self.eic_namespace
status_code, data_vals, get_outcome = self.connection.mget(
data_keys, get_option, data_vals
)
if status_code != eic.StatusCode.SUCCESS:
if status_code == eic.StatusCode.PARTIAL_FAILED:
for i, err_code in enumerate(get_outcome.status_codes):
success = err_code == eic.StatusCode.SUCCESS
if success:
logger.debug(f"eic get data {eic_keys[i]} success")
else:
logger.error(
f"eic get data {eic_keys[i]} failed, err_code {err_code}"
)
success_mask[i] = False
else:
logger.error(
f"eic mget {len(eic_keys)} keys failed, status_code {status_code}"
)
success_mask = [False] * len(keys)
return success_mask
get_data_end_time = time.perf_counter()
get_data_execution_time = (get_data_end_time - get_data_start_time) * 1e6
logger.debug(f"eic get {count} keys data cost %.2f us", get_data_execution_time)
return success_mask
def generic_batch_set(
self,
keys: List[str],
values: List[torch.Tensor],
) -> List[bool]:
assert len(keys) == len(values)
logger.debug(f"eic generic set {len(keys)} keys")
if len(keys) == 0:
return True
eic_keys = self._get_eic_key(keys)
keys_vec = eic.StringVector()
vals_vec = eic.IOBuffers()
count = len(keys)
registered = False
items = []
if self.enable_kv_set_direct:
values_data_ptrs = []
items = self.kv_cache_write_mem_pool.try_allocate_kv_cache(
self.kv_cache_shape, self.kv_cache_dtype, count
)
if items is None:
logger.warning("can not allocate tensor from pool")
for i, value in enumerate(values):
values_data_ptrs.append(
(value.data_ptr(), value.element_size() * value.numel(), False)
)
else:
objs = items
registered = True
for i, key in enumerate(eic_keys):
temp = objs[i].reshape(values[i].shape).contiguous()
temp.copy_(values[i])
if temp.data_ptr() != objs[i].data_ptr():
registered = False
temp = temp.cpu()
values_data_ptrs.append(
(
temp.data_ptr(),
temp.element_size() * temp.numel(),
registered,
)
)
for i, key in enumerate(eic_keys):
keys_vec.append(key)
data_ptr, data_size, registered = values_data_ptrs[i]
vals_vec.append(data_ptr, data_size, registered)
else:
# use tensor direct
for i, key in enumerate(eic_keys):
keys_vec.append(key)
vals_vec.append(
values[i].data_ptr(),
values[i].element_size() * values[i].numel(),
False,
)
# set options
set_option = eic.SetOption()
set_option.ns = self.eic_namespace
set_option.ttl_second = -1
status_code, set_outcome = self.connection.mset(keys_vec, vals_vec, set_option)
if status_code != eic.StatusCode.SUCCESS:
logger.error(f"eic mset {len(eic_keys)} failed, status_code {status_code}")
else:
logger.debug(f"eic mset {len(eic_keys)} success")
if self.enable_kv_set_direct and items is not None:
for item in items:
self.kv_cache_write_mem_pool.free_to_mempool(item.data_ptr())
err_code = set_outcome.status_codes[0]
if err_code != eic.StatusCode.SUCCESS:
logger.error(f"set data key {len(eic_keys)} failed, err_code {err_code}")
return [False] * len(keys)
logger.debug(f"set data key {len(eic_keys)} success")
return [True] * len(keys)
def generic_batch_get(
self, keys: List[str], buffers: List[torch.Tensor]
) -> List[bool]:
# all success or all fail
logger.debug(f"eic generic get {len(keys)} keys")
eic_keys = self._get_eic_key(keys)
get_data_start_time = time.perf_counter()
data_keys = eic.StringVector()
data_vals = eic.IOBuffers()
count = len(eic_keys)
registered = False
items = []
success_mask = [True] * len(keys)
if self.enable_kv_get_direct:
items = self.kv_cache_get_mem_pool.try_allocate_kv_cache(
self.kv_cache_shape, self.kv_cache_dtype, count
)
if items is None:
logger.warning("can not allocate tensor from pool")
for i, key in enumerate(eic_keys):
data_keys.append(key)
data_vals.append(
buffers[i].data_ptr(),
buffers[i].element_size() * buffers[i].numel(),
False,
)
else:
registered = True
for i, key in enumerate(eic_keys):
data_keys.append(key)
data_vals.append(
items[i].data_ptr(),
items[i].element_size() * items[i].numel(),
registered,
)
else:
for i, key in enumerate(eic_keys):
data_keys.append(key)
data_vals.append(
buffers[i].data_ptr(),
buffers[i].element_size() * buffers[i].numel(),
False,
)
# Get data: recv data buffer tensor
get_option = eic.GetOption()
get_option.ns = self.eic_namespace
status_code, data_vals, get_outcome = self.connection.mget(
data_keys, get_option, data_vals
)
if status_code != eic.StatusCode.SUCCESS:
if status_code == eic.StatusCode.PARTIAL_FAILED:
for i, err_code in enumerate(get_outcome.status_codes):
success = err_code == eic.StatusCode.SUCCESS
if success:
logger.debug(f"eic get data {eic_keys[i]} success")
else:
logger.error(
f"eic get data {eic_keys[i]} failed, err_code {err_code}"
)
success_mask[i] = False
else:
logger.error(
f"eic mget {len(eic_keys)} keys failed, status_code {status_code}"
)
success_mask = [False] * len(keys)
if registered:
for i, item in enumerate(items):
if success_mask[i]:
buffers[i].copy_(item)
self.kv_cache_get_mem_pool.free_to_mempool(item.data_ptr())
get_data_end_time = time.perf_counter()
get_data_execution_time = (get_data_end_time - get_data_start_time) * 1e6
logger.debug(f"eic get {count} keys data cost %.2f us", get_data_execution_time)
return success_mask
def _get_mha_zero_copy_keys(self, keys: List[str]) -> List[str]:
new_keys = []
for k in keys:
new_keys.append(f"{k}_k")
new_keys.append(f"{k}_v")
return new_keys
def _get_mha_zero_copy_values(
self, values: List[torch.Tensor]
) -> List[torch.Tensor]:
new_values = []
for value in values:
new_values.append(value[0])
new_values.append(value[1])
return new_values
def _batch_get_preprocess(self, keys, host_indices):
page_num = len(host_indices) // self.page_size
# use memory pool directly or dummy page
values = (
[
self.memory_pool_host.get_data_page(
host_indices[i * self.page_size], flat=False
)
for i in range(page_num)
]
if self.use_zero_copy
else [
self.memory_pool_host.get_dummy_flat_data_page()
for _ in range(page_num)
]
)
if self.mha_zero_copy:
keys = self._get_mha_zero_copy_keys(keys)
values = self._get_mha_zero_copy_values(values)
return keys, values
def _batch_get_postprocess(self, host_indices, values, results):
page_num = len(host_indices) // self.page_size
if self.use_zero_copy:
if not self.is_mla_model:
results = [
(results[2 * i] and results[2 * i + 1]) for i in range(page_num)
]
results = results[:page_num]
return results
# dummy page copy to host memory pool
for i in range(page_num):
if not results[i]:
break
self.memory_pool_host.set_from_flat_data_page(
host_indices[i * self.memory_pool_host.page_size], values[i]
)
return results
def batch_get_v1(
self,
keys: List[str],
host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]:
keys, values = self._batch_get_preprocess(keys, host_indices)
results = self.batch_get(keys, values)
return self._batch_get_postprocess(host_indices, values, results)
def _batch_set_preprocess(self, keys, host_indices):
page_num = len(host_indices) // self.page_size
flat = not self.use_zero_copy
values = [
self.memory_pool_host.get_data_page(
host_indices[i * self.page_size], flat=flat
)
for i in range(page_num)
]
if self.mha_zero_copy:
keys = self._get_mha_zero_copy_keys(keys)
values = self._get_mha_zero_copy_values(values)
return keys, values
def batch_set_v1(
self,
keys: List[str],
host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]:
keys, values = self._batch_set_preprocess(keys, host_indices)
results = self.batch_set(keys, values)
return results
@@ -0,0 +1,115 @@
import argparse
import os
import eic
import torch
import yaml
def pase_args():
parser = argparse.ArgumentParser(description="EIC Storage Unit Test")
parser.add_argument(
"--config",
"-c",
type=str,
default="/sgl-workspace/config/remote-eic.yaml",
help="EIC yaml config",
)
args, _ = parser.parse_known_args()
return args
def init_eic_client():
args = pase_args()
config_path = os.path.abspath(args.config)
if not os.path.exists(config_path):
raise FileNotFoundError(f"Config file not found: {config_path}")
with open(config_path, "r") as fin:
config = yaml.safe_load(fin)
remote_url = config.get("remote_url", None)
if remote_url is None:
AssertionError("remote_url is None")
endpoint = remote_url[len("eic://") :]
eic_instance_id = config.get("eic_instance_id", None)
eic_log_dir = config.get("eic_log_dir", None)
eic_log_level = config.get("eic_log_level", 2)
eic_trans_type = config.get("eic_trans_type", 3)
eic_flag_file = config.get("eic_flag_file", None)
if not os.path.exists(eic_log_dir):
os.makedirs(eic_log_dir, exist_ok=True)
eic_client = eic.Client()
init_option = eic.InitOption()
init_option.log_dir = eic_log_dir
init_option.log_level = eic.LogLevel(eic_log_level)
init_option.transport_type = eic.TransportType(eic_trans_type)
init_option.flag_file = eic_flag_file
ret = eic_client.init(eic_instance_id, endpoint, init_option)
if ret != 0:
raise RuntimeError(f"EIC Client init failed with error code: {ret}")
return eic_client
def test_set(eic_client):
test_key = ["test_key_" + str(i) for i in range(16)]
tensors = [
torch.ones([12, 6, 1, 512], dtype=torch.bfloat16, device="cpu")
for _ in range(16)
]
data_keys = eic.StringVector()
data_vals = eic.IOBuffers()
for i in range(16):
data_keys.append(test_key[i])
data_vals.append(
tensors[i].data_ptr(), tensors[i].numel() * tensors[i].element_size(), False
)
set_opt = eic.SetOption()
set_opt.ttl_second = 3
status_code, set_outcome = eic_client.mset(data_keys, data_vals, set_opt)
assert (
status_code == eic.StatusCode.SUCCESS
), f"Set failed with status code: {status_code}"
def test_get(eic_client):
test_key = ["test_key_" + str(i) for i in range(16)]
tensors = [
torch.zeros([12, 6, 1, 512], dtype=torch.bfloat16, device="cpu")
for _ in range(16)
]
data_keys = eic.StringVector()
data_vals = eic.IOBuffers()
for i in range(16):
data_keys.append(test_key[i])
data_vals.append(
tensors[i].data_ptr(), tensors[i].numel() * tensors[i].element_size(), False
)
get_opt = eic.GetOption()
status_code, data_vals, get_outcome = eic_client.mget(data_keys, get_opt, data_vals)
assert (
status_code == eic.StatusCode.SUCCESS
), f"Get failed with status code: {status_code}"
def test_exists(eic_client):
test_key = ["test_key_" + str(i) for i in range(16)]
data_keys = eic.StringVector()
for key in test_key:
data_keys.append(key)
exists_opt = eic.ExistOption()
status_code, exists_outcome = eic_client.mexist(data_keys, exists_opt)
assert (
status_code == eic.StatusCode.SUCCESS
), f"Exists failed with status code: {status_code}"
def main():
eic_client = init_eic_client()
test_set(eic_client)
test_exists(eic_client)
test_get(eic_client)
if __name__ == "__main__":
main()
@@ -0,0 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to SGLang project
"""File storage backend helpers for SGLang HiCache."""
from .lru_file_evictor import LRUFileEvictor
__all__ = [
"LRUFileEvictor",
]
@@ -0,0 +1,393 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to SGLang project
"""LRU/size-based file eviction for the HiCache file storage backend.
``HiCacheFile`` is a thin raw-bytes store: it suffixes keys, reads/writes
``.bin`` pages, and answers existence queries. Everything that bounds how much
disk those pages consume -- the LRU recency index, per-file size accounting,
free-space probing, scanning pre-existing files on startup, and unlinking
victims -- lives here so the backend stays a plain key/value store.
A backend constructs one evictor and drives it through a small lifecycle::
touch(key, path) # read hit / already-on-disk: bump recency
reserve(key, n_bytes) -> bool # admit a new write, evicting if needed
commit(key) # write landed on disk
abort(key) # write failed; release the reservation
clear() # backend wiped all files
When eviction is not configured the evictor is inert: ``reserve`` always admits
and the other calls are no-ops, so the backend behaves as unbounded storage.
"""
from __future__ import annotations
import argparse
import logging
import os
import threading
from collections import OrderedDict
from typing import Any, Callable, Optional, Set, Tuple
from sglang.srt.environ import envs
from sglang.srt.utils.common import human_readable_int
logger = logging.getLogger(__name__)
def _parse_size_to_bytes(value: Any) -> int:
"""Parse a size to bytes via human_readable_int (e.g. '200G', '1Gi', '1048576').
None / empty / '0' disables; an invalid value also disables (with a warning)."""
if value is None:
return 0
if isinstance(value, (int, float)):
return max(0, int(value))
s = str(value).strip()
if not s or s == "0":
return 0
try:
return max(0, human_readable_int(s))
except (argparse.ArgumentTypeError, ValueError):
logger.warning(f"Invalid size {value!r} for HiCacheFile; disabling.")
return 0
class LRUFileEvictor:
"""Bounds the on-disk size of a HiCacheFile directory via LRU eviction.
Tracks one ``.bin`` file per suffixed key (oldest at the front of the LRU),
enforces an optional byte cap and an optional free-space watermark, and
unlinks the least-recently-used files to stay within those bounds. Eviction
config comes from ``extra_config`` (per-backend, takes precedence) falling
back to the ``SGLANG_HICACHE_FILE_BACKEND_*`` env vars.
"""
def __init__(
self,
file_path: str,
config_suffix: str,
*,
tp_rank: int,
is_mla_model: bool,
extra_config: Optional[dict] = None,
on_evict: Optional[Callable[[str], None]] = None,
) -> None:
self.file_path = file_path
self.config_suffix = config_suffix
self._tp_rank = tp_rank
self._on_evict = on_evict
# MLA ranks share the same physical files, so centralize LRU bookkeeping
# on rank 0; non-MLA ranks each own their own files via the suffix.
self._is_storage_owner = (not is_mla_model) or (tp_rank == 0)
# suffixed_key -> file size in bytes; oldest at front.
self._lru: OrderedDict[str, int] = OrderedDict()
self._pending_writes: Set[str] = set()
self._total_bytes: int = 0
self._lock = threading.Lock()
self._load_config(extra_config or {})
self._eviction_configured = self.max_size_bytes > 0 or self.min_free_bytes > 0
self._eviction_enabled = self._eviction_configured and self._is_storage_owner
if self._eviction_configured and not self._is_storage_owner:
logger.info(
f"HiCacheFile rank {self._tp_rank} (MLA): eviction handled by rank 0; "
f"this rank skips LRU bookkeeping and will not create new files."
)
if not self._eviction_enabled:
return
# Clamp max_size to the filesystem capacity so a too-large cap can't OOM tmpfs.
fs = self._fs_stats()
if fs is not None and self.max_size_bytes > 0:
safe_max = max(0, fs[0] - self.min_free_bytes)
if self.max_size_bytes > safe_max:
logger.warning(
f"HiCacheFile max_size exceeds filesystem capacity; "
f"clamping to {safe_max} B."
)
self.max_size_bytes = safe_max
self._scan_existing_files()
with self._lock:
if self.max_size_bytes > 0 and self._total_bytes > self.max_size_bytes:
self._evict_locked(0)
if self.min_free_bytes > 0:
self._enforce_free_space_locked(0)
logger.info(
f"HiCacheFile eviction enabled: cap={self.max_size_bytes} B, "
f"watermark={self.eviction_ratio:.2f}, min_free={self.min_free_bytes} B, "
f"existing={self._total_bytes} B ({len(self._lru)} entries)"
)
def _load_config(self, extra: dict) -> None:
# extra_config (per-backend) takes precedence over env vars.
def _cfg(key, env):
val = extra.get(key)
return env.get() if val is None else val
self.max_size_bytes = _parse_size_to_bytes(
_cfg("max_size", envs.SGLANG_HICACHE_FILE_BACKEND_MAX_SIZE)
)
self.min_free_bytes = _parse_size_to_bytes(
_cfg("min_free_space", envs.SGLANG_HICACHE_FILE_BACKEND_MIN_FREE_SPACE)
)
ratio_raw = _cfg(
"eviction_ratio", envs.SGLANG_HICACHE_FILE_BACKEND_EVICTION_RATIO
)
try:
self.eviction_ratio = float(ratio_raw)
except (TypeError, ValueError):
self.eviction_ratio = 0.9
if not (0.0 < self.eviction_ratio <= 1.0):
self.eviction_ratio = 0.9
@property
def enabled(self) -> bool:
"""True when this rank actively evicts (configured AND storage owner)."""
return self._eviction_enabled
@property
def configured(self) -> bool:
"""True when a cap or free-space watermark is set (on any rank)."""
return self._eviction_configured
@property
def is_storage_owner(self) -> bool:
"""True when this rank owns (and may create/evict) the on-disk files."""
return self._is_storage_owner
def reserve(self, suffixed_key: str, value_bytes: int, key: str = "") -> bool:
"""Admit a new write of ``value_bytes``, evicting LRU victims as needed.
On success the key is pre-reserved at MRU and flagged in-flight so a
concurrent ``reserve`` won't evict it before the file is committed; the
caller must then call ``commit`` (write landed) or ``abort`` (write
failed). Returns ``False`` -- reserving nothing -- when the write is
refused: this rank is not the storage owner, the value is larger than
the cap, there is no evictable space, or the free-space watermark cannot
be met. When eviction is not configured the write is always admitted.
"""
if not self._eviction_configured:
return True # unbounded storage: nothing to enforce
if not self._is_storage_owner:
logger.warning(
f"HiCacheFile rank {self._tp_rank} is not the MLA storage owner; "
f"not caching new key {key} because file eviction is enabled."
)
return False
if self.max_size_bytes > 0 and value_bytes > self.max_size_bytes:
logger.warning(
f"HiCacheFile: value {value_bytes} B exceeds cap "
f"{self.max_size_bytes} B; not caching {key}"
)
return False
with self._lock:
# Cap-based eviction: evict, then bail if still over cap.
if (
self.max_size_bytes > 0
and (self._total_bytes + value_bytes) > self.max_size_bytes
):
self._evict_locked(value_bytes)
if (self._total_bytes + value_bytes) > self.max_size_bytes:
logger.warning(
f"HiCacheFile: no evictable space for {value_bytes} B "
f"under cap {self.max_size_bytes} B; not caching {key}"
)
return False
# Free-space watermark.
if self.min_free_bytes > 0 and not self._enforce_free_space_locked(
value_bytes
):
logger.warning(
f"HiCacheFile: filesystem hosting {self.file_path!r} "
f"would fall below min_free={self.min_free_bytes} B "
f"after writing {value_bytes} B; refusing {key} "
f"to avoid OOM/ENOSPC."
)
return False
# Pre-reserve at MRU so a concurrent evict won't grab this slot.
prev = self._lru.pop(suffixed_key, None)
if prev is not None:
self._total_bytes -= prev
self._lru[suffixed_key] = value_bytes
self._pending_writes.add(suffixed_key)
self._total_bytes += value_bytes
return True
def commit(self, suffixed_key: str) -> None:
"""Mark a reserved write as durably on disk (clears its in-flight flag)."""
if not self._eviction_enabled:
return
with self._lock:
self._pending_writes.discard(suffixed_key)
def abort(self, suffixed_key: str) -> None:
"""Release a reservation whose write failed: drop it and refund the bytes."""
if not self._eviction_enabled:
return
with self._lock:
cur = self._lru.pop(suffixed_key, None)
self._pending_writes.discard(suffixed_key)
if cur is not None:
self._total_bytes -= cur
def touch(self, suffixed_key: str, tensor_path: str) -> None:
"""Mark key as MRU, adopting an untracked on-disk file if needed."""
if not self._eviction_enabled:
return
with self._lock:
if suffixed_key in self._lru:
self._lru.move_to_end(suffixed_key, last=True)
return
# Untracked file: stat without holding the lock.
try:
size = os.path.getsize(tensor_path)
except OSError:
return
with self._lock:
if suffixed_key in self._lru:
self._lru.move_to_end(suffixed_key, last=True)
else:
self._lru[suffixed_key] = size
self._total_bytes += size
def clear(self) -> None:
"""Reset all bookkeeping after the backend has removed the files."""
with self._lock:
self._lru.clear()
self._pending_writes.clear()
self._total_bytes = 0
def _fs_stats(self) -> Optional[tuple]:
"""(total, available) bytes for the filesystem; None if unavailable."""
try:
st = os.statvfs(self.file_path)
except (OSError, AttributeError):
return None
total = st.f_blocks * st.f_frsize
free = st.f_bavail * st.f_frsize
return total, free
def _enforce_free_space_locked(self, value_bytes: int) -> bool:
"""Evict until writing value_bytes still leaves min_free_bytes free.
Caller holds _lock. Returns False if the write can't be satisfied."""
if self.min_free_bytes <= 0:
return True
fs = self._fs_stats()
if fs is None:
return True # cannot probe -> permissive, fall back to OS errors
# tmpfs frees space on unlink, so credit reclaimed bytes back to the
# estimate rather than re-probing statvfs on every eviction.
free = fs[1]
self._evict_while(
lambda reclaimed: (free + reclaimed) - value_bytes < self.min_free_bytes
)
# Re-probe: external writers may have changed free space meanwhile.
fs = self._fs_stats()
if fs is None:
return True
return fs[1] - value_bytes >= self.min_free_bytes
def _scan_existing_files(self) -> None:
"""Seed LRU index from disk on startup (oldest mtime first)."""
try:
names = os.listdir(self.file_path)
except FileNotFoundError:
return
entries = []
for fn in names:
if not fn.endswith(".bin"):
continue
stem = fn[:-4]
# Only files belonging to this rank/model.
if not stem.endswith(self.config_suffix):
continue
fp = os.path.join(self.file_path, fn)
try:
st = os.stat(fp)
except OSError:
continue
entries.append((st.st_mtime, stem, st.st_size))
entries.sort(key=lambda e: e[0]) # oldest first
for _, stem, size in entries:
self._lru[stem] = size
self._total_bytes += size
def _evict_one_lru_locked(self) -> Tuple[str, int]:
"""Evict the single oldest evictable LRU entry. Caller holds _lock.
The shared pop / skip-pending / unlink / ``_total_bytes`` step driven by
`_evict_while`. Returns ``(outcome, freed_bytes)``:
- ``("evicted", n)``: oldest entry dropped from the index; ``n`` disk
bytes reclaimed (0 if the file was already gone).
- ``("skipped", 0)``: oldest entry is an in-flight write; re-pinned at MRU
so the writer is not evicted out from under itself.
- ``("stop", 0)``: nothing evictable (empty index) or the unlink failed
(entry re-pinned at LRU); the caller should stop its eviction loop.
"""
if not self._lru:
return "stop", 0
evict_stem, evict_size = self._lru.popitem(last=False) # oldest
if evict_stem in self._pending_writes:
# Keep in-flight reservations; their file isn't committed yet.
self._lru[evict_stem] = evict_size
return "skipped", 0
tensor_path = os.path.join(self.file_path, f"{evict_stem}.bin")
try:
os.remove(tensor_path)
freed = evict_size
if self._on_evict is not None:
self._on_evict(evict_stem)
except FileNotFoundError:
freed = 0 # file already gone; still drop the stale index entry
if self._on_evict is not None:
self._on_evict(evict_stem)
except OSError as e:
logger.warning(f"HiCacheFile eviction failed for {evict_stem}: {e}")
self._lru[evict_stem] = evict_size
self._lru.move_to_end(evict_stem, last=False)
return "stop", 0
self._total_bytes -= evict_size
return "evicted", freed
def _evict_while(self, should_continue) -> int:
"""Evict oldest non-pending entries while ``should_continue(reclaimed)``.
``should_continue`` is passed the disk bytes reclaimed so far and returns
whether to keep evicting. In-flight writes are skipped; the loop is bounded
so it can't spin once every remaining entry is pending. Caller holds _lock.
Returns the total disk bytes reclaimed.
"""
reclaimed = 0
attempts_left = len(self._lru)
while self._lru and attempts_left > 0 and should_continue(reclaimed):
outcome, freed = self._evict_one_lru_locked()
if outcome == "stop":
break
if outcome == "skipped":
attempts_left -= 1
continue
# An entry left the index; reset the skip budget and bank the bytes.
reclaimed += freed
attempts_left = len(self._lru)
return reclaimed
def _evict_locked(self, needed_bytes: int) -> None:
"""Evict LRU entries until total + needed <= cap*ratio. Caller holds _lock."""
if self.max_size_bytes <= 0:
return
target = max(0, int(self.max_size_bytes * self.eviction_ratio) - needed_bytes)
reclaimed = self._evict_while(lambda _: self._total_bytes > target)
if reclaimed:
logger.debug(
f"HiCacheFile reclaimed {reclaimed} bytes; "
f"now {self._total_bytes} bytes used"
)
@@ -0,0 +1,427 @@
# FlexKV ↔ sglang integration
A `RadixCache` subclass that routes sglang's host-tier KV cache through a
FlexKV [`KVManager`](https://github.com/taco-project/FlexKV) (CPU / SSD /
Remote offload). Same integration pattern as
[`LMCRadixCache`](../lmcache/README.md): `FlexKVRadixCache` overrides
`match_prefix` / `init_load_back` / `cache_finished_req` / `evict`; a
`FlexKVConnector` façade talks to `KVManager`, `KVTPClient`, and a
3-axis (PP × CP × TP) sync context.
---
## Quick start (single H20, single GPU, Qwen3-8B)
This walks through everything the verification on H20-GPU-11 actually
exercised. Adjust paths / model / GPU as needed.
### 1. Prereqs
* `lmsysorg/sglang:dev` (or any sglang container with CUDA 12.x + torch 2.10+).
* This sglang fork (branch `feat/flexkv-main-connector`) and FlexKV
(branch `main`) checked out somewhere reachable from the container
— e.g. `/raid/fly/sglang-connector-dir/{sglang,FlexKV}`. Verified
against FlexKV main at `aa74e39` (PR #184); older commits down to
the layerwise integration also work.
### 2. Start a container with both repos mounted
```bash
docker run -d --name flexkv-sglang \
--gpus all --ipc=host --network host \
--shm-size=32g --cap-add SYS_NICE --cap-add IPC_LOCK \
-v /raid/fly:/raid/fly \
--workdir /raid/fly/sglang-connector-dir \
--entrypoint "" \
lmsysorg/sglang:dev sleep infinity
docker exec flexkv-sglang bash -c "
apt-get update -qq &&
apt-get install -y numactl libnuma-dev libxxhash-dev liburing-dev cmake ninja-build
"
```
### 3. Install sglang fork (editable) + FlexKV
```bash
docker exec flexkv-sglang bash -c '
set -e
git config --global --add safe.directory "*"
# sglang fork: install in editable mode, replacing the prebuilt sglang
cd /raid/fly/sglang-connector-dir/sglang
pip install --no-deps -e python
# FlexKV: pin to main, init the xxHash submodule, debug C++ build.
cd /raid/fly/sglang-connector-dir/FlexKV
git checkout main && git pull --ff-only
git submodule update --init third_party/xxHash
pip install -q cython ninja pybind11
FLEXKV_ENABLE_METRICS=0 bash build.sh --debug
# Smoke check
python3 -c "
import sglang, flexkv
from flexkv.kvmanager import KVManager
from sglang.srt.mem_cache.storage.flexkv import flexkv_comm
from sglang.srt.mem_cache.registry import registered_radix_cache_backends
import sglang.srt.mem_cache.storage.flexkv # registers
print(\"flexkv ok\", flexkv.__file__)
print(\"sglang ok\", sglang.__file__)
print(\"registered backends:\", registered_radix_cache_backends())
"
'
```
If the build hangs on `pip install sglang-kernel`, see
[Troubleshooting](#troubleshooting).
### 4. Minimal FlexKV YAML
```yaml
# /raid/fly/sglang-connector-dir/flexkv_min.yaml
cpu_cache_gb: 16
```
That's enough to enable a 16 GiB CPU offload pool. See
[`example_config_mp.yaml`](example_config_mp.yaml) for SSD / remote /
distributed knobs.
### 5. Launch the server (MP / synchronous mode)
```bash
docker exec -d flexkv-sglang bash -c '
cd /raid/fly/sglang-connector-dir
CUDA_VISIBLE_DEVICES=0 \
SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1 \
python3 -m sglang.launch_server \
--model-path /raid/fly/model/Qwen3-8B \
--port 30000 --tp-size 1 \
--enable-flexkv \
--flexkv-config-file /raid/fly/sglang-connector-dir/flexkv_min.yaml \
--mem-fraction-static 0.45 --max-running-requests 8 \
> /tmp/sglang.log 2>&1
'
```
`SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1` bypasses the prebuilt
`sglang-kernel` version assertion (the `lmsysorg/sglang:dev` image ships
0.4.2.post2; main expects ≥ 0.4.3). Not a FlexKV-specific issue;
remove when the container image is refreshed.
Wait ~2 min for the model load + CUDA graph capture. Confirm with:
```bash
docker exec flexkv-sglang bash -c '
grep -E "fired up|Connector ready" /tmp/sglang.log | tail -2
'
```
Expected (key lines):
```
[FlexKV] Connector ready ...: layerwise=False, prefetch=False
The server is fired up and ready to roll!
```
### 6. Send a request and observe a cache hit
```bash
# First call: priming — fresh prefill, FlexKV stores the prefix.
docker exec flexkv-sglang bash -c '
curl -s http://127.0.0.1:30000/generate -X POST \
-H "Content-Type: application/json" \
-d "{\"text\": \"The capital of France is\",
\"sampling_params\": {\"max_new_tokens\": 5, \"temperature\": 0}}"
'
# Flush the GPU radix (FlexKV CPU pool keeps the data) and re-send.
docker exec flexkv-sglang bash -c '
curl -s http://127.0.0.1:30000/flush_cache -X POST
curl -s http://127.0.0.1:30000/generate -X POST \
-H "Content-Type: application/json" \
-d "{\"text\": \"The capital of France is\",
\"sampling_params\": {\"max_new_tokens\": 5, \"temperature\": 0}}"
'
```
Look at the second response's `meta_info`:
```json
"cached_tokens": 4,
"cached_tokens_details": { "device": 0, "host": 4 },
```
`host: 4` confirms the bytes came back from FlexKV's CPU pool. The
server log should also show a matching D2H/H2D bandwidth line:
```
[FLEXKV] ... H2D transfer request: N finished transfer data size: 0.0xx GB ... 30+ GB/s
```
### 7. Layerwise mode
Add `FLEXKV_ENABLE_LAYERWISE_TRANSFER=1` before `python3 -m
sglang.launch_server`. Everything else is identical. On the second
request you'll see `cached_tokens_details: {"device": N, "host": 0}`
(in IP mode the load happens inside `match_prefix` so sglang accounts
for it as device-side) and a log line `LAYERWISE transfer request: N
finished ...`. The startup log will also include
`[FlexKV] Eventfd handshake complete ... counters=3 layers=<N>`.
---
## Correctness verification
Numerical match against a no-FlexKV baseline (greedy decoding,
deterministic). Scripts are in this repo's testing notes; the canonical
two are reproduced below.
```bash
# Phase 1: capture the no-FlexKV baseline.
docker exec -d flexkv-sglang bash -c '
CUDA_VISIBLE_DEVICES=0 SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1 \
python3 -m sglang.launch_server \
--model-path /raid/fly/model/Qwen3-8B --port 30000 --tp-size 1 \
--mem-fraction-static 0.45 > /tmp/sglang.log 2>&1
'
# ... wait until ready ...
docker exec flexkv-sglang python3 /raid/fly/sglang-connector-dir/sglang/python/sglang/srt/mem_cache/storage/flexkv/verify_outputs.py --phase baseline
docker exec flexkv-sglang bash -c "pkill -9 -f launch_server; sleep 3"
# Phase 2: relaunch with --enable-flexkv and compare.
docker exec -d flexkv-sglang bash -c '
CUDA_VISIBLE_DEVICES=0 SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1 \
python3 -m sglang.launch_server \
--model-path /raid/fly/model/Qwen3-8B --port 30000 --tp-size 1 \
--enable-flexkv --flexkv-config-file /raid/fly/sglang-connector-dir/flexkv_min.yaml \
--mem-fraction-static 0.45 > /tmp/sglang.log 2>&1
'
# ... wait until ready ...
docker exec flexkv-sglang python3 /raid/fly/sglang-connector-dir/sglang/python/sglang/srt/mem_cache/storage/flexkv/verify_outputs.py --phase test
```
Expected last line: `Total mismatches: 0`. Each prompt is run twice
(R1 fresh / R2 after `flush_cache`); both R1 and R2 outputs must
byte-equal the baseline.
Repeat the Phase-2 launch with `FLEXKV_ENABLE_LAYERWISE_TRANSFER=1`
to validate the layerwise path.
---
## Selecting the backend
Two equivalent CLI flags:
```bash
# Auto-selection chain (matches --enable-lmcache style)
python3 -m sglang.launch_server --enable-flexkv \
--flexkv-config-file /path/to/flexkv_config.yaml ...
# Explicit registry path
python3 -m sglang.launch_server --radix-cache-backend flexkv \
--flexkv-config-file /path/to/flexkv_config.yaml ...
```
Either flag also sets `FLEXKV_CONFIG_PATH` so you can omit
`--flexkv-config-file` and configure FlexKV purely through env vars.
---
## Modes
### MP (synchronous, default)
* `match_prefix` calls `FlexKVConnector.lookup_kv` only.
* When `host_hit_length > 0`, the scheduler later calls
`init_load_back`, which allocates the uncached slots and fires
`retrieve_kv` (FlexKV `launch` + `wait`).
* `cache_finished_req` runs `put_match` + `launch` and stashes the
in-flight FlexKV task id. Source-node lock is held until
`check_completed_stores` (called from `check_hicache_events` /
`evict`) signals completion.
This is the path you'll use under any non-trivial deployment topology
(DP > 1, multi-instance, multi-node, ...).
### IP / layerwise (`FLEXKV_ENABLE_LAYERWISE_TRANSFER=1`)
* `match_prefix` allocates the uncached slots and fires
`start_load_kv_layerwise` immediately.
* A `FlexKVLayerDoneCounter` is registered onto sglang's KV pool via
`register_layer_transfer_counter`; the per-layer hook blocks each
forward layer on its own eventfd until the FlexKV transfer worker
signals the layer is staged.
* Layerwise mode requires the FlexKV transfer worker's UDS socket
(`/tmp/flexkv_layerwise_eventfd.sock` by default) to be reachable —
the connector handshakes with it at startup. The socket path is
computed by FlexKV's `build_layerwise_eventfd_socket_path` from the
same dp/pp/instance settings, so configuration is taken care of as
long as you launch FlexKV consistently.
---
## Files
* `flexkv_radix_cache.py``FlexKVRadixCache(RadixCache)`. Overrides
`match_prefix`, `init_load_back`, `cache_finished_req`, `evict`,
`check_hicache_events`, `reset`.
* `flexkv_connector.py``FlexKVConnector`. Owns the `KVManager`,
`KVTPClient`, and the cross-rank sync context. Public methods:
`lookup_kv`, `retrieve_kv`, `start_load_kv_layerwise`, `store_kv`,
`check_completed_stores`, `prefetch_async`, …
* `flexkv_comm.py``FlexKVComm` (3-axis PP × CP × TP sync built on
torch.distributed) + the eventfd / `SCM_RIGHTS` shims used by the
layerwise transfer UDS handshake. **`FlexKVLayerLoadingEvent` here
carries the layerwise correctness fix** (drain stale eventfd
signals on reset, switch `wait` to `select.select` to keep blocking
semantics on a NONBLOCK fd).
* `__init__.py` — registers the `"flexkv"` factory with
`sglang.srt.mem_cache.registry`.
---
## TP / PP / CP / DP
FlexKV runs one `KVManager` per DP route (=
`instance_id * dp_size + dp_rank`). Every other rank in the same
fan-out is the "sync follower" — `FlexKVComm` broadcasts the
leader's lookup / store decisions via gloo CPU groups so non-leader
ranks know which task ids and slot mappings to use.
Supported:
* **TP** (any size) — typical sglang topology.
* **DP** (`dp_size > 1`) and multi-instance — FlexKV automatically
switches its `KVManager` to server-client mode.
* **PP** (`pp_size > 1`) — including cross-node PP. The PP receiver
forwards its slot mappings back to FlexKV's
`TransferManagerOnRemote` via the same ZMQ channel used for GPU
registration.
* **CP** (`attn_cp_size > 1`) — sync handled symmetrically with TP.
* **DP attention** (`enable_dp_attention=True`) — the inner
`attn_tp_size` is what FlexKV uses for register-side routing.
---
## Environment variables
* `FLEXKV_CONFIG_PATH` — full FlexKV YAML / JSON config (also set
automatically by `--flexkv-config-file`).
* `FLEXKV_ENABLE_LAYERWISE_TRANSFER``1` to enable layerwise mode.
* `FLEXKV_LAYERWISE_EVENTFD_SOCKET` — UDS socket path (default
`/tmp/flexkv_layerwise_eventfd.sock`); auto-suffixed per
`(pp_rank, dp_client_id)` when those dims are > 1.
* `FLEXKV_MASTER_HOST` / `FLEXKV_MASTER_PORTS` — multi-node master
endpoint for `TransferManagerOnRemote`. Default
`localhost:5556,5557,5558`. With `nnodes > 1` we also fall back to
`server_args.dist_init_addr`'s host.
* `FLEXKV_KV_CACHE_DTYPE` — override KV dtype when sglang uses
`--kv-cache-dtype auto`.
* `SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK` — bypass the prebuilt
`sglang-kernel` version assertion (not FlexKV-specific).
---
## Troubleshooting
* **`fatal: not a git repository ... third_party/xxHash`** — FlexKV's
build.sh needs an actual git checkout for the submodule. If you
rsync'd FlexKV without `.git/`, sync it: `rsync -az
/path/to/FlexKV/.git/ <remote>:<dir>/FlexKV/.git/` then
`git config --global --add safe.directory "*"`.
* **`fatal: detected dubious ownership`** — same fix:
`git config --global --add safe.directory "*"`.
* **`xxhash.h: No such file or directory`** — submodule not init'd.
`cd FlexKV && git submodule update --init third_party/xxHash`.
* **`dist/lease_meta_mempool.h: No such file or directory`** — your
rsync excluded `csrc/dist/`. The directory `FlexKV/csrc/dist/` is
source, not a build artifact; re-sync without `--exclude='dist'`.
* **`No module named 'Cython'`** — install: `pip install cython ninja pybind11`.
* **`sglang-kernel is installed with version 0.4.2.post2, which is
less than the minimum required version 0.4.3`** — either run with
`SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1` or refresh
`pip install -U sglang-kernel`. The download is ~600 MB and can
take a long time on slow links.
* **`cudaHostRegister failed with error code 100` (cudaErrorNoDevice)**
— happens when the FlexKV transfer subprocess can't init CUDA on
the assigned device. Usually a stuck previous session; restart
the container.
* **`[FlexKV] Waiting for FlexKV ready` loops > 60 s** — the
KVManager subprocess crashed at boot. Check `/tmp/sglang.log` for
the actual stack (usually a CUDA-init or torch-mp issue).
* **Layerwise mode: server hangs at "Eventfd connected attempts=..."**
— the `LayerwiseTransferWorker` hasn't started yet. Wait — it can
take 20-30 s after `Eventfd server created`. If it never advances,
check the FlexKV-side log lines beginning with `[LayerwiseWorker]`.
---
## Status
* MP (synchronous) path — verified end-to-end on Qwen3-8B (H20-3e):
output byte-equal to no-FlexKV baseline across short / medium / long
prompts. ~3046 GB/s observed for D2H stores and ~37 GB/s for H2D
loads.
* IP (layerwise) path — verified end-to-end with the fix in
`flexkv_comm.py`. ~712 GB/s per-layer (smaller per-call payload).
* PP / CP / DP / multi-node — code paths driven by `FlexKVComm`,
carried over from the production-validated `BaseKVConnector`
integration. Not exercised in single-GPU smoke tests; needs a
multi-node run before shipping.
### Known limitations
* Hybrid models (Mamba / SWA / DSV4 indexer auxiliary pools) are not
supported through this connector — only the primary KV pool is
hooked up. HiCache's multi-pool `batch_*_v2` interface would map
here but requires `PoolTransfer` + `PoolHitPolicy` plumbing in
`FlexKVConnector`.
* Write-back acks are per-request (one `dec_lock_ref` per
`cache_finished_req`), not per-page like HiCache's
`flush_write_through_acks`.
* `--radix-cache-backend=flexkv` and `--enable-flexkv` are
mutually equivalent today; we don't yet emit a deprecation
warning if both are set.
## Benchmarks
Setup: Qwen3-8B on 1× H20. Server flags:
--attention-backend triton --mem-fraction-static 0.32
--max-running-requests 32 --chunked-prefill-size 16384
--context-length 32000
Workload: 120 prompts sampled from
[`princeton-nlp/SWE-bench_Lite_oracle`](https://huggingface.co/datasets/princeton-nlp/SWE-bench_Lite_oracle)
with input length ≤ 28k tokens (p50 = 7088, max = 27961). Two passes —
pass 1 populates the host cache, pass 2 is the measured run. `qps=2.0`,
`concurrency=24`, `max_new_tokens=32`, `temperature=0`.
### Warm-pass results
| Config | TTFT avg / p50 / p90 / p99 | E2E p50 | Throughput | Output tok/s | H2D / D2H |
| --- | --- | --- | --- | --- | --- |
| baseline | 6.86 / 8.04 / 9.88 / 10.89 s | 8.15 s | 1.86 req/s | 37.7 | — |
| `--enable-hierarchical-cache` | **0.04 / 0.04 / 0.06 / 0.06 s** | 0.23 s | 2.02 req/s | 40.8 | — |
| `--enable-flexkv` | **0.05 / 0.05 / 0.07 / 0.08 s** | 0.24 s | 2.02 req/s | 40.8 | 86 / 155 |
Server-side (via `ReqTimeStats` in the sglang log): 76 / 76 non-EOS-immediate
warm-pass requests have `cached_input_len == input_len` for both `hicache`
and `flexkv` (100 % prefix recovery); baseline stays at ~59 tokens
(system-prompt header only). The 86 `H2D transfer` log lines under `flexkv`
confirm the CPU-tier loadbacks actually fired.
### Output correctness
Byte-level diff of generated text across 32 prompts, `temperature=0`:
* baseline: cold pass == warm pass (32 / 32; fully deterministic without cache)
* `hicache`: warm vs baseline warm — 29 / 32 identical, 3 diverge
* `flexkv`: warm vs baseline warm — 29 / 32 identical, 3 diverge (mostly the same 3 as `hicache`)
The ~10 % divergence at `temperature=0` is the well-known KV-cache-reuse
artifact caused by floating-point non-associativity between "prefill in
place" and "load pre-computed KV" paths; it affects the mainline
`--enable-hierarchical-cache` at the same rate and is not FlexKV-specific.
@@ -0,0 +1,87 @@
"""FlexKV-backed RadixCache integration for sglang.
Two ways to select this backend at server launch:
1. ``--enable-flexkv`` (default chain in ``default_radix_cache_factory``)
2. ``--radix-cache-backend=flexkv`` (explicit registry path)
Importing this package registers the explicit name with the registry,
so the second form is available without further wiring.
"""
from __future__ import annotations
import logging
from sglang.srt.mem_cache.registry import register_radix_cache_backend
logger = logging.getLogger(__name__)
def _flexkv_factory(ctx):
"""Build a :class:`FlexKVRadixCache` from a ``TreeCacheBuildContext``.
``TreeCacheBuildContext`` carries TP rank/size and the TP group
coordinator, but not PP/CP. We pick those up from the global
accessors in :mod:`sglang.srt.distributed.parallel_state`; FlexKV
needs them to fan out lookup/store decisions across the full TP × CP
× PP topology.
"""
from sglang.srt.distributed.parallel_state import (
get_attn_cp_group,
get_attn_tp_group,
get_pp_group,
)
from sglang.srt.mem_cache.storage.flexkv.flexkv_radix_cache import (
FlexKVRadixCache,
)
server_args = ctx.server_args
# PP group is always available; attn TP / attn CP groups may share
# the regular TP group when attn DP is off — that's fine, the
# connector treats size-1 groups as no-ops.
try:
pp_group = get_pp_group()
except (RuntimeError, AssertionError):
pp_group = None
try:
attn_tp_group = get_attn_tp_group()
except (RuntimeError, AssertionError):
attn_tp_group = ctx.tp_group
try:
attn_cp_group = get_attn_cp_group()
except (RuntimeError, AssertionError):
attn_cp_group = None
# PP / CP ranks: use the group's own rank_in_group view if available;
# fall back to 0 for single-rank dims.
pp_rank = pp_group.rank_in_group if pp_group is not None else 0
attn_cp_rank = attn_cp_group.rank_in_group if attn_cp_group is not None else 0
return FlexKVRadixCache(
params=ctx.params,
model_config=ctx.model_config,
server_args=server_args,
tp_rank=ctx.tp_rank,
tp_size=ctx.tp_size,
# ``dp_rank`` isn't carried on TreeCacheBuildContext or ServerArgs
# at construction time; the connector normalizes ``None`` to 0
# for the single-DP-rank case that this factory targets.
dp_rank=None,
pp_rank=pp_rank,
attn_cp_rank=attn_cp_rank,
tp_group=ctx.tp_group,
pp_group=pp_group,
attn_tp_group=attn_tp_group,
attn_cp_group=attn_cp_group,
)
try:
register_radix_cache_backend("flexkv", _flexkv_factory)
except ValueError as exc:
# The registry refuses duplicates. Importing this package twice
# (e.g. via both --enable-flexkv and --radix-cache-backend=flexkv)
# is fine — log and move on.
logger.debug("flexkv backend already registered: %s", exc)
@@ -0,0 +1,38 @@
# Example FlexKV YAML config (passed to sglang via --flexkv-config-file).
#
# Equivalent env vars exist for every field — see flexkv/common/config.py
# (UserConfig.from_env). This file is a minimal CPU-only setup; uncomment
# the SSD / Remote / Redis sections to enable those tiers.
# ---- CPU host-side cache ----------------------------------------------
# Size of the FlexKV CPU pool. Used to derive `num_cpu_blocks` together
# with the model dtype, head dim, num kv heads, and page size.
cpu_cache_gb: 64
# Optional: pin the CPU pool using transparent huge pages.
# use_hugepage_cpu_buffer: false
# use_hugepage_tmp_buffer: false
# hugepage_size_bytes: 2097152
# ---- SSD tier ---------------------------------------------------------
# Set ssd_cache_gb > cpu_cache_gb to enable the SSD spill tier.
# ssd_cache_gb: 256
# ssd_cache_dir: "/mnt/nvme0/flexkv;/mnt/nvme1/flexkv" # ';'-separated for striping
# enable_gds: false # cuFile / GDS path
# ---- KV cache dtype override -----------------------------------------
# When sglang is launched with --kv-cache-dtype auto, FlexKV can't tell
# which dtype the actual KV tensors use. Set explicitly here.
# kv_cache_dtype: bfloat16
# ---- Peer / distributed sharing --------------------------------------
# enable_p2p_cpu: false
# enable_p2p_ssd: false
# enable_3rd_remote: false
# ---- Redis (for distributed metadata / KV sharing) -------------------
# redis_host: 127.0.0.1
# redis_port: 6379
# redis_password: null
# node_ttl_seconds: 60
# local_ip: 10.0.0.1
@@ -0,0 +1,662 @@
"""Communication helpers for the FlexKV connector.
FlexKV runs a single KVManager per DP group (typically the TP/CP/PP
sync leader's process). Every other rank in the same KV-cache-sharing
fan-out must be told the leader's decisions: which prefix matched in
FlexKV, which task id the leader allocated, which slot mappings to
send, etc.
This file provides:
* ``FlexKVComm`` — a 3-axis (PP × CP × TP) hierarchical sync context
built on torch.distributed (gloo CPU groups). Exposes ``scatter``,
``scatter_pp``, ``barrier`` and ``all_reduce_min`` plus role flags
(``is_sync_leader`` etc.) that the connector branches on.
* libc / ``eventfd`` shims used by the layerwise transfer worker
socket handshake.
* ``FlexKVLayerLoadingEvent`` and ``FlexKVLayerDoneCounter`` — the
eventfd-backed per-layer completion structures that the FlexKV
layerwise transfer worker signals into. Hooked into sglang's
``register_layer_transfer_counter`` so each layer's forward waits
for its own host→device copy.
"""
from __future__ import annotations
import ctypes
import errno
import logging
import os
import pickle
import socket
import struct
from datetime import timedelta
from typing import Any, Dict, List
import torch
import torch.distributed as dist
from sglang.srt.distributed.parallel_state import get_world_group
logger = logging.getLogger(__name__)
# PP-channel command tags (used by ``scatter_pp`` payloads). Sender and
# receiver assert on these to catch protocol drift early.
CMD_PUT_META = 2
CMD_LAYERWISE = 3
CMD_STORE_COMPLETE = 5
class FlexKVComm:
"""3-axis (PP × CP × TP) hierarchical sync for the FlexKV connector.
Notation:
* "sync leader" is the unique rank that talks to the FlexKV
KVManager: pp_rank=0, attn_cp_rank=0, attn_tp_rank=0.
* "PP stage leader" is the (cp=0, tp=0) rank within a PP stage —
it does cross-PP P2P (``scatter_pp``).
* Every rank participates in collective layers it belongs to.
Communication strategy:
* P2P (send/recv/isend/irecv) on CPU tensors → ``world_cpu_group``
(the global gloo group). Sub-group cpu_groups have unreliable
TCP pairs for direct P2P.
* Collectives (all_reduce / barrier) → sglang's sub-group
cpu_groups (fine for collectives).
"""
# P2P tags. World group is shared with sglang's own P2P, so we pick
# 4-byte tags that won't collide.
_TAG_SCATTER = int.from_bytes(b"FxSc", byteorder="big")
_TAG_PP = int.from_bytes(b"FxPP", byteorder="big")
_TAG_CP = int.from_bytes(b"FxCP", byteorder="big")
_TAG_TP = int.from_bytes(b"FxTP", byteorder="big")
_TAG_PP_AR_MIN = int.from_bytes(b"FxA2", byteorder="big")
_TAG_PP_BARRIER = int.from_bytes(b"FxB2", byteorder="big")
_TAG_PP_BARRIER_BCAST = int.from_bytes(b"FxB3", byteorder="big")
_TAG_AR_BCAST = int.from_bytes(b"FxAR", byteorder="big")
# Adaptive async-work reaper. gloo's isend Work objects do not auto-
# advance their "completed" state on poll, so a pure poll-based reaper
# leaks. We actively wait() the oldest works with a tiny timeout;
# the watermark grows on stuck reaps (slow / asymmetric peer) and
# shrinks back on clean reaps.
_REAP_HIGH_BASE = 1024
_REAP_HIGH_MAX = 32768
_REAP_MAX_DRAIN = 512
_REAP_PROBE = timedelta(milliseconds=1)
_REAP_LOG_EVERY = 64
def __init__(
self,
rank_info,
world_rank: int,
pp_group=None,
attn_tp_group=None,
attn_cp_group=None,
):
model_config = rank_info.model_config
self.world_rank = world_rank
self._async_works: List = []
self._reap_high: int = self._REAP_HIGH_BASE
self._reap_calls: int = 0
self._reap_stuck_total: int = 0
self._reap_drained_total: int = 0
# Accept either GroupCoordinator wrappers (has ``.cpu_group``) or
# raw ProcessGroups.
self.pp_cpu_group = (
getattr(pp_group, "cpu_group", pp_group) if pp_group is not None else None
)
self.attn_tp_cpu_group = (
getattr(attn_tp_group, "cpu_group", attn_tp_group)
if attn_tp_group is not None
else None
)
self.attn_cp_cpu_group = (
getattr(attn_cp_group, "cpu_group", attn_cp_group)
if attn_cp_group is not None
else None
)
self.pp_size = model_config.pp_size
self.attn_tp_size = model_config.attn_tp_size
self.attn_cp_size = model_config.attn_cp_size
self.pp_rank = rank_info.pp_rank
self.attn_tp_rank = rank_info.attn_tp_rank
self.attn_cp_rank = rank_info.attn_cp_rank
self.is_pp_stage_leader = self.attn_tp_rank == 0 and self.attn_cp_rank == 0
self.is_sync_leader = self.pp_rank == 0 and self.is_pp_stage_leader
self.is_pp_leader = self.pp_rank == 0 and self.is_pp_stage_leader
self.is_cp_leader = self.attn_cp_rank == 0
self.is_tp_leader = self.attn_tp_rank == 0
# P2P routing tables (computed once).
stride = self.attn_tp_size * self.attn_cp_size
self._pp_stage_leader_ranks = [s * stride for s in range(self.pp_size)]
pp_stage_offset = self.pp_rank * stride
self._cp_leader_ranks = (
[
pp_stage_offset + cp * self.attn_tp_size
for cp in range(self.attn_cp_size)
]
if self.attn_cp_size > 1
else []
)
if self.attn_tp_size > 1:
if self.attn_tp_cpu_group is None:
raise RuntimeError(
f"[FlexKV] attn_tp_size={self.attn_tp_size} > 1 but "
f"attn_tp_cpu_group is None — TP CPU group is required "
f"for scatter/collectives."
)
self._tp_group_ranks = [
dist.get_global_rank(self.attn_tp_cpu_group, i)
for i in range(self.attn_tp_cpu_group.size())
]
else:
self._tp_group_ranks = []
self._pp_group_global_ranks = (
[
dist.get_global_rank(self.pp_cpu_group, i)
for i in range(self.pp_cpu_group.size())
]
if self.pp_size > 1 and self.pp_cpu_group is not None
else []
)
self._pp_stage_member_ranks = list(
range(pp_stage_offset, pp_stage_offset + stride)
)
self.needs_sync = (
self.pp_size > 1 or self.attn_tp_size > 1 or self.attn_cp_size > 1
)
self._world_cpu_group = get_world_group().cpu_group
self.pp_group = (
self.pp_cpu_group
if (self.pp_size > 1 and self.is_pp_stage_leader)
else None
)
self.is_pp_active = self.pp_size > 1
self.is_pp_sender = self.is_pp_leader
self.is_pp_receiver = self.is_pp_stage_leader and not self.is_pp_leader
self.is_cross_node_pp = self.pp_size > rank_info.pp_size_per_node
self.should_send_slot_mapping_to_remote = (
self.is_pp_receiver and self.is_cross_node_pp
)
logger.info(
"[FlexKV] Comm init: rank=%d, pp=%d/%d, tp=%d/%d, cp=%d/%d, "
"sync_leader=%s, stage_leader=%s, cross_node_pp=%s",
world_rank,
self.pp_rank,
self.pp_size,
self.attn_tp_rank,
self.attn_tp_size,
self.attn_cp_rank,
self.attn_cp_size,
self.is_sync_leader,
self.is_pp_stage_leader,
self.is_cross_node_pp,
)
# ------------------------------------------------------------------
# Public collectives
# ------------------------------------------------------------------
def scatter(self, data: Any, blocking: bool = False) -> Any:
"""Hierarchical fan-out: sync_leader → PP stage leaders →
CP leaders → TP ranks. Returns the leader's payload on every rank.
``blocking=False`` queues isends and reaps later — fine for the
hot path; ``True`` blocks until the leader's sends drain (used
on shutdown / barriers).
"""
if self.pp_size > 1 and self.is_pp_stage_leader:
data = self._scatter_group(
data,
self._pp_stage_leader_ranks,
self.is_pp_leader,
self._TAG_PP,
blocking,
)
if self._cp_leader_ranks:
data = self._scatter_group(
data,
self._cp_leader_ranks,
self.is_cp_leader,
self._TAG_CP,
blocking,
)
if self._tp_group_ranks:
data = self._scatter_group(
data,
self._tp_group_ranks,
self.is_tp_leader,
self._TAG_TP,
blocking,
)
return data
def scatter_pp(self, data: Any) -> Any:
"""PP-only fan-out across PP stages (only stage leaders participate)."""
if not self._pp_group_global_ranks:
return data
is_leader = self._pp_group_global_ranks[0] == self.world_rank
return self._scatter_group(
data,
self._pp_group_global_ranks,
is_leader,
self._TAG_SCATTER,
blocking=False,
)
def all_reduce_min(self, value: int) -> int:
"""Hierarchical all_reduce(MIN) across TP, CP, PP.
Used to align FlexKV block-count limits across all ranks that
will register GPU buffers (each rank computes the maximum it can
support, and we take the MIN to land on a value everyone can
honor).
"""
tensor = torch.tensor(value, dtype=torch.int64)
if self.attn_tp_size > 1 and self.attn_tp_cpu_group is not None:
dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=self.attn_tp_cpu_group)
if self.attn_cp_size > 1 and self.attn_cp_cpu_group is not None:
dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=self.attn_cp_cpu_group)
if self.pp_size > 1 and self.is_pp_stage_leader:
self._pp_all_reduce_min_p2p(tensor)
if self.pp_size > 1:
self._bcast_to_stage_members(tensor, self._TAG_AR_BCAST)
return int(tensor.item())
def barrier(self) -> None:
if self.attn_tp_size > 1 and self.attn_tp_cpu_group is not None:
dist.barrier(group=self.attn_tp_cpu_group)
if self.attn_cp_size > 1 and self.attn_cp_cpu_group is not None:
dist.barrier(group=self.attn_cp_cpu_group)
if self.pp_size > 1 and self.is_pp_stage_leader:
self._pp_barrier_p2p()
if self.pp_size > 1:
dummy = torch.tensor([0], dtype=torch.int64)
self._bcast_to_stage_members(dummy, self._TAG_PP_BARRIER_BCAST)
# ------------------------------------------------------------------
# Internal scatter helper
# ------------------------------------------------------------------
def _scatter_group(
self,
data: Any,
group_ranks: List[int],
is_leader: bool,
tag: int,
blocking: bool = False,
) -> Any:
if not group_ranks or self.world_rank not in group_ranks:
return data
if is_leader:
dsts = [r for r in group_ranks if r != self.world_rank]
works = []
for dst in dsts:
works.extend(self._isend(dst, data, tag, self._world_cpu_group))
if blocking:
for w in works:
w.wait()
else:
self._reap_completed_async_works()
self._async_works.extend(works)
return data
return self._recv(group_ranks[0], tag, self._world_cpu_group)
def _reap_completed_async_works(self) -> None:
n = len(self._async_works)
if n <= self._reap_high:
return
drained = 0
stuck = False
for _ in range(self._REAP_MAX_DRAIN):
if not self._async_works:
break
w = self._async_works[0]
try:
w.wait(self._REAP_PROBE)
except RuntimeError:
stuck = True
break
self._async_works.pop(0)
drained += 1
self._reap_calls += 1
self._reap_drained_total += drained
if stuck:
self._reap_stuck_total += 1
prev_high = self._reap_high
if stuck:
self._reap_high = min(self._REAP_HIGH_MAX, self._reap_high * 2)
else:
self._reap_high = max(self._REAP_HIGH_BASE, self._reap_high // 2)
if self._reap_high != prev_high:
logger.debug(
"[FlexKV] reap watermark rank=%d %d->%d "
"(stuck=%s drained=%d backlog=%d)",
self.world_rank,
prev_high,
self._reap_high,
stuck,
drained,
n,
)
if self._reap_calls % self._REAP_LOG_EVERY == 0:
logger.debug(
"[FlexKV] reap stats rank=%d calls=%d drained=%d stuck=%d "
"backlog=%d high=%d",
self.world_rank,
self._reap_calls,
self._reap_drained_total,
self._reap_stuck_total,
len(self._async_works),
self._reap_high,
)
# ------------------------------------------------------------------
# Low-level send / recv on the world cpu group
# ------------------------------------------------------------------
def _isend(self, dst: int, data: Any, tag: int = 0, group=None) -> list:
serialized = bytearray(pickle.dumps(data))
t_size = torch.tensor([len(serialized)], dtype=torch.long)
t_data = torch.frombuffer(serialized, dtype=torch.uint8)
return [
dist.isend(t_size, dst=dst, tag=tag, group=group),
dist.isend(t_data, dst=dst, tag=tag, group=group),
]
def _recv(self, src: int, tag: int = 0, group=None) -> Any:
t_size = torch.tensor([0], dtype=torch.long)
dist.irecv(t_size, src=src, tag=tag, group=group).wait()
size = int(t_size.item())
if size == 0:
return []
t_data = torch.empty(size, dtype=torch.uint8)
dist.irecv(t_data, src=src, tag=tag, group=group).wait()
return pickle.loads(t_data.numpy().tobytes())
def _send_tensor(
self, tensor: torch.Tensor, dst: int, tag: int = 0, group=None
) -> None:
dist.send(tensor, dst=dst, tag=tag, group=group)
def _recv_tensor(
self, tensor: torch.Tensor, src: int, tag: int = 0, group=None
) -> None:
dist.recv(tensor, src=src, tag=tag, group=group)
def _bcast_to_stage_members(self, tensor: torch.Tensor, tag: int) -> None:
if not self.is_pp_stage_leader:
self._recv_tensor(
tensor,
src=self._pp_stage_leader_ranks[self.pp_rank],
tag=tag,
group=self._world_cpu_group,
)
return
for rank in self._pp_stage_member_ranks:
if rank != self.world_rank:
self._send_tensor(
tensor, dst=rank, tag=tag, group=self._world_cpu_group
)
def _pp_all_reduce_min_p2p(self, tensor: torch.Tensor) -> None:
leader_rank = self._pp_stage_leader_ranks[0]
other_leaders = self._pp_stage_leader_ranks[1:]
tag = self._TAG_PP_AR_MIN
if self.world_rank == leader_rank:
result = int(tensor.item())
for src in other_leaders:
other = torch.tensor(0, dtype=torch.int64)
self._recv_tensor(other, src=src, tag=tag, group=self._world_cpu_group)
result = min(result, int(other.item()))
tensor.fill_(result)
for dst in other_leaders:
self._send_tensor(tensor, dst=dst, tag=tag, group=self._world_cpu_group)
else:
self._send_tensor(
tensor, dst=leader_rank, tag=tag, group=self._world_cpu_group
)
self._recv_tensor(
tensor, src=leader_rank, tag=tag, group=self._world_cpu_group
)
def _pp_barrier_p2p(self) -> None:
leader_rank = self._pp_stage_leader_ranks[0]
other_leaders = self._pp_stage_leader_ranks[1:]
tag = self._TAG_PP_BARRIER
dummy = torch.tensor([1], dtype=torch.int64)
if self.world_rank == leader_rank:
for src in other_leaders:
self._recv_tensor(dummy, src=src, tag=tag, group=self._world_cpu_group)
for dst in other_leaders:
self._send_tensor(dummy, dst=dst, tag=tag, group=self._world_cpu_group)
else:
self._send_tensor(
dummy, dst=leader_rank, tag=tag, group=self._world_cpu_group
)
self._recv_tensor(
dummy, src=leader_rank, tag=tag, group=self._world_cpu_group
)
# ----------------------------------------------------------------------
# libc / eventfd / SCM_RIGHTS shims for the layerwise UDS handshake
# ----------------------------------------------------------------------
_libc = ctypes.CDLL("libc.so.6", use_errno=True)
_libc.eventfd.argtypes = [ctypes.c_uint, ctypes.c_int]
_libc.eventfd.restype = ctypes.c_int
_libc.read.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_size_t]
_libc.read.restype = ctypes.c_ssize_t
_libc.write.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_size_t]
_libc.write.restype = ctypes.c_ssize_t
EFD_SEMAPHORE = 0x1
EFD_NONBLOCK = 0x800
def eventfd(initval: int = 0, flags: int = 0) -> int:
fd = _libc.eventfd(ctypes.c_uint(initval), ctypes.c_int(flags))
if fd == -1:
err = ctypes.get_errno()
raise OSError(err, os.strerror(err))
return fd
def eventfd_write(fd: int, val: int) -> None:
v = ctypes.c_uint64(val)
n = _libc.write(fd, ctypes.byref(v), ctypes.sizeof(v))
if n != ctypes.sizeof(v):
err = ctypes.get_errno()
raise OSError(err, f"eventfd write failed: {os.strerror(err)}")
def eventfd_read(fd: int) -> int:
v = ctypes.c_uint64()
n = _libc.read(fd, ctypes.byref(v), ctypes.sizeof(v))
if n != ctypes.sizeof(v):
err = ctypes.get_errno()
if err == errno.EAGAIN:
return 0
raise OSError(err, f"eventfd read failed: {os.strerror(err)}")
return v.value
def send_fds(sock: socket.socket, fds: list, extra_data: bytes = b"x") -> None:
"""SCM_RIGHTS-send a list of file descriptors over a UDS socket."""
fds_packed = struct.pack(f"{len(fds)}i", *fds)
ancdata = [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds_packed)]
sock.sendmsg([extra_data], ancdata)
# ----------------------------------------------------------------------
# Layerwise transfer signaling (eventfd-backed)
# ----------------------------------------------------------------------
class FlexKVLayerLoadingEvent:
"""One per producer slot. Holds ``num_layers`` semaphore eventfds —
the FlexKV layerwise worker writes 1 to each as the corresponding
layer's H2D copy completes; the consumer (sglang's
``register_layer_transfer_counter`` hook) reads to wait for them."""
def __init__(self, num_layers: int):
self._num_layers = num_layers
# Semaphore mode so each read consumes exactly one signal. NONBLOCK
# lets ``reset_for_new_transfer`` drain leftover counter values
# without blocking; ``wait`` re-arms the fd to blocking before
# reading so consumers still get the desired blocking semantics.
self.load_event_fds: List[int] = [
eventfd(0, EFD_SEMAPHORE | EFD_NONBLOCK) for _ in range(num_layers)
]
self._finished = True
self.wait_remaining: List[int] = [1] * num_layers
def reset_for_new_transfer(self) -> None:
"""Drain any leftover signals from prior transfers, then arm.
Without this drain, a previous transfer that wrote N eventfd
signals but only had N-K reads (e.g. because the attention
backend skipped a layer's ``get_key_buffer`` call) leaves K
pending. The next transfer's first ``wait(layer)`` returns
immediately reading one of those stale signals, even though
the FlexKV worker hasn't actually finished that layer's H2D
yet — and forward proceeds with wrong KV data.
"""
import os
for fd in self.load_event_fds:
# The fd is NONBLOCK: read until EAGAIN. Each read is 8 bytes.
while True:
try:
if not os.read(fd, 8):
break
except BlockingIOError:
break
except OSError:
break
self._finished = False
self.wait_remaining = [1] * self._num_layers
def wait(self, layer_index: int) -> None:
"""Block until the FlexKV worker signals layer ``layer_index``.
The fd was created with EFD_NONBLOCK so reset can drain it. We
re-introduce the blocking semantics with ``select.select`` on a
NONBLOCK fd: the read after select is guaranteed to consume one
signal.
"""
import os
import select
assert 0 <= layer_index < self._num_layers
fd = self.load_event_fds[layer_index]
while True:
select.select([fd], [], [])
try:
buf = os.read(fd, 8)
if buf:
break
except BlockingIOError:
# Spurious wakeup; loop and re-select.
continue
if layer_index == self._num_layers - 1:
self._finished = True
def close(self) -> None:
for fd in self.load_event_fds:
try:
os.close(fd)
except Exception:
pass
self.load_event_fds.clear()
def __del__(self) -> None:
try:
self.close()
except Exception:
pass
class FlexKVLayerDoneCounter:
"""Triple-buffered slot-based layerwise counter.
The KV pool calls ``wait_until(layer_id)`` once per layer during
forward. We track which producer slot the current task is using and
block on that slot's ``layer_id``-th eventfd. Producer rotation lets
the next prefetch start before the current one finishes consuming.
"""
def __init__(self, num_layers: int, num_counters: int = 3):
self.num_layers = num_layers
self.num_counters = num_counters
self.events: List[FlexKVLayerLoadingEvent] = [
FlexKVLayerLoadingEvent(num_layers) for _ in range(num_counters)
]
self.producer_index = -1
self.consumer_index = -1
self._task_to_producer: Dict[int, int] = {}
def register_task(self, task_id: int, producer_id: int) -> None:
self._task_to_producer[task_id] = producer_id
def register_task_with_explicit_counter_id(
self, task_id: int, counter_id: int
) -> None:
if not 0 <= counter_id < self.num_counters:
raise ValueError(
f"Invalid counter_id={counter_id}, must be in [0, {self.num_counters})"
)
self._task_to_producer[task_id] = counter_id
self.events[counter_id].reset_for_new_transfer()
def update_producer(self) -> int:
self.producer_index = (self.producer_index + 1) % self.num_counters
assert self.events[
self.producer_index
]._finished, "Producer event should be finished before reuse"
return self.producer_index
def set_consumer(self, task_id: int) -> None:
if task_id < 0:
self.consumer_index = -1
return
producer_id = self._task_to_producer.pop(task_id, None)
self.consumer_index = producer_id if producer_id is not None else -1
def wait_until(self, threshold: int) -> None:
if self.consumer_index < 0:
return
event = self.events[self.consumer_index]
if event.wait_remaining[threshold] <= 0:
return
event.wait_remaining[threshold] -= 1
event.wait(threshold)
def reset(self) -> None:
self.producer_index = -1
self.consumer_index = -1
self._task_to_producer.clear()
def __del__(self) -> None:
try:
for event in self.events:
event.close()
self.events.clear()
except Exception:
pass
@@ -0,0 +1,925 @@
"""Wrapper around FlexKV ``KVManager`` for sglang.
The public surface is small (see "Public API" below). The class owns:
* the FlexKV ``KVManager`` (server-client mode when ``dp_size > 1`` or
multi-instance; in-process otherwise — handled by FlexKV itself);
* the per-rank ``KVTPClient`` that registers this rank's GPU KV cache
with the FlexKV TransferManager;
* an optional ``FlexKVLayerDoneCounter`` plus the UDS-side handshake
that wires its eventfds into the FlexKV layerwise transfer worker.
Cross-rank sync uses :class:`FlexKVComm`. Only the **sync leader**
(rank 0 of every PP × CP × TP axis) talks to ``KVManager``; other
ranks block on broadcast / barrier.
Modes:
* **MP / synchronous** (default): ``retrieve_kv`` fires ``launch``
and blocks on ``wait`` so the device slots are ready by the time
sglang's prefill runs.
* **Layerwise** (``FLEXKV_ENABLE_LAYERWISE_TRANSFER=1``): ``launch``
is fired with ``layerwise_transfer=True`` and the per-layer hook
registered via ``register_layer_transfer_counter`` blocks each
forward layer on its own eventfd.
"""
from __future__ import annotations
import logging
import os
import socket
import struct
import time
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import torch
from sglang.srt.mem_cache.storage.flexkv.flexkv_comm import (
CMD_LAYERWISE,
CMD_PUT_META,
CMD_STORE_COMPLETE,
FlexKVComm,
FlexKVLayerDoneCounter,
send_fds,
)
try:
from flexkv.common.request import KVResponseStatus
from flexkv.common.storage import KVCacheLayout, KVCacheLayoutType
from flexkv.integration.config import FlexKVConfig
from flexkv.kvmanager import KVManager
from flexkv.server.client import KVTPClient
from flexkv.transfer.layerwise import build_layerwise_eventfd_socket_path
from flexkv.transfer_manager import TransferManagerOnRemote
except ImportError as exc: # pragma: no cover - runtime check
raise RuntimeError(
"FlexKV is not installed. Please install the FlexKV package to use "
"--enable-flexkv."
) from exc
logger = logging.getLogger(__name__)
class FlexKVConnector:
"""A FlexKV-side façade used by :class:`FlexKVRadixCache`.
This class manages connection lifecycle and provides a small,
sgl-friendly contract over FlexKV's task-based API:
* ``lookup_kv`` — page-aligned hit count + a held task id.
* ``retrieve_kv`` — synchronous load (launch + wait).
* ``start_load_kv_layerwise`` — layerwise async load.
* ``store_kv`` — page-aligned write back.
* ``check_completed_stores`` — drain async store completions.
* ``prefetch_async`` / ``check_prefetch_progress`` /
``cancel_prefetch`` — opportunistic CPU↔SSD/Remote staging.
* ``release_pending`` — cancel a held task whose load won't run.
* ``reset`` / ``shutdown``.
"""
def __init__(
self,
*,
sgl_model_config: Any,
server_args: Any,
page_size: int,
kvcache: Any,
tp_rank: int,
dp_rank: Optional[int],
pp_rank: int,
attn_cp_rank: int,
pp_group: Any = None,
attn_tp_group: Any = None,
attn_cp_group: Any = None,
) -> None:
self.page_size = int(page_size)
# 1. Resolve FlexKV config from env + sglang server args.
self.flexkv_config = FlexKVConfig.from_env()
self.rank_info = self.flexkv_config.post_init_from_sglang_config(
sglang_config=sgl_model_config,
server_args=server_args,
page_size=self.page_size,
tp_rank=tp_rank,
pp_rank=pp_rank,
dp_rank=dp_rank if dp_rank is not None else 0,
attn_cp_rank=attn_cp_rank,
)
self.model_config = self.flexkv_config.model_config
self.cache_config = self.flexkv_config.cache_config
self._label = f"[model_config={self.model_config}, rank_info={self.rank_info}]"
# 2. Cross-rank sync context.
world_rank = (
torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
)
self._sync_ctx = FlexKVComm(
rank_info=self.rank_info,
world_rank=world_rank,
pp_group=pp_group,
attn_tp_group=attn_tp_group,
attn_cp_group=attn_cp_group,
)
# 3. Align block counts across all ranks (MIN reduce) so each
# rank's KVManager registers compatible sizes.
for attr in ("num_cpu_blocks", "num_ssd_blocks", "num_remote_blocks"):
orig = getattr(self.cache_config, attr, None)
if orig is None or orig <= 0:
continue
aligned = self._sync_ctx.all_reduce_min(int(orig))
if aligned != orig:
logger.info(
"[FlexKV] Block count MIN alignment '%s': %d -> %d",
attr,
orig,
aligned,
)
setattr(self.cache_config, attr, aligned)
# 4. Extract MLA/MHA KV buffers + optional indexer buffers.
indexer_buffers = getattr(kvcache, "index_k_with_scale_buffer", None)
if hasattr(kvcache, "kv_buffer"):
# MLA: K and V share the same buffer (per-layer tensor).
kv_caches = list(kvcache.kv_buffer)
elif hasattr(kvcache, "k_buffer"):
# MHA: K buffers concatenated with V buffers, layer-first.
kv_caches = list(kvcache.k_buffer) + list(kvcache.v_buffer)
else:
raise AttributeError(
f"Unsupported KV cache type {type(kvcache).__name__}: "
f"expected kv_buffer (MLA/NSA) or k_buffer/v_buffer (MHA)."
)
self._kvcache = kvcache
# 5. On multi-node setups, every node beyond node 0 needs a
# TransferManagerOnRemote process (FlexKV side) before any rank
# on that node can register GPU buffers.
self._remote_process = None
if (
self.model_config.nnodes > 1
and self.rank_info.node_rank > 0
and self.rank_info.local_rank == 0
):
self._remote_process = TransferManagerOnRemote.create_process(
master_host=self.model_config.master_host,
master_ports=self.model_config.master_ports,
)
logger.info(
"[FlexKV] Launched TransferManagerOnRemote on node_rank=%d %s",
self.rank_info.node_rank,
self._label,
)
# 6. Bring up KVManager on the sync leader only.
self.kv_manager: Optional[KVManager] = None
if self._sync_ctx.is_sync_leader:
self.kv_manager = KVManager(
model_config=self.model_config,
cache_config=self.cache_config,
dp_client_id=self.rank_info.dp_client_id,
server_recv_port=self.flexkv_config.server_recv_port,
gpu_register_port=self.flexkv_config.gpu_register_port,
)
self.kv_manager.start()
# 7. Per-rank TP client registers this rank's GPU buffers.
self.tp_client = KVTPClient(
self.flexkv_config.gpu_register_port,
dp_client_id=self.rank_info.dp_client_id,
pp_rank=self.rank_info.pp_rank,
device_id=self.rank_info.local_rank,
)
self._register_with_retry(kv_caches, indexer_buffers)
# 8. Layerwise transfer plumbing.
self.enable_layerwise = bool(
int(os.environ.get("FLEXKV_ENABLE_LAYERWISE_TRANSFER", "0"))
)
self._layerwise_socket = build_layerwise_eventfd_socket_path(
dp_client_id=self.rank_info.dp_client_id,
pp_rank=self.rank_info.pp_rank,
model_config=self.model_config,
)
self._layerwise_eventfd_connect_max_retries = max(
360,
int(os.environ.get("FLEXKV_LAYERWISE_EVENTFD_CONNECT_MAX_RETRIES", "0")),
)
self.layer_done_counter: Optional[FlexKVLayerDoneCounter] = None
if self.enable_layerwise:
self.layer_done_counter = FlexKVLayerDoneCounter(
self.rank_info.num_layers_per_pp_stage
)
self._send_eventfds_to_worker()
# 9. Wait for the KVManager (and its remote subprocess) to be ready.
if self._sync_ctx.is_sync_leader:
self._wait_kv_manager_ready()
# 10. Per-rank in-flight tracking.
# Loads
self._pending_lookups: Dict[str, int] = {} # rid -> fkv_task_id
self._inflight_loads: Dict[int, int] = {} # producer_id -> rid hashlike
self._completed_layerwise: List[int] = []
self._launched_load_tids: List[int] = [] # leader-only, for periodic drain
# Stores
self._inflight_stores: Dict[str, int] = {} # rid -> fkv_task_id
# Prefetches
self._ongoing_prefetches: Dict[str, int] = {} # rid -> fkv_task_id
self._prefetch_enabled = bool(
self.cache_config.enable_ssd
or self.cache_config.enable_remote
or self.cache_config.enable_kv_sharing
)
logger.info(
"[FlexKV] Connector ready %s: layerwise=%s, prefetch=%s",
self._label,
self.enable_layerwise,
self._prefetch_enabled,
)
# ------------------------------------------------------------------
# Public API — lookup / load
# ------------------------------------------------------------------
def lookup_kv(
self,
token_ids: List[int],
token_mask: torch.Tensor,
rid: Optional[str] = None,
) -> Tuple[int, int]:
"""Page-aligned prefix lookup against FlexKV.
Args:
token_ids: full token id sequence we'd like to check.
token_mask: 1-D bool tensor or array, True for "this token is
*not* already on GPU and is a candidate for load-back".
rid: if set and hit > 0, the held FlexKV task id is stashed
under this key so a later ``retrieve_kv(rid, slots)`` call
can resolve it. If not set, the held task is cancelled when
hit > 0 and the caller didn't ask to track it.
Returns:
``(fkv_task_id, hit_count)``. ``hit_count`` is page-aligned
and may be smaller than the raw FlexKV match if the page
floor truncated it.
"""
fkv_task_id = -1
hit_length = 0
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
tids_np = np.asarray(token_ids, dtype=np.int64)
mask_np = self._as_numpy_mask(token_mask)
try:
res = self.kv_manager.get_match(token_ids=tids_np, token_mask=mask_np)
except Exception as exc: # noqa: BLE001
logger.warning("[FlexKV] get_match raised: %s", exc)
res = None
if res is None:
fkv_task_id = -1
hit_length = 0
else:
fkv_task_id, matched_mask = res
hit_length = int(matched_mask.sum()) if matched_mask is not None else 0
if self._sync_ctx.needs_sync:
payload = self._sync_ctx.scatter(
{"task_id": fkv_task_id, "hit": hit_length}
)
fkv_task_id = payload["task_id"]
hit_length = payload["hit"]
# Page-align: FlexKV transfers whole pages.
if hit_length > 0 and self.page_size > 1:
aligned = (hit_length // self.page_size) * self.page_size
if aligned < hit_length:
logger.debug(
"[FlexKV] lookup_kv: page-aligning hit %d -> %d (page=%d)",
hit_length,
aligned,
self.page_size,
)
hit_length = aligned
# Decide what to do with the held task. Three cases:
# 1. hit_length > 0 and rid given → stash for retrieve_kv later.
# 2. hit_length > 0 and rid is None → cancel; caller can't use it.
# 3. hit_length == 0 → no work to do; FlexKV already marked the
# empty graph COMPLETED inside get_match, cancel would warn.
if hit_length > 0 and rid is not None and fkv_task_id >= 0:
self._pending_lookups[rid] = fkv_task_id
elif hit_length > 0 and fkv_task_id >= 0 and self._sync_ctx.is_sync_leader:
assert self.kv_manager is not None
self.kv_manager.cancel([fkv_task_id])
return fkv_task_id, hit_length
def release_pending(self, rid: str) -> None:
"""Cancel the task held by an earlier ``lookup_kv(rid=...)`` that
won't be followed by a ``retrieve_kv`` (e.g. allocation failed)."""
fkv_task_id = self._pending_lookups.pop(rid, -1)
if fkv_task_id >= 0 and self._sync_ctx.is_sync_leader:
assert self.kv_manager is not None
self.kv_manager.cancel([fkv_task_id])
def retrieve_kv(
self,
rid: str,
slot_mapping: torch.Tensor,
) -> int:
"""Synchronous load: ``launch`` + ``wait``.
Returns the number of slots actually loaded. The caller is
responsible for having allocated ``slot_mapping`` of length
equal to ``hit_length`` from a prior ``lookup_kv``.
"""
fkv_task_id = self._pending_lookups.pop(rid, -1)
if fkv_task_id < 0:
return 0
slot_mapping_cpu = self._to_cpu_int64(slot_mapping)
# Cross-node PP receivers must send their slot mapping back to
# the TransferManagerOnRemote so the remote side knows where to
# land the H2D copies on its own GPUs.
if self._sync_ctx.should_send_slot_mapping_to_remote:
self._send_slot_mapping_to_remote(fkv_task_id, slot_mapping_cpu)
n = slot_mapping_cpu.numel()
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
self.kv_manager.launch(
task_ids=[fkv_task_id],
slot_mappings=[slot_mapping_cpu],
as_batch=True,
layerwise_transfer=False,
)
resp = self.kv_manager.wait([fkv_task_id], timeout=30.0)
if not (
fkv_task_id in resp
and resp[fkv_task_id].status == KVResponseStatus.SUCCESS
):
logger.warning(
"[FlexKV] retrieve_kv: task %d failed/timed out",
fkv_task_id,
)
n = 0
if self._sync_ctx.needs_sync:
self._sync_ctx.barrier()
return n
def start_load_kv_layerwise(
self,
rid: str,
slot_mapping: torch.Tensor,
) -> Tuple[int, int]:
"""Layerwise load. Fires ``launch(layerwise_transfer=True)`` and
returns ``(n_slots, producer_id)``. The caller registers
``producer_id`` with the layer hook so the KV pool blocks on
the right eventfds during forward."""
assert self.enable_layerwise and self.layer_done_counter is not None, (
"start_load_kv_layerwise called but layerwise transfer is "
"disabled. Set FLEXKV_ENABLE_LAYERWISE_TRANSFER=1."
)
fkv_task_id = self._pending_lookups.pop(rid, -1)
if fkv_task_id < 0:
return 0, -1
slot_mapping_cpu = self._to_cpu_int64(slot_mapping)
n = slot_mapping_cpu.numel()
if self._sync_ctx.should_send_slot_mapping_to_remote:
self._send_slot_mapping_to_remote(fkv_task_id, slot_mapping_cpu)
# Allocate / receive producer slot.
if self._sync_ctx.is_pp_receiver:
payload = self._sync_ctx.scatter_pp(None)
if payload.get("cmd") != CMD_LAYERWISE:
raise RuntimeError(
f"Tag mismatch: expected CMD_LAYERWISE, got "
f"{payload.get('cmd')}"
)
producer_id = int(payload["counter_id"])
self.layer_done_counter.register_task_with_explicit_counter_id(
fkv_task_id, producer_id
)
else:
producer_id = self.layer_done_counter.update_producer()
self.layer_done_counter.events[producer_id].reset_for_new_transfer()
self.layer_done_counter.register_task(fkv_task_id, producer_id)
if self._sync_ctx.is_pp_sender:
self._sync_ctx.scatter_pp(
{
"cmd": CMD_LAYERWISE,
"fkv_task_id": fkv_task_id,
"counter_id": producer_id,
}
)
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
self.kv_manager.launch(
task_ids=[fkv_task_id],
slot_mappings=[slot_mapping_cpu],
as_batch=True,
layerwise_transfer=True,
counter_id=producer_id,
)
self._launched_load_tids.append(fkv_task_id)
# Tell the layer hook which counter slot to wait on.
self.layer_done_counter.set_consumer(fkv_task_id)
return n, producer_id
def drain_launched_loads(self, threshold: int = 100) -> None:
"""Periodic non-blocking sweep on long-lived launched tasks so the
FlexKV pipe doesn't accumulate. No-op on non-leader ranks."""
if not self._sync_ctx.is_sync_leader or self.kv_manager is None:
return
if len(self._launched_load_tids) < threshold:
return
try:
self.kv_manager.try_wait(task_ids=list(self._launched_load_tids))
except Exception as exc: # noqa: BLE001
logger.debug("[FlexKV] drain_launched_loads try_wait: %s", exc)
self._launched_load_tids.clear()
# ------------------------------------------------------------------
# Public API — store
# ------------------------------------------------------------------
def store_kv(
self,
rid: str,
token_ids: List[int],
kv_indices: torch.Tensor,
) -> int:
"""Schedule a write back from GPU into FlexKV.
On the sync leader this runs ``put_match`` to discover which
tokens are NOT yet in FlexKV's CPU cache (= the "unmatched"
slice), then ``launch`` on those. On non-leaders the unmatched
mask is received over the PP fan-out so cross-node PP can
forward its slot mappings.
Returns the FlexKV task id of the in-flight store, or -1 if
nothing needed to be written.
"""
token_ids_np = np.asarray(token_ids, dtype=np.int64)
n = len(token_ids_np)
if n != len(kv_indices):
raise ValueError(
f"store_kv: token_ids has {n} entries but kv_indices "
f"has {len(kv_indices)} entries"
)
# Page-align inputs *before* put_match so the FlexKV allocator
# only reserves slots that line up with the slot_mapping we send.
if self.page_size > 1:
aligned_len = (n // self.page_size) * self.page_size
if aligned_len == 0:
self._send_pp_put_meta(-1, [])
return -1
if aligned_len < n:
token_ids_np = token_ids_np[:aligned_len]
kv_indices = kv_indices[:aligned_len]
fkv_task_id = -1
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
try:
res = self.kv_manager.put_match(token_ids=token_ids_np, token_mask=None)
except Exception as exc: # noqa: BLE001
logger.warning("[FlexKV] put_match raised: %s", exc)
res = None
if res is None:
self._send_pp_put_meta(-1, [])
return -1
fkv_task_id, unmatched_mask = res
self._send_pp_put_meta(fkv_task_id, unmatched_mask)
if int(unmatched_mask.sum()) > 0:
filtered = kv_indices[unmatched_mask]
slot_mapping_cpu = self._to_cpu_int64(filtered)
self.kv_manager.launch(
task_ids=[fkv_task_id],
slot_mappings=[slot_mapping_cpu],
as_batch=False,
layerwise_transfer=False,
)
self._inflight_stores[rid] = fkv_task_id
return fkv_task_id
return -1
# Non-leader path: receive the unmatched mask + maybe forward
# slot_mapping to the remote-side TransferManager.
if self._sync_ctx.is_pp_receiver:
payload = self._sync_ctx.scatter_pp(None)
if payload.get("cmd") != CMD_PUT_META:
raise RuntimeError(
f"Tag mismatch: expected CMD_PUT_META, got " f"{payload.get('cmd')}"
)
fkv_task_id = int(payload["fkv_task_id"])
mask_list = payload.get("unmatched_mask", [])
unmatched_mask = torch.tensor(mask_list, dtype=torch.bool)
if (
int(unmatched_mask.sum()) > 0
and fkv_task_id >= 0
and self._sync_ctx.should_send_slot_mapping_to_remote
):
filtered = kv_indices[unmatched_mask]
slot_mapping_cpu = self._to_cpu_int64(filtered)
self._send_slot_mapping_to_remote(fkv_task_id, slot_mapping_cpu)
self._inflight_stores[rid] = fkv_task_id
return fkv_task_id
def check_completed_stores(self) -> List[str]:
"""Return rids whose stores have completed since the last call."""
completed_rids: List[str] = []
completed_dict: Dict[int, Any] = {}
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
if self._inflight_stores:
fk_to_rid = {v: k for k, v in self._inflight_stores.items()}
try:
completed_dict = self.kv_manager.try_wait(
task_ids=list(fk_to_rid.keys())
)
except Exception as exc: # noqa: BLE001
logger.debug("[FlexKV] check_completed_stores: %s", exc)
completed_dict = {}
for fk_tid in completed_dict:
rid = fk_to_rid[fk_tid]
completed_rids.append(rid)
self._inflight_stores.pop(rid, None)
if self._sync_ctx.is_pp_sender:
self._sync_ctx.scatter_pp(
{
"cmd": CMD_STORE_COMPLETE,
"completed_fk_ids": list(completed_dict),
}
)
elif self._sync_ctx.is_pp_receiver:
payload = self._sync_ctx.scatter_pp(None)
if payload.get("cmd") != CMD_STORE_COMPLETE:
raise RuntimeError(
f"Tag mismatch: expected CMD_STORE_COMPLETE, got "
f"{payload.get('cmd')}"
)
fk_ids = payload.get("completed_fk_ids", [])
if fk_ids and self._inflight_stores:
fk_to_rid = {v: k for k, v in self._inflight_stores.items()}
for fk_tid in fk_ids:
if fk_tid in fk_to_rid:
rid = fk_to_rid[fk_tid]
completed_rids.append(rid)
self._inflight_stores.pop(rid, None)
if self._sync_ctx.needs_sync:
completed_rids = self._sync_ctx.scatter(completed_rids)
return completed_rids
def wait_store(self, rid: str, timeout: float = 30.0) -> bool:
"""Block until a single store task identified by ``rid`` finishes."""
fkv_task_id = self._inflight_stores.pop(rid, -1)
if fkv_task_id < 0:
return True
if not self._sync_ctx.is_sync_leader or self.kv_manager is None:
return True
try:
resp = self.kv_manager.wait([fkv_task_id], timeout=timeout)
except Exception as exc: # noqa: BLE001
logger.warning("[FlexKV] wait_store: %s", exc)
return False
return (
fkv_task_id in resp and resp[fkv_task_id].status == KVResponseStatus.SUCCESS
)
# ------------------------------------------------------------------
# Public API — prefetch
# ------------------------------------------------------------------
def prefetch_async(self, rid: str, token_ids: List[int]) -> int:
if not self._prefetch_enabled or not rid:
return -1
task_id = -1
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
try:
task_id = self.kv_manager.prefetch_async(
token_ids=np.asarray(token_ids, dtype=np.int64)
)
except Exception as exc: # noqa: BLE001
logger.debug("[FlexKV] prefetch_async: %s", exc)
task_id = -1
if self._sync_ctx.needs_sync:
payload = self._sync_ctx.scatter({"task_id": task_id})
task_id = payload["task_id"]
if task_id >= 0:
self._ongoing_prefetches[rid] = task_id
return task_id
def check_prefetch_progress(self, rid: str) -> bool:
if not self._prefetch_enabled:
return True
task_id = self._ongoing_prefetches.get(rid, -1)
if task_id < 0:
return True
done = False
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
try:
completed = self.kv_manager.try_wait(task_ids=[task_id])
except Exception: # noqa: BLE001
completed = {}
if task_id in completed:
done = True
if self._sync_ctx.needs_sync:
payload = self._sync_ctx.scatter({"done": done})
done = payload["done"]
if done:
self._ongoing_prefetches.pop(rid, None)
return done
def cancel_prefetch(self, rid: str) -> None:
self._pending_lookups.pop(rid, None)
# FlexKV doesn't currently support prefetch cancellation, but
# we still drop our tracking entry.
self._ongoing_prefetches.pop(rid, None)
# ------------------------------------------------------------------
# Layerwise transfer hooks
# ------------------------------------------------------------------
def register_layer_transfer_counter(self, kvcache: Any) -> None:
"""Register the FlexKVLayerDoneCounter onto sglang's KV pool so
each forward layer blocks on its eventfd. No-op when layerwise
is disabled."""
if (
self.layer_done_counter is None
or kvcache is None
or not hasattr(kvcache, "register_layer_transfer_counter")
):
return
kvcache.register_layer_transfer_counter(self.layer_done_counter)
# ------------------------------------------------------------------
# Lifecycle
# ------------------------------------------------------------------
def reset(self) -> None:
# Drop pending lookups (cancel their held tasks on the leader).
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
pending = [tid for tid in self._pending_lookups.values() if tid >= 0]
if pending:
try:
self.kv_manager.cancel(pending)
except Exception as exc: # noqa: BLE001
logger.debug("[FlexKV] reset cancel: %s", exc)
self._pending_lookups.clear()
self._ongoing_prefetches.clear()
self._inflight_loads.clear()
self._completed_layerwise.clear()
self._launched_load_tids.clear()
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
for fk_tid in list(self._inflight_stores.values()):
if fk_tid >= 0:
try:
self.kv_manager.wait([fk_tid], timeout=20.0)
except Exception: # noqa: BLE001
pass
self._inflight_stores.clear()
if self.layer_done_counter is not None:
self.layer_done_counter.reset()
def shutdown(self) -> None:
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
try:
self.kv_manager.shutdown()
except Exception as exc: # noqa: BLE001
logger.warning("[FlexKV] kv_manager.shutdown: %s", exc)
if self._remote_process is not None:
try:
self._remote_process.terminate()
self._remote_process.join(timeout=5.0)
if self._remote_process.is_alive():
self._remote_process.kill()
self._remote_process.join()
except Exception as exc: # noqa: BLE001
logger.warning("[FlexKV] remote process shutdown: %s", exc)
self._remote_process = None
# ------------------------------------------------------------------
# Private helpers
# ------------------------------------------------------------------
@staticmethod
def _as_numpy_mask(mask) -> np.ndarray:
if mask is None:
return None
if isinstance(mask, torch.Tensor):
return mask.detach().cpu().numpy()
return np.asarray(mask)
@staticmethod
def _to_cpu_int64(tensor: torch.Tensor) -> torch.Tensor:
if tensor.is_cuda:
tensor = tensor.cpu()
return tensor.to(torch.int64)
def _wait_kv_manager_ready(self, poll_interval: float = 10.0) -> None:
assert self.kv_manager is not None
wait_count = 0
while not self.kv_manager.is_ready():
time.sleep(poll_interval)
wait_count += 1
logger.info(
"[FlexKV] Waiting for FlexKV ready %s (waited %.0fs)",
self._label,
wait_count * poll_interval,
)
logger.info("[FlexKV] FlexKV is ready %s", self._label)
def _register_with_retry(
self,
kv_caches: List[torch.Tensor],
indexer_buffers: Optional[List[torch.Tensor]] = None,
max_retries: int = 360,
) -> None:
"""Retry GPU registration. On node_rank>0, the
TransferManagerOnRemote may not be ready immediately; retry up
to ~6 minutes."""
for attempt in range(max_retries):
try:
self._register_to_server(kv_caches, indexer_buffers)
return
except Exception as exc: # noqa: BLE001
if attempt == max_retries - 1:
raise
if attempt % 30 == 0:
logger.info(
"[FlexKV] GPU register retry %s attempt=%d/%d " "error=%s",
self._label,
attempt + 1,
max_retries,
exc,
)
time.sleep(1.0)
def _register_to_server(
self,
kv_caches: List[torch.Tensor],
indexer_buffers: Optional[List[torch.Tensor]] = None,
) -> None:
assert len(kv_caches) > 0
assert (
kv_caches[0].ndim == 3
), f"Expected 3D KV cache tensor, got shape={kv_caches[0].shape}"
is_mla = self.model_config.use_mla
num_blocks, num_kv_heads, head_size = kv_caches[0].shape
gpu_layout = KVCacheLayout(
type=KVCacheLayoutType.LAYERFIRST,
num_layer=self.rank_info.num_layers_per_pp_stage,
num_block=num_blocks // self.page_size,
tokens_per_block=self.page_size,
num_head=num_kv_heads,
head_size=head_size,
is_mla=is_mla,
)
indexer_layout = None
if indexer_buffers is not None and len(indexer_buffers) > 0:
indexer_tensor = indexer_buffers[0]
assert indexer_tensor.ndim == 2, (
f"Expected 2D indexer tensor (num_pages, page_stride_size), "
f"got shape={indexer_tensor.shape}"
)
indexer_layout = KVCacheLayout(
type=KVCacheLayoutType.LAYERFIRST,
num_layer=len(indexer_buffers),
num_block=indexer_tensor.shape[0],
tokens_per_block=1,
num_head=1,
head_size=indexer_tensor.shape[1],
is_mla=True,
)
self.tp_client.register_to_server(
kv_caches=kv_caches,
kv_layout=gpu_layout,
indexer_buffers=indexer_buffers,
indexer_layout=indexer_layout,
)
logger.info("[FlexKV] Registered KV caches to server %s", self._label)
def _send_pp_put_meta(self, fkv_task_id: int, unmatched_mask) -> None:
if not self._sync_ctx.is_pp_active:
return
if hasattr(unmatched_mask, "tolist"):
mask_list = unmatched_mask.tolist()
else:
mask_list = list(unmatched_mask)
self._sync_ctx.scatter_pp(
{
"cmd": CMD_PUT_META,
"fkv_task_id": fkv_task_id,
"unmatched_mask": mask_list,
}
)
def _send_slot_mapping_to_remote(
self, task_id: int, slot_mapping_cpu: torch.Tensor
) -> None:
np_arr = slot_mapping_cpu.numpy()
self.tp_client.set_slot_mapping(task_id, np_arr)
def _send_eventfds_to_worker(self, retry_interval: float = 1.0) -> None:
"""UDS handshake with the FlexKV layerwise transfer worker.
Sends per-counter eventfd FDs over a unix domain socket using
``SCM_RIGHTS``. Retries connect (worker may not yet be up) and
retries the whole connect+send sequence on send error.
"""
max_retries = self._layerwise_eventfd_connect_max_retries
max_send_retries = 3
last_error: Optional[BaseException] = None
assert self.layer_done_counter is not None
for send_attempt in range(max_send_retries):
sock: Optional[socket.socket] = None
try:
# Phase 1: connect (worker may not yet be up).
for attempt in range(max_retries):
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
try:
sock.connect(self._layerwise_socket)
logger.info(
"[FlexKV] Eventfd connected %s socket=%s attempts=%d",
self._label,
self._layerwise_socket,
attempt + 1,
)
break
except (FileNotFoundError, ConnectionRefusedError) as exc:
sock.close()
sock = None
if attempt == max_retries - 1:
raise RuntimeError(
f"[FlexKV] Failed to connect to eventfd socket "
f"{self._layerwise_socket} after {max_retries} attempts"
) from exc
time.sleep(retry_interval)
assert sock is not None
# Phase 2: send 16-byte metadata + per-counter FDs + read ACK.
num_counters = self.layer_done_counter.num_counters
metadata = struct.pack(
"iiii",
self.rank_info.tp_rank_per_node,
self.model_config.tp_size_per_node,
self.rank_info.num_layers_per_pp_stage,
num_counters,
)
sock.sendall(metadata)
for counter_id in range(num_counters):
fds = self.layer_done_counter.events[counter_id].load_event_fds
send_fds(sock, fds, struct.pack("i", counter_id))
sock.settimeout(30.0)
try:
ack = sock.recv(1)
except socket.timeout as exc:
raise RuntimeError(
"Timed out waiting for ACK from FlexKV layerwise worker"
) from exc
if not ack or ack[0] != 1:
raise RuntimeError(
f"FlexKV layerwise worker NACK'd eventfd transfer "
f"(ack={ack!r})"
)
logger.info(
"[FlexKV] Eventfd handshake complete %s counters=%d layers=%d",
self._label,
num_counters,
self.rank_info.num_layers_per_pp_stage,
)
return
except Exception as exc: # noqa: BLE001
last_error = exc
logger.warning(
"[FlexKV] Eventfd handshake send_attempt=%d/%d failed: %s",
send_attempt + 1,
max_send_retries,
exc,
)
finally:
if sock is not None:
sock.close()
time.sleep(retry_interval)
raise RuntimeError(
f"[FlexKV] Failed to send eventfds to {self._layerwise_socket} "
f"after {max_send_retries} attempts: {last_error}"
)
@@ -0,0 +1,511 @@
"""FlexKV-backed RadixCache for sglang.
This module exposes :class:`FlexKVRadixCache`, a subclass of
:class:`sglang.srt.mem_cache.radix_cache.RadixCache` that delegates
host-side prefix storage to a FlexKV ``KVManager``. The design mirrors
``LMCRadixCache`` (the LMCache integration) so the scheduler-side
contract is identical:
* MP (synchronous) mode — the default.
``match_prefix`` fires only a FlexKV LOOKUP and returns ``host_hit_length``;
the scheduler then calls :meth:`init_load_back` at dispatch time which
allocates slots and fires the FlexKV RETRIEVE.
* IP (layerwise) mode — enabled with ``FLEXKV_ENABLE_LAYERWISE_TRANSFER=1``.
``match_prefix`` allocates uncached slots and kicks off a layerwise
load; the per-layer hook registered via
``register_layer_transfer_counter`` then waits on each layer's
eventfd inside the model's forward pass.
Selection: ``--enable-flexkv`` on the sglang CLI routes the default
RadixCache factory here. See ``__init__.py`` in this package for the
``register_radix_cache_backend("flexkv", ...)`` entry-point that backs
the explicit ``--radix-cache-backend=flexkv`` form.
"""
from __future__ import annotations
import enum
import logging
import threading
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional, Tuple
import torch
from sglang.srt.mem_cache.base_prefix_cache import (
EvictParams,
EvictResult,
InitLoadBackParams,
MatchPrefixParams,
MatchResult,
)
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
from sglang.srt.mem_cache.storage.flexkv.flexkv_connector import FlexKVConnector
if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__)
class FlexKVMode(enum.Enum):
MP = enum.auto() # synchronous lookup → retrieve in two phases
IP = enum.auto() # in-process layerwise transfer
@dataclass
class _LoadBackMarker:
"""State carried from a hit-producing ``match_prefix`` to its
matching ``init_load_back``. The detached ``RadixKey`` is a snapshot
of the matched key at lookup time (the live request key aliases
``req.fill_ids`` which keeps growing)."""
key: RadixKey
value_numel: int # device tokens already present at lookup time
class FlexKVRadixCache(RadixCache):
"""RadixCache extended with FlexKV host-tier IO."""
def __init__(
self,
params: CacheInitParams,
model_config: Optional[ModelConfig],
server_args: ServerArgs,
tp_rank: int,
tp_size: int,
dp_rank: Optional[int],
pp_rank: int,
attn_cp_rank: int,
tp_group=None,
pp_group=None,
attn_tp_group=None,
attn_cp_group=None,
) -> None:
super().__init__(params)
kvcache = self.token_to_kv_pool_allocator.get_kvcache()
# ``tp_group`` and ``attn_tp_group`` are sometimes passed
# interchangeably by sglang's factory; prefer the explicit
# ``attn_tp_group`` when given.
attn_tp_group_eff = attn_tp_group if attn_tp_group is not None else tp_group
self.flexkv_connector = FlexKVConnector(
sgl_model_config=model_config,
server_args=server_args,
page_size=params.page_size,
kvcache=kvcache,
tp_rank=tp_rank,
dp_rank=dp_rank,
pp_rank=pp_rank,
attn_cp_rank=attn_cp_rank,
pp_group=pp_group,
attn_tp_group=attn_tp_group_eff,
attn_cp_group=attn_cp_group,
)
self._mode = (
FlexKVMode.IP if self.flexkv_connector.enable_layerwise else FlexKVMode.MP
)
if self._mode is FlexKVMode.IP:
# Register the eventfd counter onto sglang's KV pool so each
# forward layer blocks on its own eventfd.
self.flexkv_connector.register_layer_transfer_counter(kvcache)
# CUDA streams (mirroring LMCRadixCache).
self.load_stream = torch.cuda.Stream()
self.store_stream = torch.cuda.Stream()
# Two-phase MP load: stash marker between ``match_prefix`` and
# ``init_load_back``.
self._load_markers: dict[str, _LoadBackMarker] = {}
# ``store_kv`` is async — we keep a lock on the source node
# until FlexKV signals completion, draining in ``evict`` /
# ``check_hicache_events``.
self._inflight_store_nodes: dict[str, TreeNode] = {}
self._node_lock = threading.Lock()
# ------------------------------------------------------------------
# Lifecycle
# ------------------------------------------------------------------
def reset(self) -> None: # type: ignore[override]
super().reset()
if hasattr(self, "_load_markers"):
self._load_markers.clear()
if hasattr(self, "_inflight_store_nodes"):
with self._node_lock:
self._inflight_store_nodes.clear()
if hasattr(self, "flexkv_connector"):
self.flexkv_connector.reset()
def shutdown(self) -> None:
if hasattr(self, "flexkv_connector"):
self.flexkv_connector.shutdown()
# ------------------------------------------------------------------
# match_prefix
# ------------------------------------------------------------------
def match_prefix(self, params: MatchPrefixParams) -> MatchResult: # type: ignore[override]
"""Look up the longest cached prefix on host KV (FlexKV).
Dispatches to :meth:`_mp_match_prefix` or :meth:`_ip_match_prefix`
depending on whether layerwise transfer is enabled.
"""
key = params.key
if self.disable or not key:
return super().match_prefix(params)
# FlexKV operates at page granularity — round the lookup query
# down to a multiple of ``page_size`` so the hit count we report
# back to sglang matches what FlexKV can actually serve.
if self.page_size != 1:
aligned_len = (len(key) // self.page_size) * self.page_size
key = key[:aligned_len]
base_res = super().match_prefix(params)
if len(key) == 0:
return base_res
device_value: torch.Tensor = base_res.device_indices
last_node: TreeNode = base_res.last_device_node
if self._mode is FlexKVMode.MP:
if params.req is None:
return base_res
return self._mp_match_prefix(
key, base_res, device_value, last_node, params.req
)
return self._ip_match_prefix(key, base_res, device_value, last_node)
def _mp_match_prefix(
self,
key: RadixKey,
base_res: MatchResult,
device_value: torch.Tensor,
last_node: TreeNode,
req: Req,
) -> MatchResult:
"""LOOKUP-only path. Sets ``host_hit_length`` on the result so
the scheduler later invokes :meth:`init_load_back`."""
token_ids = key.raw_token_ids()
device_len = int(device_value.numel())
if device_len >= len(token_ids):
return base_res
# token_mask=True for tokens NOT on device — FlexKV decides
# which of those it can serve.
token_mask = torch.zeros(len(token_ids), dtype=torch.bool)
token_mask[device_len:] = True
fkv_task_id, hit = self.flexkv_connector.lookup_kv(
token_ids=token_ids, token_mask=token_mask, rid=req.rid
)
if hit <= 0:
return base_res
# Snapshot the matched key (the live key aliases ``req.fill_ids``).
if token_ids is key.token_ids:
token_ids_snap = token_ids[:]
else:
token_ids_snap = token_ids
self._load_markers[req.rid] = _LoadBackMarker(
key=RadixKey(token_ids_snap, key.extra_key, key.is_bigram),
value_numel=device_len,
)
return MatchResult(
device_indices=device_value,
last_device_node=last_node,
last_host_node=last_node,
best_match_node=last_node,
host_hit_length=hit,
)
def _ip_match_prefix(
self,
key: RadixKey,
base_res: MatchResult,
device_value: torch.Tensor,
last_node: TreeNode,
) -> MatchResult:
"""Layerwise path: allocate slots and fire ``start_load_kv_layerwise``
immediately. Per-layer hook waits during forward."""
token_ids = key.raw_token_ids()
device_len = int(device_value.numel())
if device_len >= len(token_ids):
return base_res
# Quick LOOKUP first to discover how many slots we'd need.
token_mask = torch.zeros(len(token_ids), dtype=torch.bool)
token_mask[device_len:] = True
# No rid here — IP mode self-pops; pass a synthetic stable key.
synthetic_rid = f"_ip_{id(key)}"
_, hit = self.flexkv_connector.lookup_kv(
token_ids=token_ids, token_mask=token_mask, rid=synthetic_rid
)
if hit <= 0:
return base_res
result = self._allocate_and_load(
key=key,
value_numel=device_len,
uncached_len=hit,
last_node=last_node,
load_fn=lambda slot_mapping: self.flexkv_connector.start_load_kv_layerwise(
synthetic_rid, slot_mapping
)[0],
)
if result is None:
return base_res
new_slots, new_node = result
return MatchResult(
device_indices=torch.cat([device_value, new_slots]),
last_device_node=new_node,
last_host_node=new_node,
best_match_node=new_node,
)
# ------------------------------------------------------------------
# init_load_back (MP RETRIEVE)
# ------------------------------------------------------------------
def init_load_back( # type: ignore[override]
self,
params: InitLoadBackParams,
) -> Tuple[torch.Tensor, Optional[TreeNode]]:
"""MP RETRIEVE. Allocates uncached slots and fires the FlexKV
load; inserts the resulting TreeNode."""
req = params.req
last_node: TreeNode = params.best_match_node
marker = self._load_markers.pop(req.rid, None)
if marker is None:
# ``match_prefix`` decided there was no work to do, but the
# scheduler still called us. Release any held task and
# return an empty load.
self.flexkv_connector.release_pending(req.rid)
return (
torch.empty((0,), dtype=torch.int64, device=self.device),
last_node,
)
result = self._allocate_and_load(
key=marker.key,
value_numel=marker.value_numel,
uncached_len=params.host_hit_length,
last_node=last_node,
load_fn=lambda slot_mapping: self.flexkv_connector.retrieve_kv(
req.rid, slot_mapping
),
)
if result is None:
# Allocation failed or load returned zero. ``retrieve_kv``
# already cancels/cleans up on failure paths; release_pending
# is idempotent for the case where allocation failed before
# we even popped the held task.
self.flexkv_connector.release_pending(req.rid)
return (
torch.empty((0,), dtype=torch.int64, device=self.device),
last_node,
)
return result
def _allocate_and_load(
self,
*,
key: RadixKey,
value_numel: int,
uncached_len: int,
last_node: TreeNode,
load_fn,
) -> Optional[Tuple[torch.Tensor, TreeNode]]:
"""Shared allocator + post-load bookkeeping for MP/IP.
Returns ``(token_slots[:fetched], new_node)`` on success.
``None`` on either allocation failure or zero retrieved (in
which case all slots are freed).
"""
if uncached_len <= 0:
return None
# Evict to make room when needed.
if self.token_to_kv_pool_allocator.available_size() < uncached_len:
self.evict(EvictParams(num_tokens=uncached_len))
token_slots = self.token_to_kv_pool_allocator.alloc(uncached_len)
if token_slots is None:
return None
# The FlexKV ``launch`` interface takes the slot indices for the
# tokens it will write — no leading ``-1`` padding (FlexKV has
# no concept of "skip these device slots, they're already
# cached"; we pass it exactly the destinations for the
# uncached tail).
num_retrieved = load_fn(token_slots.to(torch.int64))
if num_retrieved <= 0:
self.token_to_kv_pool_allocator.free(token_slots)
return None
# Free the tail of the over-allocation when FlexKV returned
# fewer than expected.
if num_retrieved < uncached_len:
self.token_to_kv_pool_allocator.free(token_slots[num_retrieved:])
fetched_slots = token_slots[:num_retrieved]
else:
fetched_slots = token_slots
new_node = TreeNode(priority=last_node.priority)
start = value_numel
end = start + num_retrieved
new_node.key = key[start:end]
new_node.value = fetched_slots
new_node.parent = last_node
last_node.children[new_node.key.child_key(self.page_size)] = new_node
self.evictable_size_ += num_retrieved
self._update_leaf_status(last_node)
self._update_leaf_status(new_node)
self._record_store_event(new_node.parent)
self._record_store_event(new_node)
return fetched_slots, new_node
# ------------------------------------------------------------------
# cache_finished_req (STORE)
# ------------------------------------------------------------------
def cache_finished_req( # type: ignore[override]
self, req: Req, is_insert: bool = True
) -> None:
"""Base cache_finished_req then fire an async FlexKV store."""
super().cache_finished_req(req, is_insert=is_insert)
if not is_insert:
self._load_markers.pop(req.rid, None)
return
# Compute the committed prefix mirroring LMCRadixCache's logic.
from sglang.srt.runtime_context import get_server_args
global_server_args = get_server_args()
topk = global_server_args.speculative_eagle_topk
enable_kv_committed_len = topk is None or topk == 1
if enable_kv_committed_len:
kv_committed_len = req.kv_committed_len
else:
kv_committed_len = len(req.origin_input_ids) + max(
len(req.output_ids) - 1, 0
)
token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
if not token_ids:
return
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len
]
# Anchor on the new last_device_node so FlexKV's lock matches
# the node we'll later unlock when the store completes.
match_result = super().match_prefix(
MatchPrefixParams(key=RadixKey(token_ids, req.extra_key))
)
new_last_node = match_result.last_device_node
if new_last_node is None:
return
self.inc_lock_ref(new_last_node)
try:
with torch.cuda.stream(self.store_stream):
fkv_task_id = self.flexkv_connector.store_kv(
rid=req.rid,
token_ids=list(token_ids),
kv_indices=kv_indices,
)
except Exception: # noqa: BLE001
self.dec_lock_ref(new_last_node)
raise
if fkv_task_id < 0:
# Nothing to write back (either everything already in
# FlexKV, or put_match failed / returned None).
self.dec_lock_ref(new_last_node)
return
with self._node_lock:
self._inflight_store_nodes[req.rid] = new_last_node
# ------------------------------------------------------------------
# evict + completion draining
# ------------------------------------------------------------------
def evict(self, params: EvictParams) -> EvictResult: # type: ignore[override]
"""Drain completed stores before letting the base evict touch
the source nodes."""
if self.disable:
return EvictResult()
self._drain_completed_stores()
# Make sure the store stream's GPU work is observed before any
# eviction frees the source slots.
self.store_stream.synchronize()
return super().evict(params)
def check_hicache_events(self) -> None: # type: ignore[override]
"""Periodic non-blocking sweep called by the scheduler tick.
Drains both store completions (so source nodes get unlocked
quickly) and the launched-load tail (so the FlexKV pipe
doesn't accumulate)."""
self._drain_completed_stores()
self.flexkv_connector.drain_launched_loads()
def _drain_completed_stores(self) -> None:
completed_rids = self.flexkv_connector.check_completed_stores()
if not completed_rids:
return
with self._node_lock:
for rid in completed_rids:
node = self._inflight_store_nodes.pop(rid, None)
if node is not None:
self.dec_lock_ref(node)
# ------------------------------------------------------------------
# Optional pass-throughs used by the scheduler
# ------------------------------------------------------------------
def release_aborted_request(self, rid: str) -> None:
"""Clean up tracking for an aborted request without invoking FlexKV."""
self._load_markers.pop(rid, None)
with self._node_lock:
node = self._inflight_store_nodes.pop(rid, None)
if node is not None:
self.dec_lock_ref(node)
self.flexkv_connector.release_pending(rid)
self.flexkv_connector.cancel_prefetch(rid)
def prefetch_from_storage(
self, rid: str, last_host_node: TreeNode, token_ids
) -> None:
"""Kick off an opportunistic prefetch (SSD/Remote → CPU)."""
try:
self.flexkv_connector.prefetch_async(rid, list(token_ids))
except Exception as exc: # noqa: BLE001
logger.debug("[FlexKV] prefetch_from_storage: %s", exc)
def check_prefetch_progress(self, rid: str) -> bool:
return self.flexkv_connector.check_prefetch_progress(rid)
def terminate_prefetch(self, rid: str) -> None:
self.flexkv_connector.cancel_prefetch(rid)
def pop_prefetch_loaded_tokens(self, rid: str) -> int:
# FlexKV doesn't expose per-rid prefetched token counts yet.
return 0
@property
def hicache_storage_pass_prefix_keys(self) -> bool:
# We pass token ids, not opaque key strings, so no prefix-key
# accounting in the scheduler.
return False
@@ -0,0 +1,180 @@
"""End-to-end correctness check for the FlexKV sglang connector.
Run twice with different server configurations:
# 1. Baseline: launch sglang WITHOUT --enable-flexkv first, then:
python verify_outputs.py --phase baseline
# 2. Restart sglang WITH --enable-flexkv, then:
python verify_outputs.py --phase test
Each prompt is requested twice in the test phase:
* R1 (fresh) — first call after server start; FlexKV may still have
state from a previous test run, but match must equal baseline.
* R2 (cached) — after /flush_cache; the GPU radix is empty but
FlexKV's CPU pool keeps the data, so R2 should be a host hit.
Both R1 and R2 output_ids must byte-equal the baseline. Any mismatch
is reported and exit code is non-zero. Run again with
``FLEXKV_ENABLE_LAYERWISE_TRANSFER=1`` set on the server to exercise
the layerwise path.
"""
from __future__ import annotations
import argparse
import json
import sys
import time
import urllib.request
PROMPTS = [
(
"PROMPT_SHORT",
"The capital of France is",
12,
),
(
"PROMPT_MEDIUM",
"List the first ten prime numbers in order: 2, 3, 5, ",
24,
),
(
"PROMPT_LONG",
# Long enough to span many KV pages.
(
"In the year 2025, a research team at a major AI lab released a "
"report describing the architecture of a new large language "
"model. The report had several sections. Section one introduced "
"the model and its training data. Section two covered the "
"attention mechanism in detail, including how the keys and "
"values were managed. Section three discussed deployment, "
"including KV cache offloading to CPU memory and to disk. "
"Section four reported evaluation results on standard "
"benchmarks. Section five concluded with a discussion of "
"future work, including improvements to the offloading layer "
"and to the radix tree used to index cached prefixes. "
"Now, summarize the report in one sentence: "
),
60,
),
]
def _post(host: str, path: str, body=None, timeout=120) -> str:
if body is None:
req = urllib.request.Request(f"http://{host}{path}", method="POST")
else:
req = urllib.request.Request(
f"http://{host}{path}",
data=json.dumps(body).encode(),
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(req, timeout=timeout) as resp:
return resp.read().decode()
def gen(host: str, text: str, max_new: int) -> dict:
raw = _post(
host,
"/generate",
{
"text": text,
"sampling_params": {
"max_new_tokens": max_new,
"temperature": 0.0,
},
},
)
return json.loads(raw)
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument(
"--host",
default="127.0.0.1:30000",
help="sglang server host:port (default 127.0.0.1:30000)",
)
ap.add_argument(
"--phase",
choices=["baseline", "test"],
required=True,
help="baseline: record golden outputs; test: compare against them",
)
ap.add_argument(
"--baseline-file",
default="/tmp/flexkv_baseline.json",
help="where to write/read the baseline outputs",
)
args = ap.parse_args()
if args.phase == "baseline":
result = {}
for name, text, max_new in PROMPTS:
r = gen(args.host, text, max_new)
meta = r["meta_info"]
print(
f"[baseline] {name}: completion={meta['completion_tokens']}, "
f"cached={meta['cached_tokens']}, text={r['text']!r}"
)
result[name] = {
"text": r["text"],
"output_ids": r["output_ids"],
"completion_tokens": meta["completion_tokens"],
}
with open(args.baseline_file, "w") as f:
json.dump(result, f, indent=2)
print(f"\nWrote baseline to {args.baseline_file}")
return 0
with open(args.baseline_file) as f:
baseline = json.load(f)
errors = 0
for name, text, max_new in PROMPTS:
b = baseline[name]
# R1 (fresh): may or may not hit FlexKV depending on prior state.
r1 = gen(args.host, text, max_new)
m1 = r1["meta_info"]
ok1 = r1["output_ids"] == b["output_ids"]
print(
f"[test/{name}] R1 fresh: cached={m1['cached_tokens']}/"
f"{m1['prompt_tokens']}, details={m1.get('cached_tokens_details')}, "
f"output_match={'OK' if ok1 else 'MISMATCH'}"
)
if not ok1:
print(f" baseline: {b['text']!r}")
print(f" r1 : {r1['text']!r}")
errors += 1
# Give the async D2H store a beat to complete before we flush.
time.sleep(2)
_post(args.host, "/flush_cache")
time.sleep(1)
# R2 (cached): GPU radix is empty; FlexKV must serve the prefix.
r2 = gen(args.host, text, max_new)
m2 = r2["meta_info"]
ok2 = r2["output_ids"] == b["output_ids"]
ratio = m2["cached_tokens"] / max(1, m2["prompt_tokens"])
print(
f"[test/{name}] R2 cached: cached={m2['cached_tokens']}/"
f"{m2['prompt_tokens']} ({ratio:.1%}), "
f"details={m2.get('cached_tokens_details')}, "
f"output_match={'OK' if ok2 else 'MISMATCH'}"
)
if not ok2:
print(f" baseline: {b['text']!r}")
print(f" r2 : {r2['text']!r}")
errors += 1
print(f"\nTotal mismatches: {errors}")
return 1 if errors else 0
if __name__ == "__main__":
sys.exit(main())
@@ -0,0 +1,71 @@
# Using HF3FS as L3 Global KV Cache
This document provides step-by-step instructions for setting up a k8s + 3FS + SGLang runtime environment from scratch, describing how to utilize deepseek-hf3fs as the L3 KV cache for SGLang.
The process consists of five main steps:
## Step 1: Install deepseek-3fs via 3fs-Operator
Refer to the [3fs-operator documentation](https://github.com/aliyun/kvc-3fs-operator/blob/main/README_en.md) to deploy 3FS components in your Kubernetes environment using the Operator with one-click deployment.
## Step 2: Launch SGLang Pod
Start your SGLang Pod while specifying 3FS-related labels in the YAML configuration. Follow the [fuse-client-creation guide](https://github.com/aliyun/kvc-3fs-operator/blob/main/README_en.md#fuse-client-creation).
## Step 3: Configure Usrbio Client in SGLang Pod
The Usrbio client is required for accessing 3FS. Install it in your SGLang Pod using either method below:
**Alternative 1 (Recommend):** Built from the source code, the following provides quick installation commands (refer to [setup_usrbio_client.md](setup_usrbio_client.md))
```
set -e; \
. /etc/os-release; \
case "$VERSION_ID" in \
"22.04") \
CLANG_VERSION="14"; \
GIT_BRANCH=main; \
GIT_COMMIT_ID=6f029c439d0d22995900ca357d51b37975c6ffb5; \
;; \
"24.04") \
CLANG_VERSION="18"; \
GIT_BRANCH="ubuntu24.04"; \
GIT_COMMIT_ID=d0cf83a42395cdb2a66d3ce83cb0a11a46bee9f3; \
;; \
*) \
echo "Unsupported Ubuntu version: $VERSION_ID"; \
exit 1; \
;; \
esac; \
apt-get update && apt-get install -y --no-install-recommends \
clang-format-$CLANG_VERSION clang-$CLANG_VERSION clang-tidy-$CLANG_VERSION lld-$CLANG_VERSION meson google-perftools \
libaio-dev libdouble-conversion-dev libdwarf-dev libgflags-dev libgmock-dev libgoogle-perftools-dev liblz4-dev liblzma-dev libuv1-dev \
&& rm -rf /var/lib/apt/lists/* \
&& apt-get clean \
&& git clone https://github.com/novitalabs/3FS.git -b $GIT_BRANCH 3fs \
&& cd 3fs \
&& git checkout $GIT_COMMIT_ID \
&& git submodule update --init --recursive \
&& ./patches/apply.sh \
&& CMAKE_BUILD_PARALLEL_LEVEL=32 python3 setup.py bdist_wheel -d dist \
&& pip install dist/*.whl \
&& cd .. \
&& rm -rf 3fs
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages
```
**Alternative 2:** Run `pip3 install hf3fs-py-usrbio` (Follow https://pypi.org/project/hf3fs-py-usrbio/#files)
## Step 4: Deploy Model Serving
### Single Node Deployment
```bash
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages
python3 -m sglang.launch_server \
--model-path /code/models/Qwen3-32B/ \
--host 0.0.0.0 --port 10000 \
--page-size 64 \
--enable-hierarchical-cache \
--hicache-ratio 2 --hicache-size 0 \
--hicache-write-policy write_through \
--hicache-storage-backend hf3fs
```
### Multi-Node Deployment (Shared KV Cache)
Follow the [deploy_sglang_3fs_multinode.md](deploy_sglang_3fs_multinode.md) guide to deploy SGLang with 3FS across multiple nodes for shared KV caching.
@@ -0,0 +1,65 @@
# 1. Startup 3fs metadata service
```bash
nohup python3 -m sglang.srt.mem_cache.storage.hf3fs.mini_3fs_metadata_server > meta.out &
```
# 2. Startup sglang engine
## HF3fs configures
```bash
vim /sgl-workspace/sglang/benchmark/hf3fs/hf3fs_config.json
{
"file_path_prefix": "/data/hicache",
"file_size": 1099511627776,
"numjobs": 16,
"entries": 8,
"metadata_server_url": "http://metaServerIp:18000"
}
```
## node1
```bash
export SGLANG_HICACHE_HF3FS_CONFIG_PATH=/sgl-workspace/sglang/benchmark/hf3fs/hf3fs_config.json
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages
rm -rf instance1.out && \
nohup python3 -m sglang.launch_server \
--model-path /code/models/Qwen3-32B/ \
--host 0.0.0.0 --port 10000 \
--page-size 64 \
--enable-hierarchical-cache \
--hicache-ratio 2 --hicache-size 0 \
--hicache-write-policy write_through \
--hicache-storage-backend hf3fs --tp 2 > instance1.out &
```
## node2
```bash
export SGLANG_HICACHE_HF3FS_CONFIG_PATH=/sgl-workspace/sglang/benchmark/hf3fs/hf3fs_config.json
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages
rm -rf instance2.out && \
nohup python3 -m sglang.launch_server \
--model-path /code/models/Qwen3-32B/ \
--host 0.0.0.0 --port 10000 \
--page-size 64 \
--enable-hierarchical-cache \
--hicache-ratio 2 --hicache-size 0 \
--hicache-write-policy write_through \
--hicache-storage-backend hf3fs --tp 2 > instance2.out &
```
# 3. Startup router
```bash
rm -rf router.out && \
nohup python -m sglang_router.launch_router --worker-urls http://node1:10000 http://node2:10000 > router.out &
```
# 4. Startup multiturn benchmark
```bash
rm -rf bench_multiturn.out && \
nohup python3 benchmark/hicache/bench_multiturn.py \
--model-path /code/models/Qwen3-32B \
--dataset-path /code/models/ShareGPT_V3_unfiltered_cleaned_split.json \
--port 30000 \
--request-length 2048 --num-clients 512 --num-rounds 5 --max-parallel 8 \
> bench_multiturn.out &
```
@@ -0,0 +1,68 @@
# HiCacheHF3FS Setup
## Build & Package
### Source Code
https://github.com/deepseek-ai/3FS/blob/main/README.md#check-out-source-code
```sh
git clone https://github.com/deepseek-ai/3fs
cd 3fs
git submodule update --init --recursive
./patches/apply.sh
```
### Build Dev Container
https://github.com/deepseek-ai/3FS/blob/main/dockerfile/dev.dockerfile
```sh
cd 3fs/dockerfile
docker build -t hf3fs:dev -f dev.dockerfile .
```
### Generate Python Wheel
```sh
docker run -it hf3fs:dev bash
# Inside the development container
git clone https://github.com/deepseek-ai/3fs
cd 3fs
git submodule update --init --recursive
./patches/apply.sh
apt-get update \
&& apt-get install -y --no-install-recommends \
python3 python3-pip \
&& apt-get clean \
&& rm -rf /var/lib/apt/lists/*
# apt install python3.12 python3.12-venv python3.12-dev
# curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py
# python3.12 get-pip.py
# Generated wheel location: dist/hf3fs_py_usrbio-1.2.9+2db69ce-cp310-cp310-linux_x86_64.whl
python3 setup.py bdist_wheel
```
## Installation
```sh
# Install Dependencies
# https://github.com/deepseek-ai/3FS/blob/main/dockerfile/dev.dockerfile
apt update && apt install -y \
libaio-dev \
libboost-all-dev \
libdouble-conversion-dev \
libdwarf-dev \
libgflags-dev \
libgmock-dev \
libgoogle-glog-dev \
libgoogle-perftools-dev \
libgtest-dev \
liblz4-dev \
liblzma-dev \
libssl-dev \
libunwind-dev \
libuv1-dev
# Install Python Package
pip install hf3fs_py_usrbio-1.2.9+394583d-cp312-cp312-linux_x86_64.whl
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages
```
@@ -0,0 +1,163 @@
import logging
import os
from abc import ABC, abstractmethod
from typing import List
import torch
class Hf3fsClient(ABC):
"""Abstract interface for HF3FS clients."""
@abstractmethod
def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
"""Initialize the HF3FS client.
Args:
path: File path for storage
size: Total size of storage file
bytes_per_page: Bytes per page
entries: Number of entries for batch operations
"""
pass
@abstractmethod
def batch_read(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
"""Batch read from storage."""
pass
@abstractmethod
def batch_write(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
"""Batch write to storage."""
pass
@abstractmethod
def check(self, offsets: List[int], tensors: List[torch.Tensor]) -> None:
"""Validate batch operation parameters."""
pass
@abstractmethod
def get_size(self) -> int:
"""Get total storage size."""
pass
@abstractmethod
def close(self) -> None:
"""Close the client and cleanup resources."""
pass
@abstractmethod
def flush(self) -> None:
"""Flush data to disk."""
pass
logger = logging.getLogger(__name__)
class Hf3fsMockClient(Hf3fsClient):
"""Mock implementation of Hf3fsClient for CI testing purposes."""
def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
"""Initialize mock HF3FS client."""
self.path = path
self.size = size
self.bytes_per_page = bytes_per_page
self.entries = entries
# Create directory if it doesn't exist
os.makedirs(os.path.dirname(self.path), exist_ok=True)
# Create and initialize the file
self.file = os.open(self.path, os.O_RDWR | os.O_CREAT)
os.ftruncate(self.file, size)
logger.info(
f"Hf3fsMockClient initialized: path={path}, size={size}, "
f"bytes_per_page={bytes_per_page}, entries={entries}"
)
def batch_read(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
"""Batch read from mock storage."""
self.check(offsets, tensors)
results = []
for offset, tensor in zip(offsets, tensors):
size = tensor.numel() * tensor.itemsize
try:
os.lseek(self.file, offset, os.SEEK_SET)
bytes_read = os.read(self.file, size)
if len(bytes_read) == size:
# Convert bytes to tensor and copy to target
bytes_tensor = torch.frombuffer(bytes_read, dtype=torch.uint8)
typed_tensor = bytes_tensor.view(tensor.dtype).view(tensor.shape)
tensor.copy_(typed_tensor)
results.append(size)
else:
logger.warning(
f"Short read: expected {size}, got {len(bytes_read)}"
)
results.append(len(bytes_read))
except Exception as e:
logger.error(f"Error reading from offset {offset}: {e}")
results.append(0)
return results
def batch_write(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
"""Batch write to mock storage."""
self.check(offsets, tensors)
results = []
for offset, tensor in zip(offsets, tensors):
size = tensor.numel() * tensor.itemsize
try:
# Convert tensor to bytes and write directly to file
tensor_bytes = tensor.contiguous().view(torch.uint8).flatten()
data = tensor_bytes.numpy().tobytes()
os.lseek(self.file, offset, os.SEEK_SET)
bytes_written = os.write(self.file, data)
if bytes_written == size:
results.append(size)
else:
logger.warning(f"Short write: expected {size}, got {bytes_written}")
results.append(bytes_written)
except Exception as e:
logger.error(f"Error writing to offset {offset}: {e}")
results.append(0)
return results
def check(self, offsets: List[int], tensors: List[torch.Tensor]) -> None:
"""Validate batch operation parameters."""
pass
def get_size(self) -> int:
"""Get total storage size."""
return self.size
def close(self) -> None:
"""Close the mock client and cleanup resources."""
try:
if hasattr(self, "file") and self.file >= 0:
os.close(self.file)
self.file = -1 # Mark as closed
logger.info(f"MockHf3fsClient closed: {self.path}")
except Exception as e:
logger.error(f"Error closing MockHf3fsClient: {e}")
def flush(self) -> None:
"""Flush data to disk."""
try:
os.fsync(self.file)
except Exception as e:
logger.error(f"Error flushing MockHf3fsClient: {e}")
@@ -0,0 +1,220 @@
import datetime
import logging
import multiprocessing
import os
import threading
from functools import wraps
from pathlib import Path
from typing import List
import torch
from torch.utils.cpp_extension import load
from sglang.srt.mem_cache.storage.hf3fs.hf3fs_client import Hf3fsClient
root = Path(__file__).parent.resolve()
hf3fs_utils = load(name="hf3fs_utils", sources=[f"{root}/hf3fs_utils.cpp"])
logger = logging.getLogger(__name__)
HF3FS_AVAILABLE = True
try:
from hf3fs_fuse.io import (
deregister_fd,
extract_mount_point,
make_ioring,
make_iovec,
register_fd,
)
except ImportError:
HF3FS_AVAILABLE = False
def rsynchronized():
def _decorator(func):
@wraps(func)
def wrapper(self, *args, **kwargs):
with self.rlock:
return func(self, *args, **kwargs)
return wrapper
return _decorator
def wsynchronized():
def _decorator(func):
@wraps(func)
def wrapper(self, *args, **kwargs):
with self.wlock:
return func(self, *args, **kwargs)
return wrapper
return _decorator
class Hf3fsUsrBioClient(Hf3fsClient):
"""HF3FS client implementation using usrbio."""
def __init__(
self,
path: str,
size: int,
bytes_per_page: int,
entries: int,
client_timeout: int,
):
if not HF3FS_AVAILABLE:
raise ImportError(
"hf3fs_fuse.io is not available. Please install the hf3fs_fuse package."
)
self.path = path
self.size = size
self.bytes_per_page = bytes_per_page
self.entries = entries
self.client_timeout = client_timeout
self.file = os.open(self.path, os.O_RDWR | os.O_CREAT)
os.ftruncate(self.file, size)
register_fd(self.file)
self.hf3fs_mount_point = extract_mount_point(path)
self.bs = self.bytes_per_page
self.shm_r = multiprocessing.shared_memory.SharedMemory(
size=self.bs * self.entries, create=True
)
self.shm_w = multiprocessing.shared_memory.SharedMemory(
size=self.bs * self.entries, create=True
)
self.shm_r_tensor = torch.frombuffer(self.shm_r.buf, dtype=torch.uint8)
self.shm_w_tensor = torch.frombuffer(self.shm_w.buf, dtype=torch.uint8)
self.numa = -1
self.ior_r = make_ioring(
self.hf3fs_mount_point,
self.entries,
for_read=True,
timeout=1,
numa=self.numa,
)
self.ior_w = make_ioring(
self.hf3fs_mount_point,
self.entries,
for_read=False,
timeout=1,
numa=self.numa,
)
self.iov_r = make_iovec(self.shm_r, self.hf3fs_mount_point)
self.iov_w = make_iovec(self.shm_w, self.hf3fs_mount_point)
self.shm_r.unlink()
self.shm_w.unlink()
self.rlock = threading.RLock()
self.wlock = threading.RLock()
@rsynchronized()
def batch_read(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
self.check(offsets, tensors)
results = [0] * len(offsets)
# prepare
current = 0
for offset, tensor in zip(offsets, tensors):
size = tensor.numel() * tensor.itemsize
try:
self.ior_r.prepare(
self.iov_r[current : current + size], True, self.file, offset
)
current += size
except Exception as e:
logger.error(f"Error preparing batch read: {e}")
return results
# submit
ionum = len(offsets)
try:
resv = self.ior_r.submit().wait(
min_results=ionum,
timeout=datetime.timedelta(seconds=self.client_timeout),
)
except Exception as e:
logger.error(f"Error submitting batch read: {e}")
return results
# results
try:
hf3fs_utils.read_shm(self.shm_r_tensor, tensors)
results = [res.result for res in resv]
except Exception as e:
logger.error(f"[Hf3fsUsrBioClient] read_shm failed: {e}", exc_info=True)
return results
return results
@wsynchronized()
def batch_write(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
self.check(offsets, tensors)
results = [0] * len(offsets)
# prepare
hf3fs_utils.write_shm(tensors, self.shm_w_tensor)
current = 0
for offset, tensor in zip(offsets, tensors):
size = tensor.numel() * tensor.itemsize
try:
self.ior_w.prepare(
self.iov_w[current : current + size], False, self.file, offset
)
current += size
except Exception as e:
logger.error(f"Error preparing batch write: {e}")
return results
# submit
ionum = len(offsets)
try:
resv = self.ior_w.submit().wait(
min_results=ionum,
timeout=datetime.timedelta(seconds=self.client_timeout),
)
except Exception as e:
logger.error(f"Error submitting batch write: {e}")
return results
# results
results = [res.result for res in resv]
return results
def check(self, offsets: List[int], tensors: List[torch.Tensor]) -> None:
sizes = [t.numel() * t.itemsize for t in tensors]
if any(
[
len(offsets) > self.entries,
len(offsets) != len(sizes),
all(
[
offset < 0 or offset + size > self.size
for offset, size in zip(offsets, sizes)
]
),
all([size > self.bytes_per_page for size in sizes]),
]
):
self.close()
raise ValueError(f"Hf3fsClient.check: {offsets=}, {sizes=}")
def get_size(self) -> int:
return self.size
def close(self) -> None:
deregister_fd(self.file)
os.close(self.file)
del self.ior_r
del self.ior_w
del self.iov_r
del self.iov_w
self.shm_r.close()
self.shm_w.close()
def flush(self) -> None:
os.fsync(self.file)
@@ -0,0 +1,35 @@
#include <torch/extension.h>
#include <cstring>
#include <vector>
void read_shm(const torch::Tensor &shm, std::vector<torch::Tensor> dst) {
py::gil_scoped_release release;
char *src_ptr = static_cast<char *>(shm.data_ptr());
size_t current = 0;
for (size_t i = 0; i < dst.size(); ++i) {
auto &t = dst[i];
size_t t_bytes = t.numel() * t.element_size();
char *dst_ptr = static_cast<char *>(t.data_ptr());
std::memcpy(dst_ptr, src_ptr + current, t_bytes);
current += t_bytes;
}
}
void write_shm(const std::vector<torch::Tensor> src, torch::Tensor &shm) {
py::gil_scoped_release release;
char *dst_ptr = static_cast<char *>(shm.data_ptr());
size_t current = 0;
for (size_t i = 0; i < src.size(); ++i) {
auto &t = src[i];
size_t t_bytes = t.numel() * t.element_size();
char *src_ptr = static_cast<char *>(t.data_ptr());
std::memcpy(dst_ptr + current, src_ptr, t_bytes);
current += t_bytes;
}
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("read_shm", &read_shm, "Read tensors from shared memory");
m.def("write_shm", &write_shm, "Write tensors to shared memory");
}
@@ -0,0 +1,532 @@
import argparse
import atexit
import json
import logging
import threading
from collections import OrderedDict
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import orjson
import requests
from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.responses import ORJSONResponse
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from sglang.srt.mem_cache.hicache_storage import PoolName
from sglang.srt.mem_cache.storage.hf3fs.storage_hf3fs import Hf3fsMetadataInterface
# --- Configuration ---
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
# --- Data Models ---
class RankMetadata:
"""Holds all metadata for a single rank."""
def __init__(self, num_pages: int):
self.lock = threading.Lock()
self.num_pages = num_pages
self.free_pages: List[int] = list(range(num_pages))
self.key_to_index: OrderedDict[str, int] = OrderedDict()
# Todo: Support multi files for HF3FS
def exists_keys(self, keys: List[str]) -> List[bool]:
"""Check if keys exist in metadata."""
with self.lock:
return [key in self.key_to_index for key in keys]
def reserve_and_allocate_page_indices(
self, keys: List[Tuple[str, str]]
) -> List[Tuple[bool, int]]:
"""Reserve and allocate page indices for keys."""
with self.lock:
results = [None] * len(keys)
new_keys_to_process = []
for i, (key, prefix_key) in enumerate(keys):
if key in self.key_to_index:
results[i] = (True, self.key_to_index[key])
self.key_to_index.move_to_end(key)
else:
new_keys_to_process.append((i, key, prefix_key))
# Todo: Implementing data eviction logic after HiCache supports prefix information pass-through
for i, key, prefix_key in new_keys_to_process:
if len(self.free_pages) > 0:
page_index = self.free_pages.pop()
else:
page_index = self.key_to_index.popitem(last=False)[1]
results[i] = (False, page_index)
return results
def confirm_write(
self,
written_keys_to_confirm: List[Tuple[str, int]],
pages_to_release: List[int],
) -> None:
"""Confirm write operations and release pages."""
with self.lock:
for key, page_index in written_keys_to_confirm:
self.key_to_index[key] = page_index
self.key_to_index.move_to_end(key)
for page_index in pages_to_release:
if page_index not in self.free_pages:
self.free_pages.append(page_index)
def delete_keys(self, keys: List[str]) -> int:
"""Delete keys and return count of deleted keys."""
with self.lock:
count = 0
for key in keys:
if key in self.key_to_index:
page_index = self.key_to_index.pop(key)
if page_index not in self.free_pages:
self.free_pages.append(page_index)
count += 1
return count
def clear_all(self) -> None:
"""Clear all metadata."""
with self.lock:
self.free_pages = list(range(self.num_pages))
self.key_to_index.clear()
def get_page_indices(self, keys: List[str]) -> List[Optional[int]]:
"""Get page indices for keys."""
with self.lock:
results = []
for key in keys:
if key in self.key_to_index:
results.append(self.key_to_index[key])
self.key_to_index.move_to_end(key)
else:
results.append(None)
return results
class GlobalMetadataState:
"""Manages the state for all ranks and persistence."""
def __init__(self, persistence_path: Optional[str], save_interval: int):
self.global_lock = threading.RLock()
self.ranks: Dict[str, RankMetadata] = {}
self.persistence_path = Path(persistence_path) if persistence_path else None
self.save_interval = save_interval
self.save_timer: Optional[threading.Timer] = None
self.is_shutting_down = False
def load_from_disk(self):
if not self.persistence_path or not self.persistence_path.exists():
logging.info("Persistence file not found. Starting with a clean state.")
return
logging.info(f"Loading state from {self.persistence_path}")
try:
with open(self.persistence_path, "r") as f:
persisted_data = json.load(f)
with self.global_lock:
for key_str, data in persisted_data.items():
if ":" not in key_str:
key_str = f"{key_str}:kv" # For backward compatibility
num_pages = data["num_pages"]
rank_meta = RankMetadata(num_pages)
rank_meta.free_pages = data["free_pages"]
rank_meta.key_to_index = OrderedDict(data["key_to_index"])
self.ranks[key_str] = rank_meta
logging.info(
f"Successfully loaded metadata for {len(self.ranks)} ranks."
)
except (json.JSONDecodeError, KeyError, TypeError) as e:
logging.error(
f"Failed to load or parse persistence file: {e}. Starting fresh.",
exc_info=True,
)
self.ranks.clear()
def save_to_disk(self):
if not self.persistence_path:
return
logging.info("Persisting metadata to disk...")
with self.global_lock:
serializable_state = {}
for key_str, rank_meta in self.ranks.items():
with rank_meta.lock:
serializable_state[key_str] = {
"num_pages": rank_meta.num_pages,
"free_pages": rank_meta.free_pages,
"key_to_index": list(rank_meta.key_to_index.items()),
}
try:
temp_path = self.persistence_path.with_suffix(".tmp")
with open(temp_path, "w") as f:
json.dump(serializable_state, f, indent=4)
temp_path.rename(self.persistence_path)
logging.info(f"Metadata successfully persisted to {self.persistence_path}")
except Exception as e:
logging.error(f"Failed to save metadata to disk: {e}", exc_info=True)
def schedule_save(self):
if self.is_shutting_down or not self.persistence_path:
return
self.save_to_disk()
self.save_timer = threading.Timer(self.save_interval, self.schedule_save)
self.save_timer.start()
def shutdown(self):
logging.info("Shutting down metadata server...")
self.is_shutting_down = True
if self.save_timer:
self.save_timer.cancel()
self.save_to_disk()
logging.info("Shutdown complete.")
# --- Global MetadataServer implementation ---
class Hf3fsMetadataServer:
"""HF3FS Metadata Server that manages metadata for multiple ranks."""
def __init__(self, persistence_path: Optional[str] = None, save_interval: int = 60):
self.state = GlobalMetadataState(persistence_path, save_interval)
self.app = FastAPI(default_response_class=ORJSONResponse)
self._setup_routes()
def _setup_routes(self):
"""Setup FastAPI routes."""
self.app.post("/{rank}/initialize")(self.initialize)
self.app.post("/{rank}/exists")(self.exists)
self.app.post("/{rank}/reserve_and_allocate_page_indices")(
self.reserve_and_allocate_page_indices
)
self.app.post("/{rank}/confirm_write")(self.confirm_write)
self.app.post("/{rank}/delete_keys")(self.delete_keys)
self.app.post("/{rank}/clear")(self.clear)
self.app.post("/{rank}/get_page_indices")(self.get_page_indices)
def _rank_key(self, rank: int, namespace: str) -> str:
"""Generate the composite key for rank+namespace."""
return f"{rank}:{namespace}"
def get_rank_metadata(self, rank: int, namespace: str = "kv") -> RankMetadata:
"""Get rank metadata with proper error handling."""
key = self._rank_key(rank, namespace)
if key not in self.state.ranks:
raise HTTPException(
status_code=404,
detail=f"Rank {rank} namespace '{namespace}' not initialized. Please call /{rank}/initialize first.",
)
return self.state.ranks[key]
async def _read_json(self, request: Request) -> dict:
"""Parse request JSON using orjson if available."""
body = await request.body()
return orjson.loads(body)
def _json_response(self, content: dict):
"""Return ORJSONResponse when available to bypass jsonable_encoder."""
return ORJSONResponse(content)
async def initialize(self, rank: int, request: Request):
"""Initialize a rank with specified number of pages."""
data = await self._read_json(request)
num_pages = data["num_pages"]
namespace = data.get("namespace", "kv")
key = self._rank_key(rank, namespace)
with self.state.global_lock:
if key in self.state.ranks:
logging.info(
f"Rank {rank} namespace '{namespace}' already exists. Initialization request ignored."
)
if self.state.ranks[key].num_pages != num_pages:
logging.warning(
f"Rank {rank} namespace '{namespace}' initialized with different num_pages. Existing: {self.state.ranks[key].num_pages}, New: {num_pages}"
)
else:
logging.info(
f"Initializing new Rank {rank} namespace '{namespace}' with {num_pages} pages."
)
self.state.ranks[key] = RankMetadata(num_pages)
return Response(status_code=204)
async def exists(self, rank: int, request: Request):
"""Check if keys exist in metadata."""
data = await self._read_json(request)
keys = data["keys"]
namespace = data.get("namespace", "kv")
metadata = self.get_rank_metadata(rank, namespace)
results = metadata.exists_keys(keys)
return self._json_response({"exists": results})
async def reserve_and_allocate_page_indices(self, rank: int, request: Request):
"""Reserve and allocate page indices for keys."""
data = await self._read_json(request)
namespace = data.get("namespace", "kv")
metadata = self.get_rank_metadata(rank, namespace)
keys = data["keys"]
results = metadata.reserve_and_allocate_page_indices(keys)
return self._json_response({"indices": results})
async def confirm_write(self, rank: int, request: Request):
"""Confirm write operations and release pages."""
data = await self._read_json(request)
namespace = data.get("namespace", "kv")
metadata = self.get_rank_metadata(rank, namespace)
success_written_keys = data.get("written_keys_to_confirm", [])
released_pages = data.get("pages_to_release", [])
metadata.confirm_write(success_written_keys, released_pages)
return Response(status_code=204)
async def delete_keys(self, rank: int, request: Request):
"""Delete keys from metadata."""
data = await self._read_json(request)
namespace = data.get("namespace", "kv")
metadata = self.get_rank_metadata(rank, namespace)
count = metadata.delete_keys(data["keys"])
return Response(status_code=204)
async def clear(self, rank: int, request: Request):
"""Clear all metadata for a rank."""
data = await self._read_json(request)
namespace = data.get("namespace", "kv")
metadata = self.get_rank_metadata(rank, namespace)
metadata.clear_all()
return Response(status_code=204)
async def get_page_indices(self, rank: int, request: Request):
"""Get page indices for keys."""
data = await self._read_json(request)
namespace = data.get("namespace", "kv")
metadata = self.get_rank_metadata(rank, namespace)
keys = data["keys"]
results = metadata.get_page_indices(keys)
return self._json_response({"indices": results})
def run(self, host: str = "0.0.0.0", port: int = 18000):
"""Run the metadata server."""
self.state.load_from_disk()
if self.state.persistence_path:
self.state.schedule_save()
atexit.register(self.state.shutdown)
import uvicorn
logging.info(f"Starting metadata server on http://{host}:{port}")
if self.state.persistence_path:
logging.info(
f"Persistence is ENABLED. Saving to '{self.state.persistence_path}' every {self.state.save_interval} seconds."
)
else:
logging.info("Persistence is DISABLED.")
uvicorn.run(self.app, host=host, port=port)
# --- Client implementation ---
class Hf3fsGlobalMetadataClient(Hf3fsMetadataInterface):
"""Global http metadata client for HF3FS."""
def __init__(self, base_url: str, max_retries: int = 3):
self.base_url = base_url.rstrip("/")
self._session = requests.Session()
retry_strategy = Retry(
total=max_retries,
backoff_factor=0.3,
status_forcelist=[500, 502, 503, 504],
allowed_methods=["GET", "POST"],
)
adapter = HTTPAdapter(
max_retries=retry_strategy, pool_connections=256, pool_maxsize=256
)
self._session.mount("http://", adapter)
def _post(self, endpoint: str, json_data: dict) -> dict:
try:
url = f"{self.base_url}/{endpoint}"
headers = {"Content-Type": "application/json"}
payload = orjson.dumps(json_data) # type: ignore[union-attr]
response = self._session.post(url, data=payload, headers=headers)
response.raise_for_status()
if response.status_code == 204 or not response.content:
return {}
return orjson.loads(response.content) # type: ignore[union-attr]
except requests.exceptions.RequestException as e:
logging.error(f"Failed to POST to {endpoint} after retries: {e}")
raise RuntimeError(f"Failed to connect to metadata server: {e}") from e
def initialize(
self, rank: int, num_pages: int, namespace: PoolName = PoolName.KV
) -> None:
self._post(
f"{rank}/initialize", {"num_pages": num_pages, "namespace": str(namespace)}
)
def reserve_and_allocate_page_indices(
self, rank: int, keys: List[Tuple[str, str]], namespace: PoolName = PoolName.KV
) -> List[Tuple[bool, int]]:
response = self._post(
f"{rank}/reserve_and_allocate_page_indices",
{"keys": keys, "namespace": str(namespace)},
)
return [tuple(item) for item in response.get("indices")]
def confirm_write(
self,
rank: int,
written_keys_to_confirm: List[Tuple[str, int]],
pages_to_release: List[int],
namespace: PoolName = PoolName.KV,
) -> None:
self._post(
f"{rank}/confirm_write",
{
"written_keys_to_confirm": written_keys_to_confirm,
"pages_to_release": pages_to_release,
"namespace": str(namespace),
},
)
def delete_keys(
self, rank: int, keys: List[str], namespace: PoolName = PoolName.KV
) -> None:
self._post(f"{rank}/delete_keys", {"keys": keys, "namespace": str(namespace)})
def exists(
self, rank: int, keys: List[str], namespace: PoolName = PoolName.KV
) -> List[bool]:
response = self._post(
f"{rank}/exists", {"keys": keys, "namespace": str(namespace)}
)
return response.get("exists", [])
def clear(self, rank: int, namespace: PoolName = PoolName.KV) -> None:
self._post(f"{rank}/clear", {"namespace": str(namespace)})
def get_page_indices(
self, rank: int, keys: List[str], namespace: PoolName = PoolName.KV
) -> List[Optional[int]]:
response = self._post(
f"{rank}/get_page_indices", {"keys": keys, "namespace": str(namespace)}
)
return response.get("indices")
class Hf3fsLocalMetadataClient(Hf3fsMetadataInterface):
"""Local metadata client that directly operates on RankMetadata in memory without metadata server."""
def __init__(self):
self._metadata: Dict[str, RankMetadata] = {} # key: "rank:namespace"
def _ns_key(self, rank: int, namespace: PoolName) -> str:
return f"{rank}:{namespace}"
def _get_metadata(self, rank: int, namespace) -> RankMetadata:
key = self._ns_key(rank, namespace)
if key not in self._metadata:
raise RuntimeError(
f"Namespace '{namespace}' for rank {rank} not initialized"
)
return self._metadata[key]
def initialize(
self, rank: int, num_pages: int, namespace: PoolName = PoolName.KV
) -> None:
key = self._ns_key(rank, namespace)
if key not in self._metadata:
self._metadata[key] = RankMetadata(num_pages)
def reserve_and_allocate_page_indices(
self, rank: int, keys: List[Tuple[str, str]], namespace: PoolName = PoolName.KV
) -> List[Tuple[bool, int]]:
"""Reserve and allocate page indices for keys."""
return self._get_metadata(rank, namespace).reserve_and_allocate_page_indices(
keys
)
def confirm_write(
self,
rank: int,
written_keys_to_confirm: List[Tuple[str, int]],
pages_to_release: List[int],
namespace: PoolName = PoolName.KV,
) -> None:
"""Confirm write operations."""
self._get_metadata(rank, namespace).confirm_write(
written_keys_to_confirm, pages_to_release
)
def delete_keys(
self, rank: int, keys: List[str], namespace: PoolName = PoolName.KV
) -> None:
"""Delete keys."""
self._get_metadata(rank, namespace).delete_keys(keys)
def exists(
self, rank: int, keys: List[str], namespace: PoolName = PoolName.KV
) -> List[bool]:
"""Check if keys exist."""
return self._get_metadata(rank, namespace).exists_keys(keys)
def clear(self, rank: int, namespace: PoolName = PoolName.KV) -> None:
"""Clear all metadata for rank."""
self._get_metadata(rank, namespace).clear_all()
def get_page_indices(
self, rank: int, keys: List[str], namespace: PoolName = PoolName.KV
) -> List[Optional[int]]:
"""Get page indices for keys."""
return self._get_metadata(rank, namespace).get_page_indices(keys)
def run_metadata_server(
host: str = "0.0.0.0",
port: int = 18000,
persistence_path: Optional[str] = None,
save_interval: int = 60,
):
"""Run the HF3FS metadata server."""
global server
server = Hf3fsMetadataServer(
persistence_path=persistence_path, save_interval=save_interval
)
server.run(host=host, port=port)
# --- Main Execution ---
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="HF3FS Metadata Server")
parser.add_argument(
"--host", type=str, default="0.0.0.0", help="Host to bind the server to."
)
parser.add_argument(
"--port", type=int, default=18000, help="Port to run the server on."
)
parser.add_argument(
"--persistence-path",
type=str,
default=None,
help="Path to the file for persisting metadata. If not provided, persistence is disabled.",
)
parser.add_argument(
"--save-interval",
type=int,
default=60,
help="Interval in seconds for periodically saving metadata to disk.",
)
args = parser.parse_args()
run_metadata_server(args.host, args.port, args.persistence_path, args.save_interval)
@@ -0,0 +1,956 @@
import atexit
import concurrent.futures
import json
import logging
import os
import signal
import threading
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import wraps
from typing import Any, List, Optional, Tuple
import torch
from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorage,
HiCacheStorageConfig,
HiCacheStorageExtraInfo,
PoolHitPolicy,
PoolName,
PoolTransfer,
PoolTransferResult,
)
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
logger = logging.getLogger(__name__)
class Hf3fsMetadataInterface(ABC):
"""Interface for HF3FS metadata operations."""
@abstractmethod
def initialize(
self, rank: int, num_pages: int, namespace: PoolName = PoolName.KV
) -> None:
"""Initialize the metadata service with specified number of pages."""
pass
@abstractmethod
def reserve_and_allocate_page_indices(
self,
rank: int,
keys: List[Tuple[str, str]],
namespace: PoolName = PoolName.KV,
) -> List[Tuple[bool, int]]:
"""
Reserve and allocate page indices for the specified keys.
Args:
rank: The rank of the process.
keys: The keys to reserve and allocate page indices for. Each tuple contains a key and the key of its prefix block.
namespace: The namespace (pool type) for the metadata.
Returns:
List[Tuple[bool, int]]: A list of tuples, where each tuple contains a boolean indicating whether the key has existed and an integer indicating the allocated page index.
"""
pass
@abstractmethod
def confirm_write(
self,
rank: int,
written_keys_to_confirm: List[Tuple[str, int]],
pages_to_release: List[int],
namespace: PoolName = PoolName.KV,
) -> None:
"""
Confirm that key-value pairs have been successfully written to storage.
Args:
rank: The rank of the process.
written_keys_to_confirm: A list of tuples, where each tuple contains a key and its corresponding page index.
pages_to_release: A list of page indices to be released.
namespace: The namespace (pool type) for the metadata.
"""
pass
@abstractmethod
def get_page_indices(
self, rank: int, keys: List[str], namespace: PoolName = PoolName.KV
) -> List[Optional[int]]:
"""
Get page indices for the specified keys.
Args:
rank: The rank of the process.
keys: A list of keys.
namespace: The namespace (pool type) for the metadata.
Returns:
List[Optional[int]]: A list of integers representing the page indices for the specified keys.
If a key is not found, the corresponding index will be None.
"""
pass
@abstractmethod
def delete_keys(
self, rank: int, keys: List[str], namespace: PoolName = PoolName.KV
) -> None:
"""Delete specified keys and their associated pages."""
pass
@abstractmethod
def exists(
self, rank: int, keys: List[str], namespace: PoolName = PoolName.KV
) -> List[bool]:
"""Check if the specified keys exist."""
pass
@abstractmethod
def clear(self, rank: int, namespace: PoolName = PoolName.KV) -> None:
"""Clear all key-value pairs and page allocations for the specified rank."""
pass
class AtomicCounter:
def __init__(self, n: int):
assert n > 0
self.n = n
self._value = 0
self._lock = threading.Lock()
def next(self) -> int:
with self._lock:
current = self._value
self._value = (current + 1) % self.n
return current
def synchronized():
def _decorator(func):
@wraps(func)
def wrapper(self, *args, **kwargs):
with self.lock:
return func(self, *args, **kwargs)
return wrapper
return _decorator
def create_hf3fs_client(
path: str,
size: int,
bytes_per_page: int,
entries: int,
client_timeout: int,
use_mock: bool = False,
) -> Hf3fsClient:
"""Factory function to create appropriate HF3FS client.
Args:
path: File path for storage
size: Total size of storage file
bytes_per_page: Bytes per page
entries: Number of entries for batch operations
use_mock: Whether to use mock client instead of real usrbio client
Returns:
"""
if use_mock:
from sglang.srt.mem_cache.storage.hf3fs.hf3fs_client import Hf3fsMockClient
logger.info(f"[Rank Using Hf3fsMockClient for testing")
return Hf3fsMockClient(path, size, bytes_per_page, entries)
else:
from sglang.srt.mem_cache.storage.hf3fs.hf3fs_usrbio_client import (
Hf3fsUsrBioClient,
)
return Hf3fsUsrBioClient(path, size, bytes_per_page, entries, client_timeout)
@dataclass
class _PoolStorageCtx:
"""Per-pool storage context for hybrid KV cache pools."""
pool_name: str
bytes_per_page: int
num_pages: int
namespace: PoolName
clients: List[Hf3fsClient]
gb_per_page: float
class HiCacheHF3FS(HiCacheStorage):
"""HiCache backend that stores KV cache pages in HF3FS files."""
default_env_var: str = "SGLANG_HICACHE_HF3FS_CONFIG_PATH"
def __init__(
self,
rank: int,
file_path: str,
file_size: int,
numjobs: int,
bytes_per_page: int,
entries: int,
client_timeout: int,
dtype: torch.dtype,
metadata_client: Hf3fsMetadataInterface,
is_mla_model: bool = False,
is_page_first_layout: bool = False,
use_mock_client: bool = False,
enable_storage_metrics: bool = False,
):
self.rank = rank
self.file_path = file_path
self.file_size = file_size
self.numjobs = numjobs
self.bytes_per_page = bytes_per_page
self.gb_per_page = bytes_per_page / (1 << 30)
self.entries = entries
self.client_timeout = client_timeout
self.dtype = dtype
self.metadata_client = metadata_client
self.is_mla_model = is_mla_model
self.is_page_first_layout = is_page_first_layout
self.enable_storage_metrics = enable_storage_metrics
self.use_mock_client = use_mock_client
self.numel = self.bytes_per_page // self.dtype.itemsize
self.num_pages = self.file_size // self.bytes_per_page
self.skip_backup = False
if self.is_mla_model and self.rank != 0:
self.skip_backup = True
self.rank = 0
self.is_zero_copy = False
logger.info(
f"[Rank {self.rank}] HiCacheHF3FS Client Initializing: "
f"file_path={self.file_path}, "
f"file_size={self.file_size / (2 ** 30):.2f} GB, "
f"num_pages={self.num_pages}, "
f"is_mla_model={self.is_mla_model}"
)
self.ac = AtomicCounter(self.numjobs)
self.clients = [
create_hf3fs_client(
self.file_path,
self.file_size,
self.bytes_per_page,
self.entries,
self.client_timeout,
use_mock_client,
)
for _ in range(numjobs)
]
self.executor = concurrent.futures.ThreadPoolExecutor(
max_workers=self.numjobs, thread_name_prefix=f"HiCacheHF3FS-Rank{self.rank}"
)
self.metadata_client.initialize(self.rank, self.num_pages)
self.lock = threading.RLock()
self._pool_storage_ctx: dict = {}
atexit.register(self.close)
signal.signal(signal.SIGINT, lambda sig, frame: self.close())
signal.signal(signal.SIGTERM, lambda sig, frame: self.close())
signal.signal(signal.SIGQUIT, lambda sig, frame: self.close())
self.prefetch_pgs = []
self.backup_pgs = []
self.prefetch_bandwidth = []
self.backup_bandwidth = []
@staticmethod
def from_env_config(
bytes_per_page: int,
dtype: torch.dtype,
storage_config: HiCacheStorageConfig = None,
) -> "HiCacheHF3FS":
"""Create a HiCacheHF3FS instance from environment configuration.
Environment:
- Uses env var stored in `HiCacheHF3FS.default_env_var` to locate a JSON config.
- Falls back to a local single-machine config when the env var is not set.
Raises:
ValueError: If MLA Model is requested without global metadata server or required keys are missing.
"""
from sglang.srt.mem_cache.storage.hf3fs.mini_3fs_metadata_server import (
Hf3fsGlobalMetadataClient,
Hf3fsLocalMetadataClient,
)
use_mock_client = False
if storage_config is not None:
rank, is_mla_model, is_page_first_layout = (
storage_config.tp_rank,
storage_config.is_mla_model,
storage_config.is_page_first_layout,
)
if storage_config.extra_config is not None:
use_mock_client = storage_config.extra_config.get(
"use_mock_hf3fs_client", False
)
else:
rank, is_mla_model, is_page_first_layout = (
0,
False,
False,
)
mla_unsupported_msg = f"MLA model is not supported without global metadata server, please refer to https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/mem_cache/storage/hf3fs/docs/deploy_sglang_3fs_multinode.md"
config_path = os.getenv(HiCacheHF3FS.default_env_var)
if not config_path:
if is_mla_model:
raise ValueError(mla_unsupported_msg)
return HiCacheHF3FS(
rank=rank,
file_path=f"/data/hicache.{rank}.bin",
file_size=1 << 40,
numjobs=16,
bytes_per_page=bytes_per_page,
entries=8,
client_timeout=5,
dtype=dtype,
metadata_client=Hf3fsLocalMetadataClient(),
is_page_first_layout=is_page_first_layout,
use_mock_client=use_mock_client,
)
try:
with open(config_path, "r") as f:
config = json.load(f)
except Exception as e:
raise RuntimeError(f"Failed to load config from {config_path}: {str(e)}")
# Check required keys (metadata_server_url is now optional)
required_keys = {
"file_path_prefix",
"file_size",
"numjobs",
"entries",
}
missing_keys = required_keys - set(config.keys())
if missing_keys:
raise ValueError(f"Missing required keys in config: {missing_keys}")
# Choose metadata client based on configuration
if config.get("metadata_server_url"):
# Use global metadata client to connect to metadata server
metadata_server_url = config["metadata_server_url"]
metadata_client = Hf3fsGlobalMetadataClient(metadata_server_url)
logger.info(
f"Using global metadata client with server url: {metadata_server_url}"
)
else:
# Enable MLA optimization only when using the global metadata client
if is_mla_model:
raise ValueError(mla_unsupported_msg)
# Use local metadata client for single-machine deployment
metadata_client = Hf3fsLocalMetadataClient()
rank_for_path = 0 if is_mla_model else rank
return HiCacheHF3FS(
rank=rank,
# Let all ranks use the same file path for MLA model
file_path=f"{config['file_path_prefix']}.{rank_for_path}.bin",
file_size=int(config["file_size"]),
numjobs=int(config["numjobs"]),
bytes_per_page=bytes_per_page,
entries=int(config["entries"]),
client_timeout=config.get("client_timeout", 5),
dtype=dtype,
metadata_client=metadata_client,
is_mla_model=is_mla_model,
is_page_first_layout=is_page_first_layout,
use_mock_client=use_mock_client,
enable_storage_metrics=storage_config.enable_storage_metrics,
)
def _batch_get(
self,
keys: List[str],
values: List[torch.Tensor],
) -> List[bool]:
page_indices = self.metadata_client.get_page_indices(self.rank, keys)
if len(page_indices) != len(keys):
logger.error(
f"[Rank {self.rank}] HiCacheHF3FS get: page_indices length {len(page_indices)} mismatch keys length {len(keys)}."
)
return [False] * len(keys)
batch_indices, file_offsets = [], []
for i, page_index in enumerate(page_indices):
if page_index is not None:
batch_indices.append(i)
file_offsets.append(page_index * self.bytes_per_page)
for target_location in values:
assert target_location.is_contiguous()
file_results = values
start_time = time.perf_counter()
futures = [
self.executor.submit(
self.clients[self.ac.next()].batch_read,
file_offsets[i : i + self.entries],
file_results[i : i + self.entries],
)
for i in range(0, len(batch_indices), self.entries)
]
read_results = [result for future in futures for result in future.result()]
end_time = time.perf_counter()
ionum = len(batch_indices)
if self.enable_storage_metrics:
self.prefetch_pgs.append(ionum)
self.prefetch_bandwidth.append(
ionum / (end_time - start_time) * self.gb_per_page
)
results = [False] * len(keys)
for batch_index, read_result in zip(batch_indices, read_results):
if read_result == self.bytes_per_page:
results[batch_index] = True
else:
logger.error(
f"[Rank {self.rank}] HiCacheHF3FS get {keys[batch_index]} failed"
)
return results
def _batch_set(
self,
keys: List[str],
values: Optional[Any] = None,
) -> List[bool]:
# In MLA backend, only one rank needs to backup the KV cache
if self.skip_backup:
return True
# Todo: Add prefix block's hash key
key_with_prefix = [(key, "") for key in keys]
indices = self.metadata_client.reserve_and_allocate_page_indices(
self.rank, key_with_prefix
)
if len(indices) != len(keys):
logger.error(
f"[Rank {self.rank}] HiCacheHF3FS batch_get: mismatched lengths {len(indices)} != {len(keys)}"
)
# free allocated pages
if indices:
self.metadata_client.confirm_write(
self.rank, [], [index[1] for index in indices]
)
return [False] * len(keys)
batch_indices, file_offsets, file_values = [], [], []
pages_to_release = []
for i, (value, (is_written, page_index)) in enumerate(zip(values, indices)):
if is_written or page_index == -1:
continue
batch_indices.append(i)
file_offsets.append(page_index * self.bytes_per_page)
assert value.is_contiguous()
file_values.append(value)
start_time = time.perf_counter()
futures = [
self.executor.submit(
self.clients[self.ac.next()].batch_write,
file_offsets[i : i + self.entries],
file_values[i : i + self.entries],
)
for i in range(0, len(batch_indices), self.entries)
]
write_results = [
result == self.bytes_per_page
for future in futures
for result in future.result()
]
end_time = time.perf_counter()
ionum = len(batch_indices)
if self.enable_storage_metrics:
self.backup_pgs.append(ionum)
self.backup_bandwidth.append(
ionum / (end_time - start_time) * self.gb_per_page
)
written_keys_to_confirm = []
results = [index[0] for index in indices]
for batch_index, write_result in zip(batch_indices, write_results):
key = keys[batch_index]
page_index = indices[batch_index][1]
if write_result:
written_keys_to_confirm.append((key, page_index))
else:
logger.error(f"[Rank {self.rank}] HiCacheHF3FS set {key} failed")
pages_to_release.append(page_index)
results[batch_index] = write_result
if len(written_keys_to_confirm) > 0 or len(pages_to_release) > 0:
self.metadata_client.confirm_write(
self.rank, written_keys_to_confirm, pages_to_release
)
return results
def delete(self, key: str) -> None:
self.metadata_client.delete_keys(self.rank, [key])
def exists(self, key: str) -> bool:
result = self.metadata_client.exists(self.rank, [key])
return result[0] if result else False
def batch_exists(
self, keys: List[str], extra_info: Optional[HiCacheStorageExtraInfo] = None
) -> int:
factor = 1
if self.mha_zero_copy:
keys = self._get_mha_zero_copy_keys(keys)
factor = 2
results = self.metadata_client.exists(self.rank, keys)
i = 0
while i < len(keys) and results[i]:
i += 1
return i // factor
def clear(self) -> None:
try:
self.metadata_client.clear(self.rank)
for ctx in getattr(self, "_pool_storage_ctx", {}).values():
self.metadata_client.clear(self.rank, namespace=ctx.namespace)
logger.info(f"Cleared HiCacheHF3FS for rank {self.rank}")
except Exception as e:
logger.error(f"Failed to clear HiCacheHF3FS: {e}")
def close(self) -> None:
try:
for c in self.clients:
c.close()
for ctx in getattr(self, "_pool_storage_ctx", {}).values():
for c in ctx.clients:
c.close()
self.executor.shutdown(wait=True)
except Exception as e:
logger.error(f"close HiCacheHF3FS: {e}")
logger.info("close HiCacheHF3FS")
def get_stats(self):
storage_metrics = StorageMetrics()
storage_metrics.prefetch_pgs.extend(self.prefetch_pgs)
storage_metrics.backup_pgs.extend(self.backup_pgs)
storage_metrics.prefetch_bandwidth.extend(self.prefetch_bandwidth)
storage_metrics.backup_bandwidth.extend(self.backup_bandwidth)
self.prefetch_pgs.clear()
self.backup_pgs.clear()
self.prefetch_bandwidth.clear()
self.backup_bandwidth.clear()
return storage_metrics
def register_mem_pool_host(self, mem_pool_host: HostKVCache):
super().register_mem_pool_host(mem_pool_host)
self.is_zero_copy = self.mem_pool_host.layout in [
"page_first",
"page_first_direct",
]
self.mha_zero_copy = self.is_zero_copy and not self.is_mla_model
logger.info(f"{self.is_zero_copy=}, layout={self.mem_pool_host.layout}")
def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name):
if host_pool_name == PoolName.KV:
return
super().register_mem_host_pool_v2(host_pool, host_pool_name)
pool_page_size = getattr(host_pool, "page_size", 1) or 1
pool_bytes_per_page = host_pool.get_ksize_per_token() * pool_page_size
pool_num_pages = self.file_size // pool_bytes_per_page
pool_file_path = f"{self.file_path}.{host_pool_name}"
namespace = host_pool_name # e.g. PoolName.MAMBA, PoolName.INDEXER
pool_clients = [
create_hf3fs_client(
pool_file_path,
self.file_size,
pool_bytes_per_page,
self.entries,
self.client_timeout,
self.use_mock_client,
)
for _ in range(self.numjobs)
]
self.metadata_client.initialize(self.rank, pool_num_pages, namespace=namespace)
self._pool_storage_ctx[host_pool_name] = _PoolStorageCtx(
pool_name=host_pool_name,
bytes_per_page=pool_bytes_per_page,
num_pages=pool_num_pages,
namespace=namespace,
clients=pool_clients,
gb_per_page=pool_bytes_per_page / (1 << 30),
)
logger.info(
f"[Rank {self.rank}] Registered hybrid pool '{host_pool_name}': "
f"bytes_per_page={pool_bytes_per_page}, num_pages={pool_num_pages}, "
f"namespace={namespace}, file={pool_file_path}"
)
def _get_mha_zero_copy_keys(self, keys: List[str]) -> List[str]:
_keys = []
for k in keys:
_keys.append(f"{k}-k")
_keys.append(f"{k}-v")
return _keys
def _get_mha_zero_copy_values(
self, values: List[torch.Tensor]
) -> List[torch.Tensor]:
_values = []
for value in values:
_values.append(value[0])
_values.append(value[1])
return _values
def _batch_get_preprocess(self, keys, host_indices):
page_num = len(host_indices) // self.mem_pool_host.page_size
# host_indices to kv_buffer
flat = not self.is_zero_copy
values = (
[
self.mem_pool_host.get_data_page(
host_indices[i * self.mem_pool_host.page_size], flat=flat
)
for i in range(page_num)
]
if self.is_zero_copy
else [
self.mem_pool_host.get_dummy_flat_data_page() for _ in range(page_num)
]
)
if self.mha_zero_copy:
keys = self._get_mha_zero_copy_keys(keys)
values = self._get_mha_zero_copy_values(values)
return keys, values
def _batch_get_postprocess(self, host_indices, values, results):
page_num = len(host_indices) // self.mem_pool_host.page_size
if self.is_zero_copy:
if not self.is_mla_model:
results = [
(results[2 * i] and results[2 * i + 1]) for i in range(page_num)
]
results = results[:page_num]
return results
for i in range(page_num):
if not results[i]:
break
self.mem_pool_host.set_from_flat_data_page(
host_indices[i * self.mem_pool_host.page_size], values[i]
)
return results
def batch_exists_v2(
self,
keys: List[str],
pool_transfers: Optional[List[PoolTransfer]] = None,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> PoolTransferResult:
kv_pages = self.batch_exists(keys, extra_info)
hit_count: dict = {PoolName.KV: kv_pages} if kv_pages else {}
final_pages = kv_pages
for transfer in pool_transfers or []:
if final_pages == 0:
break
pool_name = transfer.name
ctx = self._pool_storage_ctx.get(pool_name)
if ctx is None:
final_pages = 0
break
component_keys = [f"{key}_{pool_name}" for key in keys[:kv_pages]]
exists_results = self.metadata_client.exists(
self.rank, component_keys, namespace=ctx.namespace
)
boundary = 0
if transfer.hit_policy == PoolHitPolicy.ALL_PAGES:
try:
boundary = exists_results.index(False)
except ValueError:
boundary = kv_pages
elif transfer.hit_policy == PoolHitPolicy.TRAILING_PAGES:
trailing = max(1, len(transfer.keys) if transfer.keys else 1)
for prefix_len in range(kv_pages, 0, -1):
if all(
exists_results[i]
for i in range(max(0, prefix_len - trailing), prefix_len)
):
boundary = prefix_len
break
if boundary:
hit_count[pool_name] = boundary
final_pages = min(final_pages, boundary)
return PoolTransferResult(final_pages, hit_count)
def _pool_batch_get(self, transfer: PoolTransfer) -> List[bool]:
pool_name = transfer.name
ctx = self._pool_storage_ctx[pool_name]
host_pool = self.registered_pools[pool_name]
keys = transfer.keys
host_indices = transfer.host_indices
page_size = getattr(host_pool, "page_size", 1) or 1
page_num = len(keys)
component_keys = [f"{key}_{pool_name}" for key in keys]
page_indices = self.metadata_client.get_page_indices(
self.rank, component_keys, namespace=ctx.namespace
)
batch_indices, file_offsets, values = [], [], []
for i, page_index in enumerate(page_indices):
if page_index is not None:
batch_indices.append(i)
file_offsets.append(page_index * ctx.bytes_per_page)
values.append(host_pool.get_dummy_flat_data_page())
if not batch_indices:
return [False] * page_num
start_time = time.perf_counter()
futures = [
self.executor.submit(
ctx.clients[self.ac.next()].batch_read,
file_offsets[j : j + self.entries],
values[j : j + self.entries],
)
for j in range(0, len(batch_indices), self.entries)
]
read_results = [r for f in futures for r in f.result()]
end_time = time.perf_counter()
ionum = len(batch_indices)
if self.enable_storage_metrics:
self.prefetch_pgs.append(ionum)
self.prefetch_bandwidth.append(
ionum / (end_time - start_time) * ctx.gb_per_page
)
results = [False] * page_num
for idx, (batch_idx, read_result) in enumerate(
zip(batch_indices, read_results)
):
if read_result == ctx.bytes_per_page:
host_idx = host_indices[batch_idx * page_size].item()
host_pool.set_from_flat_data_page(host_idx, values[idx])
results[batch_idx] = True
else:
logger.error(
f"[Rank {self.rank}][Pool {pool_name.upper()}] HiCacheHF3FS get {keys[batch_idx]} failed"
)
return results
def _pool_batch_set(self, transfer: PoolTransfer) -> List[bool]:
pool_name = transfer.name
ctx = self._pool_storage_ctx[pool_name]
host_pool = self.registered_pools[pool_name]
keys = transfer.keys
host_indices = transfer.host_indices
page_size = getattr(host_pool, "page_size", 1) or 1
page_num = len(keys)
component_keys = [f"{key}_{pool_name}" for key in keys]
key_with_prefix = [(k, "") for k in component_keys]
indices = self.metadata_client.reserve_and_allocate_page_indices(
self.rank, key_with_prefix, namespace=ctx.namespace
)
if len(indices) != page_num:
logger.error(
f"[Rank {self.rank}] Pool {pool_name}: mismatched indices length"
)
if indices:
self.metadata_client.confirm_write(
self.rank, [], [idx[1] for idx in indices], namespace=ctx.namespace
)
return [False] * page_num
batch_indices, file_offsets, file_values = [], [], []
for i, (is_written, page_index) in enumerate(indices):
if is_written or page_index == -1:
continue
batch_indices.append(i)
file_offsets.append(page_index * ctx.bytes_per_page)
host_idx = host_indices[i * page_size].item()
data = host_pool.get_data_page(host_idx, flat=True)
assert data.is_contiguous()
file_values.append(data)
start_time = time.perf_counter()
futures = [
self.executor.submit(
ctx.clients[self.ac.next()].batch_write,
file_offsets[j : j + self.entries],
file_values[j : j + self.entries],
)
for j in range(0, len(batch_indices), self.entries)
]
write_results = [r == ctx.bytes_per_page for f in futures for r in f.result()]
end_time = time.perf_counter()
ionum = len(batch_indices)
if self.enable_storage_metrics:
self.backup_pgs.append(ionum)
self.backup_bandwidth.append(
ionum / (end_time - start_time) * ctx.gb_per_page
)
written_keys_to_confirm = []
pages_to_release = []
results = [idx[0] for idx in indices]
for batch_idx, write_ok in zip(batch_indices, write_results):
key = component_keys[batch_idx]
page_index = indices[batch_idx][1]
if write_ok:
written_keys_to_confirm.append((key, page_index))
else:
logger.error(
f"[Rank {self.rank}][Pool {pool_name.upper()}] HiCacheHF3FS set {keys[batch_idx]} failed"
)
pages_to_release.append(page_index)
results[batch_idx] = write_ok
if written_keys_to_confirm or pages_to_release:
self.metadata_client.confirm_write(
self.rank,
written_keys_to_confirm,
pages_to_release,
namespace=ctx.namespace,
)
return results
def batch_get_v2(
self,
transfers: List[PoolTransfer],
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> dict:
results = {}
for transfer in transfers:
results[transfer.name] = self._pool_batch_get(transfer)
return results
def batch_set_v2(
self,
transfers: List[PoolTransfer],
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> dict:
results = {}
for transfer in transfers:
results[transfer.name] = self._pool_batch_set(transfer)
return results
def batch_get_v1(
self,
keys: List[str],
host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]:
keys, values = self._batch_get_preprocess(keys, host_indices)
results = self._batch_get(keys, values)
return self._batch_get_postprocess(host_indices, values, results)
def _batch_set_preprocess(self, keys, host_indices):
page_num = len(host_indices) // self.mem_pool_host.page_size
# host_indices to kv_buffer
flat = not self.is_zero_copy
values = [
self.mem_pool_host.get_data_page(
host_indices[i * self.mem_pool_host.page_size], flat=flat
)
for i in range(page_num)
]
if self.mha_zero_copy:
keys = self._get_mha_zero_copy_keys(keys)
values = self._get_mha_zero_copy_values(values)
return keys, values
def batch_set_v1(
self,
keys: List[str],
host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]:
len_keys = len(keys)
keys, values = self._batch_set_preprocess(keys, host_indices)
results = self._batch_set(keys, values)
return results
# Deprecated
def get(
self,
key: str,
target_location: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> torch.Tensor | None:
pass
# Deprecated
def batch_get(
self,
keys: List[str],
target_locations: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> List[torch.Tensor | None] | int:
pass
# Deprecated
def set(
self,
key: str,
value: Optional[Any] = None,
target_location: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> bool:
pass
# Deprecated
def batch_set(
self,
keys: List[str],
values: Optional[Any] = None,
target_locations: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> bool:
pass
@@ -0,0 +1,44 @@
import multiprocessing.shared_memory
import sys
from pathlib import Path
import pytest
import torch
from torch.utils.cpp_extension import load
from tqdm import tqdm
root = Path(__file__).parent.resolve()
hf3fs_utils = load(
name="hf3fs_utils", sources=[f"{root}/hf3fs_utils.cpp"], verbose=True
)
def test_rw_shm():
numel = 8 << 20
dtype = torch.bfloat16
page_num = 128
page_bytes = numel * dtype.itemsize
shm = multiprocessing.shared_memory.SharedMemory(
size=page_num * page_bytes, create=True
)
tshm = torch.frombuffer(shm.buf, dtype=torch.uint8)
a = [
torch.randn(numel, dtype=dtype)
for _ in tqdm(range(page_num), desc="prepare input")
]
b = [
torch.empty(numel, dtype=dtype)
for _ in tqdm(range(page_num), desc="prepare output")
]
hf3fs_utils.write_shm(a, tshm)
hf3fs_utils.read_shm(tshm, b)
for _a, _b in tqdm(zip(a, b), desc="assert_close"):
torch.testing.assert_close(_a, _b)
del tshm
shm.close()
shm.unlink()
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
@@ -0,0 +1,71 @@
# LMCache Connector for SGLang
This document describes how to use LMCache as KV Cache Management Backend for SGLang engine.
For more details about LMCache, please refer to: https://lmcache.ai
## Install LMCache
### Method 1: with pip
```bash
pip install lmcache
```
### Method 2: from source
Clone LMCache project:
```bash
git clone https://github.com/LMCache/LMCache
```
Install:
```bash
cd LMCache
pip install -e . --no-build-isolation
```
## Use LMCache
LMCache supports two transport modes. **MP (multi-process, default)** issues a single blocking retrieve over ZMQ to a standalone daemon that owns the KV store and survives SGLang restarts. **IP (in-process)** uses an embedded layerwise connector — the cache lives and dies with the SGLang process. Mode selection is currently a code-level setting in `LMCRadixCache.__init__` (`self._mode`); only MP is reachable by default.
### MP mode (default): multi-process daemon
Uses `LMCacheMPConnector`. Daemon host/port come from the LMCache YAML config (`mp_host`, `mp_port`).
Terminal 1 — start the LMCache daemon:
```bash
lmcache server \
--host 127.0.0.1 --port 5556 \
--l1-size-gb 4 \
--eviction-policy LRU
```
Use the bundled `example_config_mp.yaml` (or any YAML setting `mp_host` / `mp_port`):
Terminal 2 — start SGLang:
```bash
python -m sglang.launch_server \
--model-path MODEL \
--enable-lmcache \
--lmcache-config-file example_config_mp.yaml
```
For full LMCache config options see https://docs.lmcache.ai/api_reference/configurations.html.
### IP mode: in-process
Uses `LMCacheLayerwiseConnector`. KV transfer happens per layer inside the SGLang process; the cache lives and dies with the server. To enable, edit `LMCRadixCache.__init__` and set `self._mode = LMCacheMode.IP`.
The LMCache config still controls chunk_size and storage; `mp_host` / `mp_port` are ignored on this path. Use the bundled `example_config_ip.yaml`:
```bash
python -m sglang.launch_server \
--model-path MODEL \
--enable-lmcache \
--lmcache-config-file example_config_ip.yaml
```
@@ -0,0 +1,7 @@
# Basic configurations
chunk_size: 256
# CPU offloading configurations
local_cpu: true
use_layerwise: true
max_local_cpu_size: 10 # number of CPU backend GB
@@ -0,0 +1,3 @@
# MP mode: SGLang dials the standalone `lmcache server` at this host/port.
mp_host: 127.0.0.1
mp_port: 5556
@@ -0,0 +1,503 @@
from __future__ import annotations
import enum
import logging
import threading
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional, Tuple
import torch
from sglang.srt.mem_cache.base_prefix_cache import (
EvictParams,
EvictResult,
InitLoadBackParams,
MatchPrefixParams,
MatchResult,
)
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
from sglang.srt.runtime_context import get_server_args
try:
from lmcache.integration.sglang.multi_process_adapter import LMCacheMPConnector
from lmcache.integration.sglang.sglang_adapter import (
LMCacheLayerwiseConnector,
LoadMetadata,
StoreMetadata,
)
from lmcache.integration.sglang.utils import lmcache_get_config
except ImportError as e:
raise RuntimeError(
"LMCache is not installed. Please install it by running `pip install lmcache`"
) from e
if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
logger = logging.getLogger(__name__)
@dataclass
class _LMCacheLoadBackMarker:
"""Carries the data ``init_load_back`` needs from the
``match_prefix`` call in MP mode.
"""
key: RadixKey # detached snapshot of the matched key (the live query key
# aliases the req's growing fill_ids and must not be retained)
value_numel: int # number of tokens already in radix at match time
class LMCacheMode(enum.Enum):
MP = enum.auto() # multi-process mode
IP = enum.auto() # in-process mode
class LayerTransferCounter:
"""Minimal adapter that lets the memory pool notify LMCache per-layer.
The KV pool calls `wait_until(layer_id)` after finishing a layer, which we
translate into a `load_kv_layerwise(layer_id)` call on the LMCache connector
within the provided CUDA stream.
"""
def __init__(
self,
num_layers: int,
load_stream: torch.cuda.Stream,
lmc_connector: LMCacheLayerwiseConnector,
printable: bool = False,
):
self.num_layers = num_layers
self.load_stream = load_stream
self.lmc_connector = lmc_connector
def wait_until(self, layer_id: int):
# Ensure ordering of the async loads wrt compute stream(s).
self.load_stream.synchronize()
with self.load_stream:
self.lmc_connector.load_kv_layerwise(layer_id)
class LMCRadixCache(RadixCache):
"""RadixCache + LMCache IO.
IP mode keeps the existing layerwise connector and
its per-layer transfer hook: ``match_prefix`` kicks off the load via
``start_load_kv`` and SGLang's per-layer KV-pool hook drives subsequent
layers during forward.
MP mode uses ``LMCacheMPConnector`` with a two-phase
load: ``match_prefix`` fires LOOKUP only (``connector.lookup_kv``) and
returns ``host_hit_length`` on the ``MatchResult``; the SGLang
scheduler then calls `init_load_back` at dispatch time,
which fires the actual RETRIEVE (``connector.retrieve_kv``) into
pre-allocated GPU slots.
"""
def __init__(
self,
params: CacheInitParams,
model_config: Optional[ModelConfig] = None,
tp_size: int = 1,
rank: int = 0,
tp_group: Optional[torch.distributed.ProcessGroup] = None,
):
super().__init__(params)
cli_lmc_cfg = get_server_args().lmcache_config_file or ""
kvcache = self.token_to_kv_pool_allocator.get_kvcache()
connector_kwargs = dict(
sgl_config=model_config,
tp_size=tp_size,
rank=rank,
# NOTE: The original implementation accessed private buffers via
# `_kvcache.k_buffer` / `.v_buffer`. We prefer public accessors when
# available; fall back to private fields if needed.
k_pool=getattr(
kvcache,
"k_buffer",
getattr(self.token_to_kv_pool_allocator._kvcache, "k_buffer"),
),
v_pool=getattr(
kvcache,
"v_buffer",
getattr(self.token_to_kv_pool_allocator._kvcache, "v_buffer"),
),
tp_group=tp_group.device_group if tp_group is not None else None,
)
self.load_stream = torch.cuda.Stream()
self.store_stream = torch.cuda.Stream()
# MP is the default. To use the in-process layerwise connector,
# set ``self._mode = LMCacheMode.IP`` here.
self._mode = LMCacheMode.MP
if self._mode is LMCacheMode.MP:
if not cli_lmc_cfg:
raise ValueError(
"MP mode requires --lmcache-config-file (the YAML "
"supplies mp_host / mp_port)."
)
lm_cfg = lmcache_get_config(cli_lmc_cfg)
self.lmcache_connector = LMCacheMPConnector(
page_size=params.page_size,
host=lm_cfg.mp_host,
port=lm_cfg.mp_port,
**connector_kwargs,
)
elif self._mode is LMCacheMode.IP:
self.lmcache_connector = LMCacheLayerwiseConnector(
config_file=cli_lmc_cfg, **connector_kwargs
)
# Per-layer hook
self.layer_done_executor = LayerTransferCounter(
num_layers=(
model_config.num_hidden_layers if model_config is not None else 0
),
load_stream=self.load_stream,
lmc_connector=self.lmcache_connector,
)
kvcache.register_layer_transfer_counter(self.layer_done_executor)
self._in_flight_nodes: list[TreeNode] = []
self._node_lock = threading.Lock()
self._mp_load_back_markers: dict[str, _LMCacheLoadBackMarker] = {}
def reset(self):
super().reset()
if hasattr(self, "_in_flight_nodes"):
with self._node_lock:
self._in_flight_nodes.clear()
if hasattr(self, "_mp_load_back_markers"):
self._mp_load_back_markers.clear()
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
"""Dispatch to the mode-specific match_prefix.
MP mode → ``_mp_match_prefix`` (fires LOOKUP only).
IP mode → ``_ip_match_prefix`` (single-shot ``start_load_kv``
plus per-layer hook).
"""
key = params.key
if self.disable or not key:
return super().match_prefix(params)
if self.page_size != 1:
aligned_len = len(key) // self.page_size * self.page_size
key = key[:aligned_len]
base_res = super().match_prefix(params)
value: torch.Tensor = base_res.device_indices
last_node: TreeNode = base_res.last_device_node
if self._mode is LMCacheMode.MP:
if params.req is None:
return base_res
return self._mp_match_prefix(key, base_res, value, last_node, params.req)
elif self._mode is LMCacheMode.IP:
return self._ip_match_prefix(key, base_res, value, last_node)
return base_res
def _mp_match_prefix(
self,
key: RadixKey,
base_res: MatchResult,
value: torch.Tensor,
last_node: TreeNode,
req: Req,
) -> MatchResult:
"""MP LOOKUP
Returns a ``MatchResult`` with ``host_hit_length`` set when
LMCache has tokens beyond radix. Otherwise releases
the held read locks and returns the radix-only result.
"""
token_ids = key.raw_token_ids()
matched = self.lmcache_connector.lookup_kv(token_ids, req.rid)
if matched <= value.numel():
# Release the read locks; keep the pending session for end_session.
self.lmcache_connector.release_pending(req.rid)
return base_res
if token_ids is key.token_ids:
token_ids = token_ids[:]
self._mp_load_back_markers[req.rid] = _LMCacheLoadBackMarker(
key=RadixKey(token_ids, key.extra_key, key.is_bigram),
value_numel=int(value.numel()),
)
return MatchResult(
device_indices=value,
last_device_node=last_node,
last_host_node=last_node,
best_match_node=last_node,
host_hit_length=matched - int(value.numel()),
)
def _ip_match_prefix(
self,
key: RadixKey,
base_res: MatchResult,
value: torch.Tensor,
last_node: TreeNode,
) -> MatchResult:
"""IP mode: ``start_load_kv`` + per-layer hook.
Allocates slots for the page-aligned uncached tail and kicks off
the layerwise load. Returns ``base_res`` if there's nothing to
fetch or alloc/load fails.
"""
if value.numel() == len(key):
return base_res
uncached_len = len(key) - value.numel()
if uncached_len == 0:
return base_res
token_ids = key.raw_token_ids()
result = self._load_back(
key=key,
value_numel=int(value.numel()),
uncached_len=uncached_len,
last_node=last_node,
load_fn=lambda sm, pp: self._ip_load_back(
token_ids=token_ids,
value_numel=int(value.numel()),
slot_mapping=sm,
prefix_pad=pp,
),
)
if result is None:
return base_res
new_slots, new_node = result
return MatchResult(
device_indices=torch.cat([value, new_slots]),
last_device_node=new_node,
last_host_node=new_node,
best_match_node=new_node,
)
def init_load_back(
self, params: InitLoadBackParams
) -> Tuple[torch.Tensor, Optional[TreeNode]]:
"""MP RETRIEVE.
Called by the scheduler when ``match_prefix`` returned
``host_hit_length > 0``. Uses the cached LOOKUP result to
allocate slots and fire RETRIEVE, inserts the resulting
TreeNode into the radix tree, and returns
``(new_indices, new_last_node)``.
"""
req = params.req
marker = self._mp_load_back_markers.pop(req.rid)
last_node: TreeNode = params.best_match_node
result = self._load_back(
key=marker.key,
value_numel=marker.value_numel,
uncached_len=params.host_hit_length,
last_node=last_node,
load_fn=lambda sm, pp: self._mp_load_back(
marker=marker,
request_id=req.rid,
slot_mapping=sm,
prefix_pad=pp,
),
)
if result is None:
# Either alloc failed (locks still held by lookup_kv) or
# retrieve returned nothing (locks already released by
# retrieve_kv). release_pending is idempotent on locks_held.
self.lmcache_connector.release_pending(req.rid)
return (
torch.empty((0,), dtype=torch.int64, device=self.device),
last_node,
)
return result
def _load_back(
self,
*,
key: RadixKey,
value_numel: int,
uncached_len: int,
last_node: TreeNode,
load_fn, # Callable[[torch.Tensor, int], int] — (slot_mapping, prefix_pad) -> num_retrieved
) -> Optional[Tuple[torch.Tensor, TreeNode]]:
"""Alloc slots, run ``load_fn``, attach a TreeNode for what was loaded.
Returns ``(slots, new_node)`` on success, ``None`` if alloc fails
or the load returned zero (slots are freed in either case).
"""
chunk_size = self.lmcache_connector.chunk_size()
prefix_pad = value_numel % chunk_size
if self.token_to_kv_pool_allocator.available_size() < uncached_len:
self.evict(EvictParams(num_tokens=uncached_len))
token_slots = self.token_to_kv_pool_allocator.alloc(uncached_len)
if token_slots is None:
return None
slot_mapping = torch.empty(
value_numel + token_slots.numel(),
dtype=torch.int64,
device=self.device,
)
slot_mapping[:value_numel].fill_(-1)
slot_mapping[value_numel:].copy_(token_slots)
num_retrieved = load_fn(slot_mapping, prefix_pad)
logger.debug("num_retrieved_tokens: %s", num_retrieved)
if num_retrieved > 0:
self.token_to_kv_pool_allocator.free(
token_slots[(num_retrieved - prefix_pad) :]
)
else:
self.token_to_kv_pool_allocator.free(token_slots)
if num_retrieved > 0:
fetched = num_retrieved - prefix_pad
new_node = TreeNode(priority=last_node.priority)
start = value_numel
end = start + fetched
new_node.key = key[start:end]
new_node.value = token_slots[:fetched]
new_node.parent = last_node
last_node.children[new_node.key.child_key(self.page_size)] = new_node
self.evictable_size_ += fetched
self._update_leaf_status(last_node)
self._update_leaf_status(new_node)
self._record_store_event(new_node.parent)
self._record_store_event(new_node)
return token_slots[:fetched], new_node
return None
def _mp_load_back(
self,
*,
marker: _LMCacheLoadBackMarker,
request_id: str,
slot_mapping: torch.Tensor,
prefix_pad: int,
) -> int:
"""MP non-layerwise loader: fire ``retrieve_kv`` and wait for the
load_stream so the compute stream observes the writes.
"""
self.load_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(self.load_stream):
n = self.lmcache_connector.retrieve_kv(
LoadMetadata(
token_ids=marker.key.token_ids,
slot_mapping=slot_mapping,
offset=marker.value_numel - prefix_pad,
prefix_pad=prefix_pad,
request_id=request_id,
)
)
torch.cuda.current_stream().wait_stream(self.load_stream)
return n
def _ip_load_back(
self,
*,
token_ids: list[int],
value_numel: int,
slot_mapping: torch.Tensor,
prefix_pad: int,
) -> int:
"""IP layerwise loader: kick off ``start_load_kv`` on ``self.load_stream``.
``start_load_kv`` enqueues the first layer's transfer; the
``LayerTransferCounter`` hook drives the rest during forward.
"""
with torch.cuda.stream(self.load_stream):
return self.lmcache_connector.start_load_kv(
LoadMetadata(
token_ids=token_ids,
slot_mapping=slot_mapping,
offset=value_numel - prefix_pad,
)
)
def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:
"""On request completion, insert device KV into radix and store to LMCache."""
super().cache_finished_req(req, is_insert=is_insert)
if not is_insert:
if self._mode is LMCacheMode.MP:
self._mp_load_back_markers.pop(req.rid, None)
self.lmcache_connector.end_session(req.rid)
return
global_server_args = get_server_args()
topk = global_server_args.speculative_eagle_topk
enable_kv_committed_len = topk is None or topk == 1
if enable_kv_committed_len:
kv_committed_len = req.kv_committed_len
else:
kv_committed_len = len(req.origin_input_ids) + max(
len(req.output_ids) - 1, 0
)
token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len
]
# Use super() to avoid a redundant LOOKUP — we only need new_last_node from radix.
match_result = super().match_prefix(
MatchPrefixParams(key=RadixKey(token_ids, req.extra_key))
)
new_last_node = match_result.last_device_node
assert new_last_node is not None
self.inc_lock_ref(new_last_node)
store_md = StoreMetadata(
last_node=new_last_node,
token_ids=token_ids,
kv_indices=kv_indices,
offset=0,
request_id=req.rid,
)
with torch.cuda.stream(self.store_stream):
self.lmcache_connector.store_kv(store_md)
if self._mode is LMCacheMode.MP:
# MP store_kv blocks until the daemon's signal event fires, so the slots are safe to evict immediately.
self._mp_load_back_markers.pop(req.rid, None)
self.dec_lock_ref(new_last_node)
self.lmcache_connector.end_session(req.rid)
elif self._mode is LMCacheMode.IP:
# Layerwise store is async on store_stream; defer the unlock to evict()'s store_stream.synchronize().
with self._node_lock:
self._in_flight_nodes.append(new_last_node)
def evict(self, params: EvictParams) -> EvictResult:
"""Before base eviction, wait for any outstanding stores and release locks."""
if self.disable:
return EvictResult()
self.store_stream.synchronize()
with self._node_lock:
for node in self._in_flight_nodes:
self.dec_lock_ref(node)
self._in_flight_nodes.clear()
return super().evict(params)
def pretty_print(self):
super().pretty_print()
try:
logger.debug(
"evictable=%d protected=%d", self.evictable_size_, self.protected_size_
)
except Exception: # pragma: no cover
pass
@@ -0,0 +1,118 @@
try:
from lmcache.integration.sglang.sglang_adapter import (
LMCacheLayerwiseConnector,
LoadMetadata,
StoreMetadata,
)
except ImportError:
raise RuntimeError(
"LMCache is not installed. Please install it by running `pip install lmcache` in the root directory of LMCache"
)
import torch
from sglang.srt.configs.model_config import ModelConfig
def test_load_store_metadata():
model_config = ModelConfig(
model_path="Qwen/Qwen3-4B",
)
# Generate Dummy KV Cache
head_num = model_config.num_key_value_heads
head_dim = model_config.head_dim
layer_num = model_config.num_hidden_layers
buffer_size = 256
input_id_len = 16
k_buffer = [
torch.randn(buffer_size, head_num, head_dim, dtype=torch.bfloat16).cuda()
for _ in range(layer_num)
]
v_buffer = [
torch.randn(buffer_size, head_num, head_dim, dtype=torch.bfloat16).cuda()
for _ in range(layer_num)
]
connector = LMCacheLayerwiseConnector(
model_config, 1, 0, k_buffer, v_buffer, config_file="example_config_ip.yaml"
)
fake_token_ids = torch.randint(0, model_config.vocab_size, (input_id_len,)).tolist()
fake_kv_indices = torch.randint(0, buffer_size, (input_id_len,))
offset = 0
store_metadata = StoreMetadata(
last_node=None,
token_ids=fake_token_ids,
kv_indices=fake_kv_indices,
offset=offset,
)
load_metadata = LoadMetadata(
token_ids=fake_token_ids,
slot_mapping=fake_kv_indices,
offset=offset,
)
current_stream = torch.cuda.current_stream()
retrieve_token_num = connector.start_load_kv(load_metadata)
assert retrieve_token_num == 0
connector.store_kv(store_metadata)
current_stream.synchronize()
# check retrieve
gt_key_buffer = [
torch.zeros(input_id_len, head_num, head_dim, dtype=torch.bfloat16).cuda()
for _ in range(layer_num)
]
gt_value_buffer = [
torch.zeros(input_id_len, head_num, head_dim, dtype=torch.bfloat16).cuda()
for _ in range(layer_num)
]
for i in range(layer_num):
gt_key_buffer[i] = k_buffer[i][fake_kv_indices]
gt_value_buffer[i] = v_buffer[i][fake_kv_indices]
# clear the k_buffer and v_buffer
for _ in range(layer_num):
k_buffer[i].zero_()
v_buffer[i].zero_()
retrieve_token_num = connector.start_load_kv(load_metadata)
assert retrieve_token_num == input_id_len
for i in range(layer_num):
current_stream.synchronize()
connector.load_kv_layerwise(i)
current_stream.synchronize()
test_key_buffer = [
torch.zeros(input_id_len, head_num, head_dim, dtype=torch.bfloat16).cuda()
for _ in range(layer_num)
]
test_value_buffer = [
torch.zeros(input_id_len, head_num, head_dim, dtype=torch.bfloat16).cuda()
for _ in range(layer_num)
]
for i in range(layer_num):
test_key_buffer[i] = k_buffer[i][fake_kv_indices]
test_value_buffer[i] = v_buffer[i][fake_kv_indices]
for i in range(layer_num):
assert torch.allclose(test_key_buffer[i], gt_key_buffer[i])
assert torch.allclose(test_value_buffer[i], gt_value_buffer[i])
print("================================================")
print("TEST_LOAD_STORE_METADATA PASSED!")
print("================================================")
connector.close()
if __name__ == "__main__":
test_load_store_metadata()
@@ -0,0 +1,493 @@
# Mooncake as L3 KV Cache
This document describes how to use Mooncake as the L3 KV cache for SGLang.
Related documentation:
* [Quick Start: SGLang HiCache with Mooncake Backend](https://kvcache-ai.github.io/Mooncake/getting_started/examples/sglang-integration/hicache-quick-start.html)
* [Complete Guide: SGLang HiCache with Mooncake Backend](https://kvcache-ai.github.io/Mooncake/getting_started/examples/sglang-integration/hicache-integration-v1.html)
* [Mooncake x SGLang HiCache System Design](https://kvcache-ai.github.io/Mooncake/design/hicache-design.html)
* [HiCache System Design and Optimization](https://docs.sglang.io/advanced_features/hicache_design.html)
* [SGLang HiCache with Mooncake Backend Benchmark](https://kvcache-ai.github.io/Mooncake/performance/sglang-hicache-benchmark-results-v1.html)
## About Mooncake
Mooncake aims to enhance the inference efficiency of large language models (LLMs), especially in slow object storage environments, by constructing a multi-level caching pool on high-speed interconnected DRAM/SSD resources. Compared to traditional caching systems, Mooncake utilizes (GPUDirect) RDMA technology to transfer data directly in a zero-copy manner, while maximizing the use of multi-NIC resources on a single machine.
For more details about Mooncake, please refer to [Mooncake project](https://github.com/kvcache-ai/Mooncake) and [Mooncake documents](https://kvcache-ai.github.io/Mooncake/).
### Mooncake & SGLang HiCache
Mooncake serves as a high-performance L3 storage backend for SGLang HiCache, enabling distributed KV cache storage across multiple servers with RDMA-accelerated data transfer. This integration addresses the capacity limitations of traditional GPU-only or GPU+CPU caching by providing virtually unlimited cache storage through a distributed memory pool.
When a cache miss occurs in L1 and L2, HiCache automatically fetches the required KV cache from Mooncake's distributed memory pool. The system uses intelligent prefetching strategies to minimize latency, and utilize RDMA technology and zero-copy technique to ensure high-bandwidth, low-latency data transfer between SGLang instances and Mooncake storage nodes.
**Key Advantages:**
- **Scalable Capacity**: Aggregate memory across entire clusters into large distributed pools.
- **Cache Sharing**: KV caches can be shared by all SGLang instances in the cluster.
- **RDMA Acceleration**: Direct memory access eliminates CPU overhead and reduces latency.
- **Zero Copy**: Direct data transfer between L2 and Mooncake without intermediate copying, maximizing throughput.
- **Fault Tolerance**: Distributed architecture provides resilience against individual node failures.
This integration is particularly valuable for production deployments involving long-context models, multi-turn conversations, and high-throughput serving scenarios where traditional caching approaches become capacity-constrained.
## Install Mooncake
**Method 1: with pip**
```bash
pip install mooncake-transfer-engine
```
**Method 2: from source**
Clone Mooncake project:
```bash
git clone https://github.com/kvcache-ai/Mooncake --recursive
```
Install dependencies:
```bash
cd Mooncake
bash dependencies.sh
```
Build the project:
```bash
mkdir build
cd build
cmake ..
make -j
```
Install Mooncake:
```bash
sudo make install
```
For more details, please refer to [Mooncake official installation guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html).
## Deployment
**Mooncake** is a distributed system that efficiently aggregates memory resources across multiple servers. It can also be deployed on a single server for simpler setups.
When integrated with **SGLang**, the system conceptually consists of four key components: `the master service`, `metadata service` (Optional), `store service` (Optional), and the `SGLang server`. Among them, the `master service` and `metadata service` are responsible for object and metadata maintenance. The `store service` manages a contiguous memory segment that contributes to the distributed KV cache, making its memory accessible to both local and remote `SGLang servers`. Data transfer occurs directly between the `store service` and `SGLang servers`, bypassing the `master service`.
### Single Server Deployment
**Launch Mooncake `metadata service` (Optional):**
```bash
python -m mooncake.http_metadata_server
```
This service is responsible for centralized metadata management including internal connection status and related metadata.
Deployment of the `metadata service` can be skipped in the following cases:
* Mooncake supports non-centralized metadata management via a P2P handshake mechanism to exchange metadata. When using this mode, deployment of the `metadata service` can be skipped.
* Mooncake also supports embedding `mededata service` into `master service`. In this case, only the `master service` needs to be started.
**Launch Mooncake `master service`:**
The `master service` orchestrates the logical storage space pool across the entire cluster, managing KV cache space allocation and eviction.
To start `mooncake_master`:
```bash
mooncake_master --eviction_high_watermark_ratio=0.95
```
To start `mooncake_master` with embedded `metadata service` (so that a separate `metadata service` deployment can be skipped):
```bash
mooncake_master --enable_http_metadata_server=true --http_metadata_server_port=8080 --eviction_high_watermark_ratio=0.95
```
**Understanding `eviction_high_watermark_ratio`:**
When a `PutStart` request fails due to insufficient memory, or when the eviction thread detects that space usage has reached the configured high watermark ratio, an eviction task is triggered to free up space by evicting a portion of objects.
Due to memory fragmentation, allocation failures may occur even when memory usage has not yet reached 100%. The actual threshold depends on the workload. This [benchmark document](https://kvcache-ai.github.io/Mooncake/performance/allocator-benchmark-result.html) provides memory allocation efficiency results under different scenarios. if excessive allocation failures are observed, consider lowering this parameter accordingly.
**Launch Mooncake `store service` (Optional):**
First, create and save a configuration file in JSON format. For example:
```json
{
"local_hostname": "localhost",
"metadata_server": "http://127.0.0.1:8080/metadata",
"master_server_address": "127.0.0.1:50051",
"protocol": "rdma",
"device_name": "",
"global_segment_size": "4gb",
"local_buffer_size": 0
}
```
Note: If the `metadata service` is not deployed, set this field to:
```json
"metadata_server": "P2PHANDSHAKE",
```
Then start the `store service`:
```bash
python -m mooncake.mooncake_store_service --config=[config_path] --port=8081
```
Mooncake `store service` configuration can also be provided via environment variables:
```bash
MOONCAKE_LOCAL_HOSTNAME="localhost" \
MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata" \
MOONCAKE_MASTER="127.0.0.1:50051" \
MOONCAKE_PROTOCOL="rdma" \
MOONCAKE_DEVICE="" \
MOONCAKE_GLOBAL_SEGMENT_SIZE="4gb" \
MOONCAKE_LOCAL_BUFFER_SIZE=0 \
python -m mooncake.mooncake_store_service --port=8081
```
**Parameter Explanation:**
* `local_hostname`, `MOONCAKE_LOCAL_HOSTNAME`: The hostname of the `store service`.
* `metadata_server`, `MOONCAKE_TE_META_DATA_SERVER` : The network address of the `metadata service`. The default port is 8080. If the `metadata service` is not deployed, set this field to: `"metadata_server": "P2PHANDSHAKE"`.
* `master_server_address`, `MOONCAKE_MASTER`: The network address of the `master service`. The default port is 50051.
* `protocol`, `MOONCAKE_PROTOCOL`: The protocol used by Mooncake. Supported values are `"rdma"` or `"tcp"`. For optimal performance, `"rdma"` is recommended.
* `device_name`, `MOONCAKE_DEVICE`: The RDMA devices used by Mooncake. This field can usually be left empty, as Mooncake automatically discovers available NICs by default. This parameter is required only when the protocol is set to `"rdma"` **and** a specific set of NICs needs to be used. Example: `"device_name": "mlx5_0,mlx5_1"`. To list available devices, run `ibv_devices`. **Note:** If the environment variable `MC_MS_AUTO_DISC` is set to `1`, any `device_name` or `MOONCAKE_DEVICE` configuration will be overridden, and Mooncake will switch to auto-discovery mode.
- For tensor parallel deployments where different ranks should use different devices, you can specify device configurations using JSON format:
```json
{
"device_name": "{0: \"ib0,ib1\", 1: \"ib2,ib3\", 2: \"ib4,ib5\"}"
}
```
- Or in environment variables:
```bash
MOONCAKE_DEVICE="{\"0\": \"ib0,ib1\", \"1\": \"ib2,ib3\", \"2\": \"ib4,ib5\"}"
```
* `global_segment_size`, `MOONCAKE_GLOBAL_SEGMENT_SIZE`: The amount of memory contributed to the global memory pool. Accepts either bytes (integer) or a string with the `gb` suffix, e.g., `"4294967296"` or `"4gb"`. A larger value allows Mooncake to cache more KV tensors.
* `local_buffer_size`, `MOONCAKE_LOCAL_BUFFER_SIZE`: Local buffer is used to do request operations such as `Get` or `Put`. In this case, it is set to 0 because the instance functions solely as a storage server, contributing memory to the global pool without issuing any request operations.
**Important: Understanding Global Segment Size**
`global_segment_size` and `MOONCAKE_GLOBAL_SEGMENT_SIZE`: This parameter specifies the amount of memory each instance contributes to the distributed memory pool. The total memory available for KV cache storage across the cluster is the sum of the memory contributed by all instances.
Adjust this value according to systems available memory and expected cache requirements.
Note: If `MOONCAKE_GLOBAL_SEGMENT_SIZE` is set to a non-zero value when starting the `SGLang server`, launching the `store service` can be skipped. In this case, the `SGLang server` also takes on the role of the `store service`, which simplifies deployment but couples the two components together. Users can choose the deployment approach that best fits their needs.
**Start the `SGLang server` with Mooncake enabled:**
There are three ways to configure Mooncake:
1. Via extra configuration passed through sglang parameters
2. Using JSON configuration files
3. Using environment variables
Mooncake loads configuration in the following priority order:
1. If Mooncake-specific options are provided in `--hicache-storage-backend-extra-config`, they are used first.
2. If not, Mooncake checks whether the environment variable `DEFAULT_MOONCAKE_CONFIG_PATH_ENV` is set, and loads the JSON config file from that path.
3. If neither of the above is provided, Mooncake falls back to environment variables.
**Using extra-config of sglang arguments to configure Mooncake**
```bash
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend mooncake \
--model-path [model_path] \
--hicache-storage-backend-extra-config '{"master_server_address": "127.0.0.1:50051", "local_hostname": "localhost", "metadata_server": "http://127.0.0.1:8080/metadata", "global_segment_size": "4gb", "protocol": "rdma", "device_name": ""}'
```
**Using JSON file to configure Mooncake**
SGLang server can load Mooncake config from `SGLANG_HICACHE_MOONCAKE_CONFIG_PATH`.
```bash
export SGLANG_HICACHE_MOONCAKE_CONFIG_PATH=/sgl-workspace/sglang/benchmark/hicache/mooncake_config.json
echo '{
"local_hostname": "localhost",
"metadata_server": "http://127.0.0.1:8080/metadata",
"master_server_address": "127.0.0.1:50051",
"protocol": "rdma",
"device_name": "",
"global_segment_size": "4gb"
}' > ${SGLANG_HICACHE_MOONCAKE_CONFIG_PATH}
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend mooncake \
--model-path [model_path]
```
**Using env variables to configure Mooncake**
```bash
MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata" \
MOONCAKE_MASTER="127.0.0.1:50051" \
MOONCAKE_PROTOCOL="rdma" \
MOONCAKE_DEVICE="" \
MOONCAKE_GLOBAL_SEGMENT_SIZE="4gb" \
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend mooncake\
--model-path [model_path]
```
**Parameter Explanation:**
The Mooncake parameters used here are essentially the same as those configured for the `store service`.
In particular, for the `global segment size`, if at least one `store service` instance is running, this value can be set to `0`. In this case, the SGLang server will not contribute any memory to the system. Note that KV tensors stored in this contributed memory will be lost when the process exits; however, this will **not** cause any system errors.
**Important:** when `tp > 1`, each Tensor Parallel (TP) rank launches its own Mooncake backend instance and contributes `1/global_segment_size` memory. Therefore, the total memory consumption equals `global segment size`.
**SSD Offload (`enable_ssd_offload`):**
When `enable_ssd_offload` is set to `true`, SGLang will request that Mooncake enable SSD offloading for the KV cache. This allows Mooncake to spill overflow data from DRAM to local SSDs, effectively expanding the available L3 cache capacity.
If you need to explicitly control the SSD spill directory, set `ssd_offload_path` or the `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` environment variable. SGLang forwards this value to `MooncakeDistributedStore.setup(..., ssd_offload_path=...)`, while other SSD offload tuning parameters continue to be read directly by the Mooncake C++ library.
You can enable it in any of the three supported configuration methods:
- **Via `--hicache-storage-backend-extra-config`:**
```bash
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend mooncake \
--model-path [model_path] \
--hicache-storage-backend-extra-config '{"master_server_address": "127.0.0.1:50051", "enable_ssd_offload": true, "ssd_offload_path": "/mnt/mooncake-ssd"}'
```
- **Via JSON config file (`SGLANG_HICACHE_MOONCAKE_CONFIG_PATH`):**
```json
{
"master_server_address": "127.0.0.1:50051",
"enable_ssd_offload": true,
"ssd_offload_path": "/mnt/mooncake-ssd"
}
```
- **Via environment variable:**
```bash
MOONCAKE_MASTER="127.0.0.1:50051" \
MOONCAKE_ENABLE_SSD_OFFLOAD=1 \
MOONCAKE_OFFLOAD_FILE_STORAGE_PATH="/mnt/mooncake-ssd" \
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend mooncake \
--model-path [model_path]
```
> **Note:** `enable_ssd_offload` requires a Mooncake version that supports the `enable_ssd_offload` parameter in `MooncakeDistributedStore.setup()`. If the installed version does not support it, SGLang will automatically fall back to the old behavior and print a warning.
**Mooncake Group Semantics (`enable_group_semantics`):**
When `enable_group_semantics` is set to `true`, SGLang passes Mooncake `group_ids` for physical objects derived from the same logical HiCache page. This allows Mooncake to apply group-aware metadata routing, lease refresh, and eviction behavior to related KV objects such as MHA K/V pairs, split-head shards, MLA objects, and supported sidecar objects.
This option is disabled by default. It requires a Mooncake version that exposes `ReplicateConfig.group_ids`. If the installed Mooncake package does not support it, SGLang automatically falls back to the existing write path and prints a warning.
Example:
```bash
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend mooncake \
--model-path [model_path] \
--hicache-storage-backend-extra-config '{"master_server_address": "127.0.0.1:50051", "enable_group_semantics": true}'
```
**HiCache Related Parameters for SGLang Server**
For a comprehensive overview of HiCache-related parameters, please refer to [this document](https://docs.sglang.io/advanced_features/hicache_design.html#related-parameters).
Note that, for `--hicache-mem-layout {layer_first,page_first,page_first_direct}`,
the regular Mooncake backend path still uses `page_first` or `page_first_direct`.
When HiSparse provides an MLA host KV pool or DeepSeek V4 C4 side pool with
layer-first page metadata, Mooncake Store uses Mooncake's multi-buffer zero-copy
APIs (`batch_put_from_multi_buffers` / `batch_get_into_multi_buffers`) to store
each logical page across its per-layer buffers.
### Distributed Deployment
Distributed deployment of Mooncake is straightforward. Similar to the single-node setup, start one `metadata service` and one `master service` for this cluster. Then start a `store service` on each server.
Mooncake also supports high availability mode. This mode enhances fault tolerance by running the `master service` as a cluster of multiple master nodes coordinated through an `etcd` cluster. The master nodes use `etcd` to elect a leader, which is responsible for handling client requests. For more details about how to deploy in this mode, please refer to our [documents](https://kvcache-ai.github.io/Mooncake/).
### Deployment with Dummy Client (Experimental)
In addition to the standard deployment where SGLang acts as a full Mooncake node, you can use the **Dummy Client** mode. In this mode, SGLang connects to a local **Mooncake Store Service** (Real Client) via RPC/IPC. This decouples the SGLang process from the heavy RDMA and memory management, potentially improving stability and allowing the cache to persist even if the SGLang process restarts.
**Architecture:**
* **Mooncake Master**: Manages the cluster topology (same as standard).
* **Mooncake Store Service (Real Client)**: Manages the actual memory pool and RDMA connections. Must be running locally.
* **SGLang Server (Dummy Client)**: Connects to the local Store Service to access the cache.
#### 1. Launch Services (Master & Store)
First, start the `master service` and the `store service`. The `store service` acts as the Real Client.
**Start Master:**
```bash
mooncake_master --eviction_high_watermark_ratio=0.95
```
**Start Store Service (Real Client):** Crucially, the default port (50052) is used for internal RPC, which the Dummy Client will connect to.
```bash
mooncake_client --global_segment_size=4GB
```
**Parameter Explanation:**
- **`host`**: (string, default: "0.0.0.0"): The hostname of the client.
- **`port`**: (int, default: 50052): The port number the client service listens on.
- **`global_segment_size`**: (string, default: "4GB"): The size of the global segment to be allocated by the client.
- **`master_server_address`**: (string, default: "localhost:50051"): The address of the Master Service.
- **`metadata_server`**: (string, default: "http://localhost:8080/metadata"): The address of the metadata service.
- **`protocol`**: (string, default: "tcp"): The protocol used by the Transfer Engine.
- **`device_name`**: (string, default: ""): The device name used by the Transfer Engine.
- **`threads`**: (int, default: 1): The number of threads used by the client.
#### 2. Launch SGLang (Dummy Client)
Configure SGLang to connect to the Real Client using the client_server_address parameter.
**Using extra-config of sglang arguments to configure Mooncake**
```bash
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend mooncake \
--model-path [model_path] \
--hicache-storage-backend-extra-config '{"standalone_storage": true, "client_server_address": "127.0.0.1:50052"}'
```
**Using JSON file to configure Mooncake**
SGLang server can load Mooncake config from `SGLANG_HICACHE_MOONCAKE_CONFIG_PATH`.
```bash
export SGLANG_HICACHE_MOONCAKE_CONFIG_PATH=/sgl-workspace/sglang/benchmark/hicache/mooncake_config.json
echo '{
"standalone_storage": true,
"client_server_address": "127.0.0.1:50052"
}' > ${SGLANG_HICACHE_MOONCAKE_CONFIG_PATH}
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend mooncake \
--model-path [model_path]
```
**Using env variables to configure Mooncake**
```bash
MOONCAKE_STANDALONE_STORAGE=1
MOONCAKE_CLIENT="127.0.0.1:50052"
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend mooncake \
--model-path [model_path]
```
### Prefill/Decode Disaggregation
In **PD disaggregation**, the configurations for the `metadata service`, `mooncake master`, and the optional `store service` remain the same as described above. The difference is that SGLang introduces three distinct roles: `prefill worker`, `decode worker`, and `router`.
Among these, the `prefill worker` supports enabling **HiCache**. To run with PD disaggregation, start from the [PD configuration](https://kvcache-ai.github.io/Mooncake/getting_started/examples/sglang-integration-v1.html), and add the HiCache-related parameters (as previously described for the `SGLang server`) to the `prefill worker`.
In the example below, one `prefill worker`, one `decode worker`, and one `router` are launched. HiCache is enabled on the `prefill worker` to optimize prefill performance.
**Prefill worker**:
```bash
MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata" \
MOONCAKE_MASTER=127.0.0.1:50051 \
MOONCAKE_PROTOCOL="rdma" \
MOONCAKE_DEVICE="mlx5_1" \
MOONCAKE_GLOBAL_SEGMENT_SIZE=4294967296 \
python -m sglang.launch_server \
--model-path [model_path] \
--page-size 64 \
--enable-hierarchical-cache \
--hicache-storage-prefetch-policy timeout \
--hicache-storage-backend mooncake \
--disaggregation-mode prefill \
--disaggregation-ib-device "mlx5_1" \
--base-gpu-id 0 \
--port 30000
```
**Decode worker**:
```bash
python -m sglang.launch_server \
--model-path [model_path] \
--page-size 64 \
--disaggregation-mode decode \
--disaggregation-ib-device "mlx5_1" \
--base-gpu-id 1 \
--port 30001
```
**Router**:
```bash
python -m sglang_router.launch_router \
--pd-disaggregation \
--prefill "http://127.0.0.1:30000" \
--decode "http://127.0.0.1:30001" \
--host 0.0.0.0 \
--port 8000
```
## Troubleshooting
**RDMA Registration Failure:**
* In some environments, RDMA registration may require root privileges. In this case, try running the program as root.
* In certain environments (e.g., eRDMA), there is an upper limit on the total amount of RDMA memory that can be registered. Once this limit is exceeded, registration will fail. To resolve this, you can lower the value of `MOONCAKE_GLOBAL_SEGMENT_SIZE`, or reduce the host memory allocated to HiCache in the `SGLang server` (since this memory is fully registered with RDMA to enable zero-copy).
**HiCache CPU Memory Usage:**
When using HiCache, the default L2 host DRAM (CPU memory) size for KV cache is **2 times** the size of the L1 device memory (GPU memory) for KV cache.
If the model is small but the GPU memory is large — especially in multi-TP (tensor parallel) setups — this may cause the L1 KV cache to become very large, which in turn can consume excessive CPU DRAM.
In such cases, you should manually configure an appropriate L2 cache size based on your hardware. This can be done by setting `--hicache-ratio` or `--hicache-size`.
**More Information:**
Additional troubleshooting information can be found [here](https://kvcache-ai.github.io/Mooncake/troubleshooting/troubleshooting.html).
## Test Mooncake Store
This test is intended for developers to quickly verify that the MooncakeStore class interfaces are functioning correctly.
First, start the `metadata service` and `master service`. Then run the `test_mooncake_store.py`. 16MB global segments size is enough to run this test.
```bash
MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata" \
MOONCAKE_MASTER=127.0.0.1:50051 \
MOONCAKE_PROTOCOL="rdma" \
MOONCAKE_GLOBAL_SEGMENT_SIZE=16777216 \
python3 [path of test_mooncake_store.py]
```
If all tests pass, the message "✅ All tests passed" will be printed at the end.
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,110 @@
import logging
from typing import Any, List
from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import MooncakeBaseStore
logger = logging.getLogger(__name__)
class MooncakeEmbeddingStore(MooncakeBaseStore):
def __init__(
self,
storage_config: Any = None,
):
super().__init__()
MooncakeDistributedStore = self._import_mooncake_store()
self.store = MooncakeDistributedStore()
self.config = self._load_config(storage_config)
ret_code = self.store.setup(
self.config.local_hostname,
self.config.metadata_server,
self.config.global_segment_size,
16 * 1024 * 1024, # Internal local buffer size
self.config.protocol,
self.config.device_name,
self.config.master_server_address,
)
if ret_code != 0:
raise RuntimeError(f"Failed to setup Mooncake Embedding Store: {ret_code}")
logger.info("Mooncake Embedding Store initialized successfully.")
def get_key(self, mm_hash: str) -> str:
return f"emb_{mm_hash}"
def batch_get(
self, hashes: List[str], ptrs: List[int], sizes: List[int]
) -> List[bool]:
keys = [self.get_key(h) for h in hashes]
results = self.store.batch_get_into(keys, ptrs, sizes)
return [res > 0 for res in results]
def batch_put(
self, hashes: List[str], ptrs: List[int], sizes: List[int]
) -> List[bool]:
keys = [self.get_key(h) for h in hashes]
exists = self.store.batch_is_exist(keys)
put_keys, put_ptrs, put_sizes, indices = [], [], [], []
success_map = [True] * len(hashes)
for i, status in enumerate(exists):
if status != 1:
put_keys.append(keys[i])
put_ptrs.append(ptrs[i])
put_sizes.append(sizes[i])
indices.append(i)
if put_keys:
results = self.store.batch_put_from(put_keys, put_ptrs, put_sizes)
for i, res in enumerate(results):
success_map[indices[i]] = res == 0
return success_map
def batch_is_exist(self, hashes: List[str]) -> List[bool]:
keys = [self.get_key(h) for h in hashes]
results = self.store.batch_is_exist(keys)
return [res == 1 for res in results]
def batch_get_into_multi_buffers(
self,
hashes: List[str],
ptrs: List[List[int]],
sizes: List[List[int]],
) -> List[bool]:
keys = [self.get_key(h) for h in hashes]
results = self.store.batch_get_into_multi_buffers(keys, ptrs, sizes)
return [res > 0 for res in results]
def batch_put_from_multi_buffers(
self,
hashes: List[str],
ptrs: List[List[int]],
sizes: List[List[int]],
) -> List[bool]:
keys = [self.get_key(h) for h in hashes]
# Skip keys that already exist in Mooncake
exists = self.store.batch_is_exist(keys)
put_keys = []
put_ptrs = []
put_sizes = []
put_indices = []
success_map = [True] * len(hashes)
for i, status in enumerate(exists):
if status != 1:
put_keys.append(keys[i])
put_ptrs.append(ptrs[i])
put_sizes.append(sizes[i])
put_indices.append(i)
if not put_keys:
return success_map
results = self.store.batch_put_from_multi_buffers(put_keys, put_ptrs, put_sizes)
for i, res in enumerate(results):
success_map[put_indices[i]] = res == 0
return success_map
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,195 @@
import logging
import uuid
import torch
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import MooncakeStore
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
def make_hicache_storage_config(
*,
is_mla_model: bool,
tp_rank: int,
tp_size: int,
) -> HiCacheStorageConfig:
return HiCacheStorageConfig(
tp_rank=tp_rank,
tp_size=tp_size,
pp_rank=0,
pp_size=1,
attn_cp_rank=0,
attn_cp_size=1,
is_mla_model=is_mla_model,
enable_storage_metrics=False,
is_page_first_layout=True,
model_name=None,
)
def generate_batch_query_keys(kv_num: int):
return ["test_" + str(uuid.uuid4()) for _ in range(kv_num)]
def create_mock_host_kv_cache(
buffer_size,
entries_per_page=2,
page_elements=1,
dtype=torch.float32,
):
"""Create a mock HostKVCache-like object for testing."""
buffer = torch.randn(buffer_size, dtype=dtype)
class MockHostKVCache:
def __init__(self, buffer, entries_per_page, page_elements):
self.kv_buffer = buffer
self.layout = "page_first"
self.page_size = 1 # Simple page size for testing
self.entries_per_page = entries_per_page
self.page_elements = page_elements
def get_page_buffer_meta(self, indices):
"""Mock implementation of get_page_buffer_meta."""
ptr_list = []
element_size_list = []
for idx in indices:
page_idx = int(idx)
page_offset = page_idx * self.entries_per_page * self.page_elements
for entry_idx in range(self.entries_per_page):
offset = page_offset + entry_idx * self.page_elements
ptr_list.append(self.kv_buffer[offset:].data_ptr())
element_size_list.append(
self.page_elements * self.kv_buffer.element_size()
)
return ptr_list, element_size_list
def get_ksize_per_token(self):
return (
self.entries_per_page
* self.page_elements
* self.kv_buffer.element_size()
)
return MockHostKVCache(buffer, entries_per_page, page_elements), buffer
def test_single_operation():
"""Test the set API with a single key-value pair."""
print("=" * 100)
print("Testing single operation")
buffer_size = 1024 * 1024 * 16 # 16MB
value_elements = 1024
store = MooncakeStore(
make_hicache_storage_config(is_mla_model=False, tp_rank=0, tp_size=1)
)
mock_host_kv_cache, buffer = create_mock_host_kv_cache(
buffer_size,
entries_per_page=2,
page_elements=value_elements,
)
# Register the memory pool host - this is the proper workflow
store.register_mem_pool_host(mock_host_kv_cache)
value_size = value_elements * buffer.element_size()
key = str(uuid.uuid4())
set_slice = buffer[:value_elements]
get_slice = buffer[value_elements : 2 * value_elements]
set_location = set_slice.data_ptr()
get_location = get_slice.data_ptr()
# Test set operation
result = store.set(key, target_location=set_location, target_sizes=value_size)
assert result is True, f"❌set operation failed for key: {key}"
# Test exists operation
assert store.exists(key), f"❌key {key} should exist after set operation"
# Test get operation
result = store.get(key, target_location=get_location, target_sizes=value_size)
assert result is True, f"❌get operation failed for key: {key}"
# Compare the data using proper tensor indices
assert torch.allclose(
set_slice, get_slice, atol=1e-6
), f"❌get operation failed for key: {key}"
logger.info(f"✅ Single operation passed")
def test_batch_operation(config: HiCacheStorageConfig):
"""Test the batch set/get APIs with multiple key-value pairs."""
print("=" * 100)
print(f"Testing batch operation with config: {config}")
buffer_size = 1024 * 1024 * 16 # 16MB
value_elements = 256
kv_num = 13
entries_per_page = 1 if config.is_mla_model else 2
store = MooncakeStore(config)
mock_host_kv_cache, buffer = create_mock_host_kv_cache(
buffer_size,
entries_per_page=entries_per_page,
page_elements=value_elements,
)
store.register_mem_pool_host(mock_host_kv_cache)
keys = generate_batch_query_keys(kv_num)
set_slices = [
buffer[i * value_elements : (i + 1) * value_elements]
for i in range(kv_num * entries_per_page)
]
set_indices = torch.arange(kv_num)
# Test batch set operation
result = store.batch_set_v1(keys, set_indices)
assert all(result), "batch set operation failed"
# Test batch exists operation
assert (
store.batch_exists(keys) == kv_num
), "keys should exist after batch set operation"
# Test batch get operation
get_slices = [
buffer[
(kv_num * entries_per_page + i)
* value_elements : (kv_num * entries_per_page + i + 1)
* value_elements
]
for i in range(kv_num * entries_per_page)
]
get_indices = torch.arange(kv_num, 2 * kv_num)
result = store.batch_get_v1(keys, get_indices)
assert all(result), "❌batch get operation failed"
for i in range(kv_num * entries_per_page):
assert torch.allclose(
set_slices[i], get_slices[i], atol=1e-6
), f"❌batch get operation failed for key: {keys[i // entries_per_page]}"
logger.info(f"✅ Batch operation passed")
if __name__ == "__main__":
test_single_operation()
test_batch_operation(
make_hicache_storage_config(is_mla_model=False, tp_rank=0, tp_size=1)
)
test_batch_operation(
make_hicache_storage_config(is_mla_model=True, tp_rank=0, tp_size=1)
)
test_batch_operation(
make_hicache_storage_config(is_mla_model=False, tp_rank=1, tp_size=4)
)
test_batch_operation(
make_hicache_storage_config(is_mla_model=True, tp_rank=3, tp_size=8)
)
logger.info(f"✅ All tests passed")
@@ -0,0 +1,614 @@
# NIXL Integration for HiCache
This directory contains the **NIXL (NVIDIA Inference Xfer Library)** integration for **HiCache**, enabling high-performance storage across multiple backends.
NIXL provides a unified API for accessing various storage plugins, including but not limited to:
- POSIX for file based operations, including AIO / io_uring / POSIX AIO.
- **Deepseek's 3FS APIs** for high-throughput file operations
- **GPU Direct Storage (GDS)** for direct data movement between storage and GPU memory, bypassing CPU memory copies
- **Amazon S3-compatible object storage** for key-value access patterns
NIXL also supports additional backends such as **AZURE_BLOB**, **GUSLI**, and **UCX**. Additional backend integrations are planned for future releases.
## NIXL Resources
- **Project Repository**: [NIXL on GitHub](https://github.com/ai-dynamo/nixl)
- **Documentation**: [NIXL Documentation](https://github.com/ai-dynamo/nixl/tree/main/docs)
## Overview
The NIXL integration consists of these main files:
- **`hicache_nixl.py`** - Main HiCache storage connector using NIXL
- **`nixl_utils.py`** - Utility classes for backend selection, registration, and file management
- **`nixl_cleaner.py`** - Background FILE-backend disk cleaner
At runtime, HiCache uses NIXL as a transfer layer between host memory and either:
- **FILE-backed storage plugins** such as 3FS / POSIX / GDS / GDS_MT
- **OBJ-backed storage plugins** such as S3-compatible object stores
The connector supports both the legacy tensor-oriented API (`get` / `set`) and the newer page-oriented API (`batch_get_v1` / `batch_set_v1`) used by modern HiCache backends.
## Components
### HiCacheNixl
The main storage connector that provides:
- Single and batch tensor set/get operations
- Automatic backend selection (3FS > POSIX > GDS_MT > GDS > OBJ)
- High-performance file-based (or) object based storage access using NIXL
- Automatic zero-copy enablement when HiCache host memory layout is `page_first` or `page_first_direct`
- MLA-aware storage naming and backend-local MLA backup skipping on non-zero TP ranks
- Runtime diagnostics for mem-pool type, MLA mode, TP rank, and backup-skip state
### NixlUtils
Consolidated utility classes:
- **NixlBackendSelection** - Handles backend selection and creation
- **NixlBackendConfig** - Handles backend configuration
- **NixlFileManager** - Handles file system operations
### NixlRegistry (`nixl_registry.py`)
Owns the `(agent, mem_type, file_manager)` triple and exposes `host(...)` and `storage(...)` context managers that register on entry, yield the NIXL `xfer_descs`, and deregister + close fds on exit. Internally composes two single-resource primitives (`_open_files` and `_registered`) so leak-freeness is verifiable per primitive.
The current implementation performs per-transfer registration for file / object targets and explicitly closes FILE descriptors after registration / transfer setup to avoid descriptor leaks.
### L3 Cleaner (`nixl_cleaner.py`)
For FILE-backed plugins, TP rank 0 starts a best-effort background cleaner that scans the bucketed storage directories and deletes the oldest logical cache-key groups when disk usage exceeds the configured high watermark. Deleted files are handled by the cache layer as ordinary storage misses and can be recomputed.
Set the top-level `l3_cleaner_enabled` config key to `false` when an external cleaner is responsible for L3 cache eviction.
## Using NIXL as the HiCache Storage Backend
### 1. How Backend Plugin Selection Works
The NIXL backend can support **multiple storage plugins** (e.g., POSIX, GDS, GDS_MT, 3FS, object store, etc).
* Each plugin has its own configuration section in the TOML file.
* The connector accepts configuration in two forms:
* a **fully qualified** form such as `{"plugin": {"posix": {...}, "gds": {...}}}`
* a **flat** form such as `{"use_uring": "true"}`, which applies to the selected plugin
* A plugin is considered **usable** if:
* Its required library is available on the system (POSIX, GDS, GDS_MT are natively supported by NIXL).
* Its configuration is valid.
* It is marked as `active = true` in the configuration file (if applicable).
* Some plugins (e.g., 3FS, GDS) require additional system libraries or hardware support.
* If the config explicitly enables multiple plugins, the connector chooses the **first active plugin** in the config.
* If no plugin is explicitly selected in config, the connector falls back to the environment variable `SGLANG_HICACHE_NIXL_BACKEND_PLUGIN`, and finally to `auto`.
* In `auto` mode, NIXL selects the backend based on **internal priority and availability**.
If a plugin is configured but its dependencies are missing, it will be skipped.
### 2. Setting the Storage Directory (Optional)
For POSIX / GDS / GDS_MT file-based backends, the default storage location is `/tmp/hicache_storage`. However, you can customize where cached data is stored:
```bash
# When specifying multiple storage directories. SGLang routes each cache object to one
# directory with a stable hash.
export SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR=/path/to/storage/dir1,/path/to/storage/dir2,/path/to/storage/dir3
```
These directories are used only for **FILE-backed** plugins. **OBJ-backed** plugins use object keys instead of local files.
### 3. How to Provide Configuration for Backends
There are three ways to specify configurations for the backends: default config, file based config, and command-line (JSON string based) config.
#### 1. Using Default Configuration
To enable HiCache with the NIXL backend, start the SGLang server with:
```bash
python3 -m sglang.launch_server \
--model-path <model> \
--host <ip> \
--port <port> \
--page-size 64 \
--enable-hierarchical-cache \
--hicache-ratio 2 \
--hicache-size 64 \
--hicache-write-policy write_through \
--hicache-storage-backend nixl
```
By default, NIXL will use its internal backend selection logic to choose an available storage plugin (and use default configs for the selected storage plugin).
For object storage backends, make sure the bucket is configured either in `--hicache-storage-backend-extra-config` or via:
```bash
export AWS_DEFAULT_BUCKET=<bucket-name>
```
#### 2. Using a Configuration File (Recommended)
For non-trivial setups with complex configurations, it is recommended to use a **TOML configuration file** to define which backend plugin to use and its configurations, via `--hicache-storage-backend-extra-config`:
Below is an example command (note: detailed configs are defined in the config file):
```bash
python3 -m sglang.launch_server \
--model-path <model> \
--host <ip> \
--port <port> \
--page-size 64 \
--enable-hierarchical-cache \
--hicache-ratio 2 \
--hicache-size 64 \
--hicache-write-policy write_through \
--hicache-storage-backend nixl \
--hicache-storage-backend-extra-config "@config.nixl.toml"
```
> **Important**
>
> * The `@` prefix tells SGLang to load the configuration from a file.
> * The file can be in **TOML format** (other formats, JSON / YAML, are also supported).
> * This is the preferred way to configure NIXL storage backends.
The structure of the config file is described in further details in [Configuration File Spec](#Configuration-File-Specification).
#### 3. Using Command-line JSON String
For debugging or quick testing, you may pass a **JSON-style string** directly via `--hicache-storage-backend-extra-config`.
This requires explicitly specifying the plugin type via an environment variable, and this method can be applicable to **only a few** plugins (e.g., POSIX, GDS, GDS_MT)
The below example shows how to use command-line string to use the POSIX plugin where URING is enabled for async POSIX storage, with O_DIRECT enabled (the default).
```bash
export SGLANG_HICACHE_NIXL_BACKEND_PLUGIN=POSIX
python3 -m sglang.launch_server \
--model-path <model> \
--host <ip> \
--port <port> \
--page-size 64 \
--enable-hierarchical-cache \
--hicache-ratio 2 \
--hicache-size 64 \
--hicache-write-policy write_through \
--hicache-storage-backend nixl \
--hicache-storage-backend-extra-config '{"use_uring": "true"}'
```
To disable O_DIRECT (e.g. for debugging or unsupported filesystems), set the top-level `use_direct_io` key:
```bash
export SGLANG_HICACHE_NIXL_BACKEND_PLUGIN=POSIX
python3 -m sglang.launch_server \
... \
--hicache-storage-backend-extra-config '{"use_direct_io": false, "use_uring": "true"}'
```
⚠️ **Note**:
This method is convenient for testing / experimenting. For production or multi-plugin setups, it is always recommended to use the config file based approach.
Also note that the flat inline config form is interpreted as plugin-specific parameters for the selected plugin.
## Running Unit Tests
### Prerequisites
- NIXL library installed and available (latest main required for supporting object query)
- PyTorch installed
- Python 3.8+
### Unit tests from current directory
From the current directory run:
#### Run all NIXL tests:
```bash
PYTHONPATH=. python -m pytest test_hicache_nixl_storage.py -o asyncio_mode=strict
```
#### Run with verbose output:
```bash
PYTHONPATH=. python -m pytest test_hicache_nixl_storage.py -v -o asyncio_mode=strict
```
Note: The `-v` flag provides more detailed output, showing each test case name and its result.
#### Run a specific test:
```bash
PYTHONPATH=. python -m pytest test_hicache_nixl_storage.py -v -k test_single_set_get -o asyncio_mode=strict
```
Note: The `-o asyncio_mode=strict` flag is added to suppress warnings about asyncio configuration. This is not required for test functionality but provides cleaner output.
## Test Coverage
Tests for this integration, a test suite can be found at `test_hicache_nixl_storage.py` which covers:
### HiCache Integration Tests
- Single tensor set/get operations
- Batch tensor set/get operations
- Mixed single and batch operations
- Data integrity for various tensor types
### File Management Tests
- Basic file operations
- NIXL tuple creation
- Error handling in file operations
### Registration and MLA / Query Tests
- Tensor registration with memory type detection
- File registration using file paths
- MLA backup-skip behavior for `batch_set_v1`
- Zero-copy `batch_exists()` accounting for MLA and MHA
## Expected Output
When tests run successfully, you should see:
- NIXL agent initialization messages
- Backend selection messages (e.g., "Backend POSIX was instantiated")
- Test results with "ok" for passed tests
- Summary showing "Ran X tests in Y seconds" and "OK"
## Troubleshooting
### Import Errors
If you encounter `ModuleNotFoundError`, ensure:
- You're running from the correct directory
- `PYTHONPATH` is set correctly
- NIXL library is properly installed
### NIXL Errors
If NIXL operations fail:
- Check that NIXL is properly installed
- Verify that required plugins are available
- Ensure file permissions are correct for test directories
- For OBJ plugins, verify `bucket` or `AWS_DEFAULT_BUCKET` is set
- Check the NIXL diagnostic log emitted when the mem pool is registered; it includes:
- `mem_pool_device_type`
- `is_mla_model`
- `tp_rank`
- `backup_skip`
### MLA Write Behavior
For MLA models, the NIXL backend now mirrors HF3FS's backend-local protection:
- TP rank 0 performs the actual storage write
- non-zero TP ranks skip backup writes locally in `batch_set` / `batch_set_v1`
- MLA storage names omit TP rank so all ranks refer to the same logical storage object or file
## File Structure
```text
python/sglang/srt/mem_cache/storage/nixl/
├── hicache_nixl.py # Main HiCache storage connector
├── nixl_cleaner.py # Background FILE-backend disk cleaner
├── nixl_utils.py # NIXL utility classes
├── test_hicache_nixl_storage.py # Unit tests
├── nixl.config.toml.sample # Example configuration
└── README.md # This file
```
## Dependencies
- **NIXL**: NVIDIA Inference Xfer Library (version 0.4 or later)
- Required plugins: POSIX (minimum), 3FS/GDS (optional for better performance)
- See [NIXL Installation Guide](https://github.com/ai-dynamo/nixl/blob/main/README.md#installation)
- **PyTorch**: For tensor operations (version 1.8 or later)
- **Python 3.8+**: For type hints and modern features
## Supported Features
### Memory Types
- **Tensor side**: multi-dimensional tensors of all numeric types (int32, int64, float32, float64) are supported.
- Tensors can be on CPU or GPU (as long as a GPU capable backend such as GDS_MT is available).
- Currently each tensor is mapped to a file or key, but it can be extended to support multiple keys per file or key.
- The page-oriented `*_v1` path also supports zero-copy transfers using `(address, size)` metadata from the host memory pool.
- **Storage side**: file and object are supported through their relevant backends (e.g., 3FS or OBJ).
### HiCache / NIXL Data Model
- **FILE backends** use local file paths under `SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR`. When multiple comma-separated directories are configured, each logical cache key is routed to one base directory with a stable hash and stored as `base_dir/<bucket>/<key>`.
- **OBJ backends** use object keys directly
- **MHA naming** includes TP rank and TP size, so each rank stores its own KV data
- **MLA naming** omits TP rank, so all ranks refer to one shared logical KV object / file
- In zero-copy mode:
- **MHA** expands each logical page into `_k` and `_v` entries
- **MLA** expands each logical page into a single `_k` entry because MLA stores one interleaved KV representation
- The L3 cleaner groups physical files by the logical base key after removing TP-rank and zero-copy `_k` / `_v` suffixes. This keeps MHA, MLA, and DSA file cleanup aligned with the names emitted by `HiCacheNixl`.
### Zero-Copy Behavior
- Zero-copy is enabled automatically when the HiCache host layout is `page_first` or `page_first_direct`
- The connector uses `mem_pool_host.get_page_buffer_meta(...)` to obtain `(address, size)` metadata
- `batch_exists()` uses the same logical key expansion rules as `batch_get_v1()` / `batch_set_v1()`
- Non-zero MLA TP ranks skip `batch_set` / `batch_set_v1()` locally as a backend-side fallback guard
### Backend Priority
The NIXL backend selection follows this priority order:
1. **3FS** - Highest performance (if available)
- Best for high-throughput file operations using Deepseek 3FS APIs
2. **POSIX** - Standard file I/O (fallback)
- Universal compatibility
- Good for development and testing - Leverages both libaio/liburing
3. **GDS_MT** - Multi-threaded GDS (if available)
- Optimized for concurrent operations
- Supports GPU Direct storage with multiple light weight threads
4. **GDS** - GPU Direct Storage (if available)
- Direct GPU-storage data path
- Best for filesystems benefiting from batch operations and smaller IOs.
5. **OBJ** - Amazon S3 based Object Storage
- Key-value based storage
The system automatically selects the best available backend, with POSIX as the default fallback.
## Configuration File Specification
This section defines the structure, supported sections, configuration keys, data types, defaults, and semantics for the NIXL HiCache backend configuration file (`config.nixl.toml`).
The configuration file is written in **TOML** and consists of multiple **plugin-specific sections** under the `plugin.*` namespace. Each section configures one storage backend plugin. Only one plugin should be enabled via setting `active = true` in the corresponding plugin-specific section.
An example of the configuration is provided in [`nixl.config.toml.sample`](./nixl.config.toml.sample).
### 1. General Structure
```toml
[plugin.<backend_name>]
<key> = <value>
```
* `<backend_name>` identifies the storage backend plugin.
* Each plugin is configured independently.
* Plugins are selected at runtime based on:
* Availability of required libraries/hardware
* Plugin configuration validity
* Internal backend priority rules
* Unless otherwise stated, all configuration keys are **optional** and have sensible defaults.
For object storage, `bucket` may also be omitted from the config if `AWS_DEFAULT_BUCKET` is already defined in the environment.
### 1a. Top-Level Configuration Keys
The following keys are placed at the **top level** of the config file (not inside any `[plugin.*]` section) and apply globally to the NIXL backend:
| Key | Type | Default | Description |
| ---------------- | ------- | -------- | ----------- |
| `use_direct_io` | boolean | `true` | Open cache files with `O_DIRECT` to bypass the OS page cache. Reduces memory pressure and improves NVMe throughput. Falls back to buffered I/O with a warning if `O_DIRECT` is unavailable on the current OS. Can also be overridden via the `SGLANG_HICACHE_NIXL_USE_DIRECT_IO` environment variable. |
| `l3_cleaner_enabled` | boolean | `true` | Enable the built-in background cleaner for FILE-backed L3 storage. Set to `false` when using an external cleaner. |
| `l3_cleaner_high_watermark` | float | `80.0` | Start cleanup when the built-in cleaner is enabled and the filesystem containing a configured storage directory reaches this disk-usage percentage. |
| `l3_cleaner_low_watermark` | float | `70.0` | Stop cleanup after hot filesystems drop below this disk-usage percentage. Must be lower than `l3_cleaner_high_watermark`. |
**Page-alignment and `O_DIRECT`**
When `use_direct_io = true` with any file-based backend (POSIX, GDS, GDS_MT, 3FS), the kernel requires every I/O buffer pointer to be OS-page-aligned (4 KiB). SGLang handles this automatically:
* **Zero-copy mode** (`page_first` / `page_first_direct` layout): the host memory pool is always mmap-backed and therefore page-aligned. If the per-page stride is also a multiple of 4 KiB, zero-copy transfers are used as-is.
* **Copy mode** (all other layouts, or if stride alignment cannot be satisfied): SGLang pre-allocates page-aligned bounce buffers via `mmap` and falls back to copy mode, logging a warning. No user action is required -- this is fully automatic.
To disable `O_DIRECT` (e.g. for debugging or when the filesystem does not support it):
```toml
use_direct_io = false
[plugin.posix]
use_uring = "true"
active = true
```
or via environment variable: `SGLANG_HICACHE_NIXL_USE_DIRECT_IO=0`.
To tune FILE-backend cleanup watermarks:
```toml
l3_cleaner_enabled = true
l3_cleaner_high_watermark = 85.0
l3_cleaner_low_watermark = 75.0
[plugin.posix]
use_uring = "true"
active = true
```
To use an external cleaner instead of the built-in cleaner:
```toml
l3_cleaner_enabled = false
[plugin.posix]
use_uring = "true"
active = true
```
### 2. POSIX File System Backend (`plugin.posix`)
#### Section
```toml
[plugin.posix]
```
#### Description
Configures the POSIX file-system-based backend.
This backend supports multiple asynchronous I/O mechanisms and automatically selects the most performant option supported by the system.
**Backend priority (highest to lowest):**
1. Linux AIO
2. `io_uring`
3. POSIX AIO
#### Configuration Keys
| Key | Type | Default | Description |
| --------------- | ------- | --------- | -------------------------------------------------------------------------------------------------------- |
| `use_uring` | string | `"false"` | Enables Linux `io_uring` for asynchronous I/O when set to `"true"`. Recommended on modern Linux kernels. |
| `use_posix_aio` | string | `"false"` | Enables POSIX AIO as an alternative async I/O mechanism. |
| `use_aio` | string | `"false"` | Enables generic Linux AIO. |
| `active` | boolean | N/A | Controls whether this plugin is eligible for backend selection. |
**Notes**
* Boolean-like options use **string values** (`"true"` / `"false"`) for compatibility.
* **Only one backend** (i.e., only one of `use_uring`, `use_aio`, `use_posix_aio`) should be included in the config.
### 3. NVIDIA GPUDirect Storage Backend (`plugin.gds`)
#### Section
```toml
[plugin.gds]
```
#### Description
Configures NVIDIA GPUDirect Storage (GDS) backend.
This backend enables direct data transfers between storage and GPU memory.
**Requirements**
* NVIDIA GPU with GDS support
* Compatible NVIDIA driver and CUDA runtime
* Supported filesystem
#### Configuration Keys
| Key | Type | Default | Description |
| ------------------ | ------- | ------------------ | ------------------------------------------------------ |
| `batch_pool_size` | integer | `128` | Number of I/O requests maintained in the request pool. |
| `batch_limit` | integer | `128` | Maximum number of requests issued in a single batch. |
| `max_request_size` | integer | `16777216` (16 MB) | Maximum size (in bytes) of a single I/O request. |
| `active` | boolean | N/A | Controls whether this plugin is eligible for backend selection.|
### 4. Multi-Threaded GDS Backend (`plugin.gds_mt`)
#### Section
```toml
[plugin.gds_mt]
```
#### Description
Configures the multi-threaded variant of the NVIDIA GDS backend, allowing parallel request processing using multiple CPU threads.
#### Configuration Keys
| Key | Type | Default | Description |
| -------------- | ------- | ------- | ----------------------------------------------------- |
| `thread_count` | integer | `4` | Number of worker threads used to submit GDS requests. |
| `active` | boolean | N/A | Controls whether this plugin is eligible for backend selection. |
### 5. 3FS Backend (`plugin.3fs`)
#### Section
```toml
[plugin.3fs]
```
#### Description
Configures the 3FS (third-party filesystem) backend.
**Requirements**
* 3FS client library installed
* Filesystem mounted and accessible on the host
#### Configuration Keys
| Key | Type | Default | Description |
| ------------- | ------- | -------- | ---------------------------------- |
| `mount_point` | string | *none* | Mount point of the 3FS filesystem. |
| `mem_config` | string | `"dram"` | Memory configuration mode. |
| `iopool_size` | integer | `64` | Size of the I/O pool. |
| `active` | boolean | N/A | Controls whether this plugin is eligible for backend selection. |
##### `mem_config` Valid Values
| Value | Description |
| --------- | --------------------------------------------------- |
| `dram` | Use DRAM for buffering |
| `dram_zc` | Use DRAM with zero-copy support |
| `auto` | Automatically select based on platform capabilities |
##### `iopool_size` Constraints
* Valid range: **[2⁶, 2²⁰]**
* Values outside this range may cause initialization failure.
### 6. Object Storage Backend (`plugin.obj`)
#### Section
```toml
[plugin.obj]
```
#### Description
Configures an object storage backend compatible with S3 APIs (e.g., AWS S3, MinIO, Ceph).
#### Configuration Keys
| Key | Type | Default | Description |
| ------------------------ | ------- | ------------ | ---------------------------------------------- |
| `num_threads` | integer | `4` | Number of client worker threads. |
| `endpoint_override` | string | `""` | Custom endpoint URL (for non-AWS S3 services). |
| `scheme` | string | `"http"` | Connection scheme (`http` or `https`). |
| `region` | string | `""` | Cloud region (if applicable). |
| `req_checksum` | string | `"required"` | Request checksum behavior. |
| `ca_bundle` | string | `""` | Path to a custom CA bundle. |
| `access_key` | string | `""` | Access key credential. |
| `secrete_key` | string | `""` | Secret key credential. |
| `session_token` | string | `""` | Session token (optional). |
| `use_virtual_addressing` | string | `"true"` | Enables virtual-hosted-style addressing. |
| `bucket` | string | `""` | Default bucket name. |
| `active` | boolean | N/A | Controls whether this plugin is eligible for backend selection. |
##### `req_checksum` Valid Values
| Value | Description |
| ----------- | ---------------------------------------------- |
| `required` | Always include a checksum |
| `supported` | Include checksum when supported by the backend |
### 7. Notes and Best Practices
* All plugin sections are optional.
* Multiple plugins may be configured in a single file. However, it is recommended that **only one plugin** is configured `active = true`.
* Plugins whose dependencies are unavailable will be skipped.
* Use a TOML configuration file instead of inline JSON for:
* Multi-plugin setups
* Production deployments
* Clear validation and maintainability
## Note
This is v0 of the NIXL connector. The current implementation favors correctness and compatibility with the existing HiCache API:
- file / object targets are registered per transfer
- FILE descriptors are explicitly cleaned up after registration / transfer setup
- MLA uses shared storage naming and backend-local write skipping on non-zero TP ranks
- zero-copy is driven by HiCache host-memory layout rather than a separate NIXL flag
Future versions will focus on further performance optimizations such as memory pre-registration (pre-allocating and registering memory buffers to reduce registration overhead during transfers) and block merging (combining related blocks as offsets within the same file to reduce file operations and improve throughput). These optimizations require changes at a higher layer, as the current HiCache API doesn't expose information like block relationships or hash patterns that would enable these optimizations.
@@ -0,0 +1,622 @@
import logging
import os
import time
import uuid
from typing import Any, List, Optional
import torch
from sglang.srt.environ import envs
from sglang.srt.mem_cache.hicache_storage import (
STORAGE_BATCH_SIZE,
HiCacheStorage,
HiCacheStorageConfig,
HiCacheStorageExtraInfo,
)
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.nixl.nixl_cleaner import HiCacheL3Cleaner
from .nixl_registry import NixlRegistry
from .nixl_utils import NixlBackendConfig, NixlBackendSelection, NixlFileManager
try:
from nixl._api import nixl_agent, nixl_agent_config, nixlBind
except ImportError as e:
raise ImportError(
"Please install NIXL by following the instructions at "
"https://github.com/ai-dynamo/nixl/blob/main/README.md "
"to use HiCacheNixl storage backend."
) from e
logger = logging.getLogger(__name__)
def _parse_storage_dirs(raw: Optional[str]) -> List[str]:
"""Split NIXL FILE storage directory config into ordered unique paths."""
if not raw:
return []
candidates = [path.strip() for path in raw.split(",")]
candidates = [path for path in candidates if path]
seen: dict[str, str] = {}
ordered: List[str] = []
for path in candidates:
real_path = os.path.realpath(path)
if real_path in seen:
raise ValueError(
"SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR contains duplicate "
f"path {path!r} (same mount as {seen[real_path]!r})."
)
seen[real_path] = path
ordered.append(path)
return ordered
class HiCacheNixl(HiCacheStorage):
"""HiCacheNixl provides high-performance storage using NIXL plugins."""
def __init__(
self,
storage_config: HiCacheStorageConfig,
file_path: str = "/tmp/hicache_storage",
):
"""Initialize NIXL storage connector."""
# create nixlconfig from the --hicache-storage-backend-extra-config
nixlconfig = NixlBackendConfig(storage_config.extra_config)
# select the NIXL backend plugin from extra_config or environment variable
plugin = nixlconfig.get_specified_plugin()
use_direct_io = nixlconfig.get_use_direct_io()
# Might be better to be unified across HiCache backends and moved to HiCacheController
storage_dirs = _parse_storage_dirs(
envs.SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR.get() or file_path
)
self.file_manager = (
NixlFileManager(storage_dirs, use_direct_io=use_direct_io)
if plugin not in NixlBackendSelection.OBJ_PLUGINS
else None
)
tp_rank, tp_size, model_name = (
storage_config.tp_rank,
storage_config.tp_size,
storage_config.model_name,
)
self.is_mla_model = storage_config.is_mla_model
self.is_zero_copy = False
self.storage_config = storage_config
self.backup_skip = self.is_mla_model and storage_config.tp_rank != 0
model_name = "-".join(model_name.split("/")) if model_name else ""
if self.is_mla_model:
self.config_suffix = f"_{model_name}"
else:
self.config_suffix = f"_{model_name}_{tp_rank}_{tp_size}"
sync_mode = getattr(
nixlBind, "NIXL_THREAD_SYNC_RW", nixlBind.NIXL_THREAD_SYNC_STRICT
)
agent_config = nixl_agent_config(backends=[])
self.agent_name = f"hicache_nixl_{str(uuid.uuid4())}"
self.agent = nixl_agent(self.agent_name, agent_config)
bind_cfg = nixlBind.nixlAgentConfig()
bind_cfg.useProgThread = agent_config.enable_pthread
bind_cfg.useListenThread = agent_config.enable_listen
bind_cfg.listenPort = agent_config.port
bind_cfg.syncMode = sync_mode
bind_cfg.pthrDelay = 0
bind_cfg.lthrDelay = 100000
bind_cfg.captureTelemetry = agent_config.capture_telemetry
self.agent.agent = nixlBind.nixlAgent(self.agent_name, bind_cfg)
self.agent.plugin_list = self.agent.agent.getAvailPlugins()
self.backend_selector = NixlBackendSelection(plugin, nixlconfig)
if not self.backend_selector.create_backend(self.agent):
raise RuntimeError("Failed to create NIXL backend")
self.registry = NixlRegistry(
self.agent,
self.backend_selector.mem_type,
self.file_manager,
)
# O_DIRECT requires OS-page-aligned I/O buffers on all file-based backends
# (POSIX, GDS, GDS_MT, 3FS). OBJ backends never open files so they are exempt
# (file_manager is None for OBJ).
self.needs_page_alignment = use_direct_io and self.file_manager is not None
if self.needs_page_alignment:
logger.info(
"HiCacheNixl: O_DIRECT is active with a file-based backend (%s). "
"Page-aligned host buffers are required (needs_page_alignment=True).",
self.backend_selector.backend_name,
)
# Pre-registered host regions (set by register_mem_pool_host):
# zero-copy: one registration covering mem_pool_host.kv_buffer
# non-zero-copy: two registrations, one bounce buffer per direction
# (set/get) so the two storage threads never share slots.
self._host_regs: List[Any] = []
self._bounce_set: Optional[torch.Tensor] = None
self._bounce_get: Optional[torch.Tensor] = None
self._bounce_page_bytes: Optional[int] = None
cleanup_dirs = (
self.file_manager.iter_all_base_dirs()
if self.file_manager is not None
else []
)
cleaner_config = nixlconfig.get_l3_cleaner_config()
self._l3_cleaner: Optional[HiCacheL3Cleaner] = (
HiCacheL3Cleaner(
cleanup_dirs,
tp_rank,
high_watermark=cleaner_config["high_watermark"],
low_watermark=cleaner_config["low_watermark"],
)
if (
cleanup_dirs
and self.file_manager is not None
and cleaner_config["enabled"]
)
else None
)
if self._l3_cleaner is not None:
self._l3_cleaner.start()
def _get_suffixed_key(self, key: str) -> str:
return key + self.config_suffix
def _create_query_tuple(self, key: str) -> tuple:
"""Build the NIXL query_memory tuple for a single key."""
if self.backend_selector.mem_type == "FILE":
return (0, 0, 0, self.file_manager.get_file_path(key))
return (0, 0, 0, key)
def _xfer_and_wait(
self,
host_descs: Any,
storage_descs: Any,
direction: str,
) -> bool:
"""Initialize and poll a NIXL transfer to completion."""
try:
xfer_req = self.agent.initialize_xfer(
direction, host_descs, storage_descs, self.agent_name
)
except Exception as e:
logger.error(f"Failed to create transfer request: {e}")
return False
try:
state = self.agent.transfer(xfer_req)
while state != "DONE":
state = self.agent.check_xfer_state(xfer_req)
if state == "ERR":
logger.error("Transfer failed")
return False
# Best would be to have a better notification mechanism from NIXL,
# but we only have polling for now.
time.sleep(0.0001)
return True
except Exception as e:
logger.error(f"Failed to execute transfer: {e}")
import traceback
logger.error(f"Traceback: {traceback.format_exc()}")
return False
finally:
self.agent.release_xfer_handle(xfer_req)
def _xfer_pre_registered(
self,
host_buffers: List[tuple],
keys: List[str],
direction: str,
) -> bool:
"""Run a transfer where the host side is already pre-registered.
``host_buffers`` is a list of ``(addr, size)`` tuples within the
pre-registered host region (kv_buffer for zero-copy, bounce buffer
otherwise). Only the storage side is registered per transfer.
"""
if len(host_buffers) != len(keys):
logger.error("Mismatch between number of host buffers and keys")
return False
host_descs = self.agent.get_xfer_descs(
[(addr, size, 0) for (addr, size) in host_buffers], "DRAM"
)
if host_descs is None:
logger.error("Failed to build host xfer descs")
return False
with self.registry.storage(host_buffers, keys, direction) as storage_descs:
if storage_descs is None:
return False
return self._xfer_and_wait(host_descs, storage_descs, direction)
def get(
self,
key: str,
target_location: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> torch.Tensor | None:
raise NotImplementedError("deprecated; use batch_get_v1")
def batch_get(
self,
keys: List[str],
target_locations: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> List[torch.Tensor | None]:
raise NotImplementedError("deprecated; use batch_get_v1")
def set(
self,
key: str,
value: Optional[Any] = None,
target_location: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> bool:
raise NotImplementedError("deprecated; use batch_set_v1")
def batch_set(
self,
keys: List[str],
values: Optional[Any] = None,
target_locations: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> bool:
raise NotImplementedError("deprecated; use batch_set_v1")
def register_mem_pool_host(self, mem_pool_host: HostKVCache):
super().register_mem_pool_host(mem_pool_host)
# enable zero-copy automatically if mem layout is page_first or page_first_direct
self.is_zero_copy = self.mem_pool_host.layout in [
"page_first",
"page_first_direct",
]
if self.needs_page_alignment and self.is_zero_copy:
# Check that the kv_buffer base AND per-page strides are multiples of
# the OS page size so every pointer passed to NIXL (base + p * stride)
# is page-aligned. The base is whatever torch.empty() happened to give
# us -- it is not guaranteed to be page-aligned. Fall back to copy mode
# if either condition fails.
# 4096: O_DIRECT alignment is FS-dependent (some allow 512 B); 4 KiB
# is the safe lower bound all known FSes accept, and real page-sizes meet it.
if not self.mem_pool_host.is_stride_page_aligned(4096):
logger.warning(
"HiCacheNixl: O_DIRECT is active but the host kv_buffer is "
"not OS-page-aligned (base or per-page stride). Falling back "
"to copy mode for this pool."
)
self.is_zero_copy = False
if self.is_zero_copy:
kv = mem_pool_host.kv_buffer
self._pre_register_host(
kv.data_ptr(), kv.numel() * kv.element_size(), "kv_buffer"
)
else:
# One bounce buffer per direction so set/get run lock-free across
# the prefetch and backup threads. Sized from get_dummy_flat_data_page()
# so each slot matches what the v1 path would otherwise allocate.
sample = mem_pool_host.get_dummy_flat_data_page()
page_numel = sample.numel()
self._bounce_page_bytes = page_numel * sample.element_size()
del sample
pin_memory = bool(getattr(mem_pool_host, "pin_memory", False))
self._bounce_set = self._alloc_registered(
page_numel, mem_pool_host.dtype, pin_memory, "bounce_set"
)
self._bounce_get = self._alloc_registered(
page_numel, mem_pool_host.dtype, pin_memory, "bounce_get"
)
logger.info(
f"HiCacheNixl: pre-registered host regions for "
f"layout={mem_pool_host.layout} zero_copy={self.is_zero_copy}"
)
def _alloc_registered(
self,
page_numel: int,
dtype: torch.dtype,
pin_memory: bool,
kind: str,
) -> torch.Tensor:
"""Allocate a ``(STORAGE_BATCH_SIZE, page_numel)`` bounce buffer and
pre-register it as a DRAM region with NIXL. Uses alloc_mmap so the
buffer is page-aligned -- required when O_DIRECT is on for any
file-based backend (POSIX/GDS/GDS_MT/3FS). pin_memory is currently
unused (alloc_mmap does not support it)."""
buf = alloc_mmap((STORAGE_BATCH_SIZE, page_numel), dtype)
self._pre_register_host(buf.data_ptr(), buf.numel() * buf.element_size(), kind)
return buf
def _pre_register_host(self, base_addr: int, total_size: int, kind: str) -> None:
"""Register a single DRAM region up-front and remember the handle."""
reg_descs = self.agent.get_reg_descs([(base_addr, total_size, 0, "")], "DRAM")
if reg_descs is None:
raise RuntimeError(f"Failed to build reg descs for host {kind}")
try:
self._host_regs.append(self.agent.register_memory(reg_descs))
except Exception as e:
raise RuntimeError(f"Failed to pre-register host {kind} with NIXL") from e
def clear(self) -> None:
if self.file_manager is None:
return
self.file_manager.clear()
def close(self):
if self._l3_cleaner is not None:
self._l3_cleaner.stop()
self._l3_cleaner = None
while self._host_regs:
reg = self._host_regs.pop()
try:
self.agent.deregister_memory(reg)
except Exception as e:
logger.debug("deregister of pre-registered host region failed: %s", e)
self._bounce_set = None
self._bounce_get = None
self._bounce_page_bytes = None
def __del__(self):
try:
self.close()
except Exception:
pass
def exists(self, key: str) -> bool:
results = self.batch_exists([key])
return results > 0
def batch_exists(
self,
keys: List[str],
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> int:
if self.is_zero_copy:
key_list = self._get_key_list_from_meta(keys)
key_denominator = (
1 if self.is_mla_model else 2
) # MLA: 1 key per page (_k only), non-MLA: 2 NIXL keys per page (_k + _v)
else:
key_list = [self._get_suffixed_key(key) for key in keys]
key_denominator = 1
tuples = [self._create_query_tuple(key) for key in key_list]
query_res = self.agent.query_memory(
tuples,
self.backend_selector.backend_name,
mem_type=self.backend_selector.mem_type,
)
for i in range(len(query_res)):
if query_res[i] is None:
return i // key_denominator
return len(query_res) // key_denominator
def _get_key_list_from_meta(self, keys: List[str]) -> List[str]:
# Each key maps to a `_k` entry, plus a `_v` entry on non-MLA models
# (MLA stores k/v interleaved in a single buffer).
key_list = []
for key in keys:
suffixed_key = self._get_suffixed_key(key)
key_list.append(f"{suffixed_key}_k")
if not self.is_mla_model:
key_list.append(f"{suffixed_key}_v")
return key_list
def _get_location_and_size_list_from_meta(
self, keys: List[str], host_indices: torch.Tensor
):
# zero copy: mem_pool_host.get_data_page() does not work due to non-contiguous tensors, causing issues for NIXL transfer
ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(
host_indices
)
key_list = self._get_key_list_from_meta(keys)
if len(key_list) != len(ptr_list):
logger.error(
f"HiCacheNixl: mismatch between number of keys and number of buffer meta entries, keys: {len(keys)}, key_list: {len(key_list)}, buffer meta entries: {len(ptr_list)}"
)
return [], [], []
return key_list, ptr_list, element_size_list
def _bounce_slot_buffers(self, buf: torch.Tensor, page_num: int) -> List[tuple]:
"""Return ``page_num`` ``(addr, size)`` tuples pointing at the first
``page_num`` slots of ``buf``.
"""
base = buf.data_ptr()
return [
(base + i * self._bounce_page_bytes, self._bounce_page_bytes)
for i in range(page_num)
]
def _batch_preprocess(self, keys: List[str], host_indices: torch.Tensor, op: str):
"""Build (key_list, host_buffers) for the v1 path.
For zero-copy: ``host_buffers`` are ``(addr, size)`` tuples inside the
pre-registered ``kv_buffer``.
For non-zero-copy: ``host_buffers`` are slots of the direction-specific
pre-registered bounce buffer (``_bounce_set`` for set, ``_bounce_get``
for get); for ``op == "set"`` we copy the host pages into those slots
here so the subsequent transfer reads from the bounce buffer.
Returns ``([], [])`` on validation failure.
"""
page_size = self.mem_pool_host.page_size
page_num = len(host_indices) // page_size
if len(keys) == 0 or len(keys) != page_num:
logger.warning(
f"HiCacheNixl: empty keys or mismatch in keys and host_indices lengths. keys: {len(keys)}, host_indices: {len(host_indices)}, page_size: {page_size}"
)
return [], []
if self.is_zero_copy:
key_list, ptr_list, size_list = self._get_location_and_size_list_from_meta(
keys, host_indices
)
host_buffers = list(zip(ptr_list, size_list))
return key_list, host_buffers
if page_num > STORAGE_BATCH_SIZE:
logger.error(
f"HiCacheNixl: batch size {page_num} exceeds bounce buffer capacity {STORAGE_BATCH_SIZE}"
)
return [], []
bounce = self._bounce_set if op == "set" else self._bounce_get
if op == "set":
for i in range(page_num):
src = self.mem_pool_host.get_data_page(
host_indices[i * page_size], flat=True
)
bounce[i].copy_(src)
host_buffers = self._bounce_slot_buffers(bounce, page_num)
key_list = [self._get_suffixed_key(key) for key in keys]
return key_list, host_buffers
def _batch_xfer(
self,
keys: List[str],
key_strs: List[str],
host_buffers: List[tuple],
direction: str,
) -> List[bool]:
"""Run a batch READ or WRITE for the v1 path against the pre-registered
host region (no per-transfer host registration).
"""
if not key_strs or not host_buffers:
return [False] * len(keys)
if len(key_strs) != len(host_buffers):
logger.error("Mismatch between number of key_strs and host_buffers")
return [False] * len(keys)
if self.backend_selector.mem_type == "FILE":
file_paths = [self.file_manager.get_file_path(key) for key in key_strs]
success = self._xfer_pre_registered(host_buffers, file_paths, direction)
else: # mem_type == "OBJ"
success = self._xfer_pre_registered(host_buffers, key_strs, direction)
# READ results are consumed by _batch_get_postprocess, which pairs
# entries 2*i / 2*i+1 for non-MLA zero-copy: it needs one bool per
# key_str (i.e. per `_k`/`_v` buffer). WRITE results map 1:1 to
# pages, i.e. to `keys`.
result_len = len(key_strs) if direction == "READ" else len(keys)
return [success] * result_len
def _batch_get_postprocess(
self,
host_indices: torch.Tensor,
results: List[bool],
) -> List[bool]:
page_size = self.mem_pool_host.page_size
page_num = len(host_indices) // page_size
if self.is_zero_copy:
# zero copy: update final results based on the boolean results from NIXL transfer
if self.is_mla_model:
return results
return [(results[2 * i] and results[2 * i + 1]) for i in range(page_num)]
# non zero copy: copy data from the get-side bounce buffer to mem_pool_host
for i in range(page_num):
if not results[i]:
break
self.mem_pool_host.set_from_flat_data_page(
host_indices[i * page_size], self._bounce_get[i]
)
return results
def _log_xfer_stats(
self,
op_name: str,
num_keys: int,
host_indices: torch.Tensor,
buffer_sizes: List[int],
elapsed_ms: float,
) -> None:
total_bytes = sum(s for s in buffer_sizes if s is not None)
bw = total_bytes / (elapsed_ms / 1000) / (1024 * 1024) if elapsed_ms else 0.0
logger.debug(
f"HiCacheNixl {op_name} transferred: {num_keys} keys (pages), "
f"{host_indices.numel()} host_indices, {total_bytes} bytes, "
f"total time: {elapsed_ms:.3f} ms, effective bandwidth: {bw:.2f} MB/s"
)
def batch_get_v1(
self,
keys: List[str],
host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]:
if not self._host_regs:
logger.error(
"HiCacheNixl batch_get_v1: register_mem_pool_host must be called first"
)
return [False] * len(keys)
key_strs, host_buffers = self._batch_preprocess(keys, host_indices, "get")
if not key_strs or not host_buffers:
return [False] * len(keys)
start_time = time.perf_counter()
results = self._batch_xfer(keys, key_strs, host_buffers, "READ")
elapsed_ms = (time.perf_counter() - start_time) * 1000
self._log_xfer_stats(
"batch_get_v1",
len(keys),
host_indices,
[s for _, s in host_buffers],
elapsed_ms,
)
return self._batch_get_postprocess(host_indices, results)
def batch_set_v1(
self,
keys: List[str],
host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]:
# skip on MLA backup rank
if self.backup_skip:
return [True] * len(keys)
if len(keys) == 0:
return []
if not self._host_regs:
logger.error(
"HiCacheNixl batch_set_v1: register_mem_pool_host must be called first"
)
return [False] * len(keys)
key_strs, host_buffers = self._batch_preprocess(keys, host_indices, "set")
if not key_strs or not host_buffers:
return [False] * len(keys)
start_time = time.perf_counter()
results = self._batch_xfer(keys, key_strs, host_buffers, "WRITE")
elapsed_ms = (time.perf_counter() - start_time) * 1000
self._log_xfer_stats(
"batch_set_v1",
len(keys),
host_indices,
[s for _, s in host_buffers],
elapsed_ms,
)
return results
@@ -0,0 +1,133 @@
################################################################################
# IMPORTANT
# 1. to enable a plugin, add "active = true" in the corresponding section
# 2. the configs inside plugin.posix (i.e., use_aio, use_uring, use_posix_aio)
# are mutually exclusive
################################################################################
########################################
# GLOBAL NIXL HICACHE SETTINGS
########################################
# Open FILE-backend cache files with O_DIRECT when supported.
use_direct_io = true
# Built-in background cleaner for FILE-backed L3 storage. Set to false when an
# external cleaner is responsible for eviction. OBJ plugins ignore this.
l3_cleaner_enabled = true
# Background cleaner watermarks for FILE-backed L3 storage. OBJ plugins ignore these.
l3_cleaner_high_watermark = 80.0
l3_cleaner_low_watermark = 70.0
########################################
# POSIX FILE SYSTEM BACKEND
########################################
[plugin.posix]
# Configuration for the POSIX file-based backend.
#
# The supported backends include:
# 1. AIO (`use_aio = "true"`)
# 2. io_uring (`use_uring = "true"`)
# 3. POSIX AIO (`use_posix_aio = "true"`)
#
# If not specified, NIXL will automatically detect and use available backends based on the following default priority: AIO > io_uring > POSIX AIO
# Enable Linux io_uring for async I/O (recommended if supported)
use_uring = "true"
# Enable POSIX AIO (alternative async I/O mechanism)
# use_posix_aio = "true"
# Enable generic AIO
# use_aio = "true"
# Whether this plugin is eligible for selection
active = true
########################################
# NVIDIA GDS (GPUDirect Storage) BACKEND
########################################
[plugin.gds]
# Configuration for NVIDIA GPUDirect Storage
# Requires compatible GPU, driver, and filesystem support
# Number of requests per batch, default: 128
batch_pool_size = 128
# Maximum number of requests issued at once, default: 128
batch_limit = 128
# Maximum size of a single request (bytes), default: 16MB
max_request_size = 16777216 # 16 MB
########################################
# MULTI-THREADED GDS BACKEND
########################################
[plugin.gds_mt]
# Multi-threaded variant of the GDS backend
# Number of worker threads, default: 4
thread_count = 4
########################################
# 3FS (THIRD-PARTY FILE SYSTEM) BACKEND
########################################
[plugin.3fs]
# Configuration for 3FS backend
# Requires the 3FS library to be installed and mounted
# Mount point of the 3FS filesystem
mount_point = "/mnt/3fs"
# Memory configuration mode:
# dram - use DRAM
# dram_zc - DRAM with zero-copy
# auto - let backend decide
mem_config = "dram"
# Size of the I/O pool
# Valid range: [2^6, 2^20]
iopool_size = 64
########################################
# OBJECT STORAGE BACKEND (S3 / COMPATIBLE)
########################################
[plugin.obj]
# Object storage backend (e.g., S3-compatible services)
# Number of client worker threads, default: 4
num_threads = 4
# Override endpoint (useful for non-AWS S3 services)
endpoint_override = ""
# Connection scheme: http or https, default: http
scheme = "http"
# Cloud region (if applicable)
region = ""
# Request checksum behavior:
# required, supported
req_checksum = "required"
# Custom CA bundle path (if needed)
ca_bundle = ""
# Credentials
access_key = ""
secrete_key = ""
session_token = ""
# Use virtual-hosted-style addressing (true/false), default: true
use_virtual_addressing = "true"
# Default bucket name
bucket = ""
@@ -0,0 +1,270 @@
"""Background disk cleaner for the NIXL FILE HiCache backend."""
from __future__ import annotations
import concurrent.futures
import logging
import os
import re
import threading
import time
from dataclasses import dataclass, field
from typing import Iterable, Optional
from sglang.srt.mem_cache.storage.nixl.nixl_routing import BUCKET_HEX_CHARS
logger = logging.getLogger(__name__)
_DEFAULT_INTERVAL_SEC = 30.0
_DEFAULT_RECHECK_GROUPS = 50
_DEFAULT_HIGH_WATERMARK = 80.0
_DEFAULT_LOW_WATERMARK = 70.0
_RANK_SUFFIX_RE = re.compile(r"_(\d+)_(\d+)$")
_KV_SUFFIXES = ("_k", "_v")
_BUCKET_NAME_RE = re.compile(rf"^[0-9a-f]{{{BUCKET_HEX_CHARS}}}$")
@dataclass
class _GroupInfo:
"""Metadata for one logical cache key group in a cleaner tick."""
base_key: str
mtime: float = 0.0
size: int = 0
paths: set[str] = field(default_factory=set)
# Physical files in one logical group can hash to different base dirs
# because TP-rank and zero-copy K/V suffixes are part of the routed key.
base_dirs: set[str] = field(default_factory=set)
def _parse_group_key(name: str) -> str:
"""Return the logical cache-key group for one NIXL FILE object name.
The physical names are produced by ``HiCacheNixl._get_suffixed_key`` and
``HiCacheNixl._get_key_list_from_meta``.
"""
stem = name
for suffix in _KV_SUFFIXES:
if stem.endswith(suffix):
stem = stem[: -len(suffix)]
break
match = _RANK_SUFFIX_RE.search(stem)
if match is not None:
stem = stem[: match.start()]
return stem
def _safe_unlink(path: str) -> tuple[bool, int]:
"""Best-effort unlink; returns whether a file was removed and its size."""
try:
size = os.stat(path).st_size
os.unlink(path)
return True, size
except FileNotFoundError:
logger.debug("NIXL L3 file already removed before cleanup: %s", path)
return False, 0
except OSError:
logger.debug("Failed to unlink NIXL L3 file %s", path, exc_info=True)
return False, 0
class HiCacheL3Cleaner:
"""Delete old NIXL FILE cache entries when disk usage exceeds watermarks.
Cleanup operates on physical file names only, so it is compatible with MHA,
MLA, and DSA naming as long as the keys are generated by ``HiCacheNixl``.
"""
def __init__(
self,
storage_dirs: list[str] | str,
tp_rank: int,
*,
high_watermark: float = _DEFAULT_HIGH_WATERMARK,
low_watermark: float = _DEFAULT_LOW_WATERMARK,
interval_sec: float = _DEFAULT_INTERVAL_SEC,
recheck_groups: int = _DEFAULT_RECHECK_GROUPS,
unlink_workers: Optional[int] = None,
) -> None:
if isinstance(storage_dirs, str):
storage_dirs = [storage_dirs] if storage_dirs else []
self.storage_dirs = [path for path in storage_dirs if path]
self.tp_rank = tp_rank
self.high_watermark = high_watermark
self.low_watermark = low_watermark
self.interval_sec = interval_sec
self.recheck_groups = max(1, recheck_groups)
self.unlink_workers = unlink_workers or max(
8, 8 * max(len(self.storage_dirs), 1)
)
if self.low_watermark >= self.high_watermark:
raise ValueError(
"L3 cleaner low_watermark must be lower than high_watermark "
f"(low_watermark={self.low_watermark}, "
f"high_watermark={self.high_watermark})"
)
self._stop = threading.Event()
self._thread: Optional[threading.Thread] = None
def start(self) -> None:
"""Start the cleaner thread on TP rank 0."""
if self.tp_rank != 0 or not self.storage_dirs:
return
if self._thread is not None and self._thread.is_alive():
return
self._thread = threading.Thread(
target=self._loop, name="hicache-l3-cleaner", daemon=True
)
self._thread.start()
logger.info(
"HiCacheL3Cleaner started: dirs=%s high=%.1f%% low=%.1f%% "
"interval=%.1fs unlink_workers=%d",
self.storage_dirs,
self.high_watermark,
self.low_watermark,
self.interval_sec,
self.unlink_workers,
)
def stop(self) -> None:
"""Stop the cleaner thread if it was started."""
self._stop.set()
if self._thread is not None:
self._thread.join(timeout=5.0)
self._thread = None
def _disk_usage_pct(self, path: str) -> float:
try:
stat = os.statvfs(path)
except OSError:
return 0.0
total = stat.f_blocks * stat.f_frsize
if total == 0:
return 0.0
available = stat.f_bavail * stat.f_frsize
return 100.0 * (total - available) / total
def _loop(self) -> None:
while not self._stop.is_set():
try:
self._tick()
except Exception:
logger.warning("NIXL L3 cleaner tick failed", exc_info=True)
if self._stop.wait(self.interval_sec):
break
def _tick(self) -> bool:
initial_pcts = {path: self._disk_usage_pct(path) for path in self.storage_dirs}
hot_dirs = {
path for path, pct in initial_pcts.items() if pct >= self.high_watermark
}
if not hot_dirs:
return False
scan_start = time.perf_counter()
groups: dict[str, _GroupInfo] = {}
for base_dir in self.storage_dirs:
self._scan_base_dir(base_dir, groups)
ordered = sorted(
(group for group in groups.values() if group.base_dirs & hot_dirs),
key=lambda group: group.mtime,
)
if not ordered:
return False
deleted_groups = 0
deleted_files = 0
bytes_deleted = 0
with concurrent.futures.ThreadPoolExecutor(
max_workers=self.unlink_workers,
thread_name_prefix="hicache-l3-unlink",
) as pool:
idx = 0
while idx < len(ordered) and not self._stop.is_set():
batch = ordered[idx : idx + self.recheck_groups]
paths = list(self._iter_group_paths(batch))
for removed, removed_bytes in pool.map(_safe_unlink, paths):
if removed:
deleted_files += 1
bytes_deleted += removed_bytes
deleted_groups += len(batch)
idx += len(batch)
if all(
self._disk_usage_pct(path) < self.low_watermark for path in hot_dirs
):
break
final_pcts = {path: self._disk_usage_pct(path) for path in self.storage_dirs}
logger.info(
"NIXL L3 cleanup: deleted %d groups / %d files (%.2f GiB) in %.1fs, "
"initial_hot=%s final=%s",
deleted_groups,
deleted_files,
bytes_deleted / (1024**3),
time.perf_counter() - scan_start,
{path: f"{initial_pcts[path]:.1f}%" for path in hot_dirs},
{path: f"{pct:.1f}%" for path, pct in final_pcts.items()},
)
return True
def _scan_base_dir(self, base_dir: str, groups: dict[str, _GroupInfo]) -> None:
if not os.path.isdir(base_dir):
return
try:
bucket_entries = list(os.scandir(base_dir))
except OSError:
logger.warning("NIXL L3 cleaner failed to scan %s", base_dir, exc_info=True)
return
for bucket_entry in bucket_entries:
try:
is_bucket_dir = bucket_entry.is_dir(follow_symlinks=False)
except OSError:
continue
if (
not is_bucket_dir
or _BUCKET_NAME_RE.fullmatch(bucket_entry.name) is None
):
continue
self._scan_bucket(base_dir, bucket_entry.path, groups)
def _scan_bucket(
self, base_dir: str, bucket_path: str, groups: dict[str, _GroupInfo]
) -> None:
try:
entries = list(os.scandir(bucket_path))
except OSError:
logger.debug(
"NIXL L3 cleaner skipped bucket %s", bucket_path, exc_info=True
)
return
for entry in entries:
try:
if not entry.is_file(follow_symlinks=False):
continue
stat = entry.stat(follow_symlinks=False)
except OSError:
continue
group_key = _parse_group_key(entry.name)
group = groups.setdefault(group_key, _GroupInfo(group_key))
group.paths.add(entry.path)
group.base_dirs.add(base_dir)
group.size += stat.st_size
group.mtime = max(group.mtime, stat.st_mtime)
def _iter_group_paths(self, groups: Iterable[_GroupInfo]) -> Iterable[str]:
seen: set[str] = set()
for group in groups:
for path in sorted(group.paths):
if path in seen:
continue
seen.add(path)
yield path
@@ -0,0 +1,193 @@
"""NIXL memory-registration helpers, exposed as context managers.
A ``NixlRegistry`` instance bundles the agent, the memory type, and
(optionally) the file manager. Its ``storage(...)`` method is a context
manager that performs the entire register-and-build-descs sequence for
the storage side of a transfer on entry, yields the ``xfer_descs`` (or
None on failure), and unwinds ``agent.deregister_memory`` plus any
``os.close(fd)`` on exit.
The host side is pre-registered up front by ``HiCacheNixl`` and is not
touched per transfer.
"""
import logging
import threading
from contextlib import contextmanager
from typing import List, Optional
from .nixl_utils import NixlFileManager
logger = logging.getLogger(__name__)
def _buffer_sizes(buffers) -> Optional[List[int]]:
"""Per-buffer byte sizes for ``(addr, len)`` tuple inputs."""
if not buffers or not isinstance(buffers[0], tuple):
return None
return [b[1] for b in buffers]
class NixlRegistry:
"""Owns the (agent, mem_type, file_manager) triple and provides a
context manager for the storage side of a transfer.
A single instance is created once per HiCacheNixl in __init__ and
reused for every transfer.
"""
def __init__(
self,
agent,
mem_type: str,
file_manager: Optional[NixlFileManager] = None,
):
self.agent = agent
self.mem_type = mem_type
self.file_manager = file_manager
# OBJ devIds key a process-wide map in the NIXL OBJ plugin
# (devIdToObjKey_) that is not protected by a lock, so concurrent
# OBJ registrations must use disjoint devId ranges. Allocate them
# from a single monotonic counter.
self._obj_devid_lock = threading.Lock()
self._obj_devid_next = 1
self.path_mode = mem_type == "FILE" and self._probe_path_mode()
if mem_type == "FILE" and self.path_mode:
logger.info("HiCacheNixl: path-mode FILE registration active.")
elif mem_type == "FILE":
# TODO: NIXL 1.3.0 adds path-mode support; remove this fd fallback once 1.3.0 is widely installed.
logger.info(
"HiCacheNixl: the installed NIXL build does not "
"support path-mode FILE registration; using legacy "
"fd registration."
)
@contextmanager
def _open_files(self, paths: List[str], create: bool):
"""Open fds for ``paths``; close all of them on exit.
Yields the list of fds, or None if any open fails (already-opened
fds are closed before returning by the same ``finally``).
"""
fds: List[int] = []
try:
for path in paths:
fd = self.file_manager.open_file(path, create=create)
if fd is None:
yield None
return
fds.append(fd)
yield fds
finally:
for fd in fds:
self.file_manager.close_file(fd)
@contextmanager
def _registered(self, items: List[tuple], mem_type: str):
"""Register ``items`` with NIXL; deregister on exit.
Yields the registration handle, or None if registration fails.
"""
reg = None
if items:
reg_descs = self.agent.get_reg_descs(items, mem_type)
if reg_descs is not None:
try:
reg = self.agent.register_memory(reg_descs)
except Exception as e:
logger.error(f"Failed to register memory of type {mem_type}: {e}")
try:
yield reg
finally:
if reg is not None:
try:
self.agent.deregister_memory(reg)
except Exception as e:
logger.debug("deregister_memory skipped: %s", e)
def _probe_path_mode(self) -> bool:
"""Probe whether NIXL honours path-mode metaInfo.
Register a FILE_SEG with a valid path-mode string pointing at a
nonexistent path (no 'create' flag). A path-mode-capable NIXL tries
to open() the path, fails with NIXL_ERR_BACKEND, and raises. A
pre-path-mode NIXL ignores metaInfo and returns NIXL_SUCCESS.
Error from register_memory => path mode supported.
"""
reg_descs = self.agent.get_reg_descs(
[(0, 4096, 1, "rw:/nonexistent-nixl-probe")], "FILE"
)
if reg_descs is None:
return False
try:
reg = self.agent.register_memory(reg_descs)
if reg is not None:
try:
self.agent.deregister_memory(reg)
except Exception:
pass
return False
except Exception:
return True
@contextmanager
def storage(self, buffers, keys, direction):
"""Open + register the storage side; deregister and close fds on exit.
Yields the storage xfer_descs, or None on failure. For the FILE
backend, files are created (O_CREAT) when ``direction == "WRITE"``.
"""
sizes = _buffer_sizes(buffers)
if sizes is None:
yield None
return
if self.mem_type == "FILE":
if self.path_mode:
parts = ["rw", "create"] if direction == "WRITE" else ["ro"]
if self.file_manager.use_direct_io:
parts.append("direct")
spec = ",".join(parts)
tuples = [
(0, sizes[i], i + 1, f"{spec}:{keys[i]}") for i in range(len(keys))
]
with self._registered(tuples, "FILE") as reg:
if reg is None:
yield None
return
yield reg.trim()
else:
with self._open_files(keys, create=(direction == "WRITE")) as fds:
if fds is None:
yield None
return
tuples = [(0, sizes[i], fds[i], keys[i]) for i in range(len(keys))]
with self._registered(tuples, "FILE") as reg:
if reg is None:
yield None
return
yield self.agent.get_xfer_descs(
[(0, sizes[i], fds[i]) for i in range(len(fds))], "FILE"
)
else: # OBJ
# Reg tuple: (addr=0, size, devId, metaInfo=key).
# Xfer tuple: (addr=0, size, devId). devId links each xfer desc
# back to its registered object's metaInfo, so devIds must be
# unique within the list AND globally unique across concurrent
# storage() calls (the OBJ plugin's devIdToObjKey_ map is shared
# and unlocked). NIXL's pybind layer requires position 3 to be
# int, hence the key goes in metaInfo (position 4).
n = len(keys)
with self._obj_devid_lock:
base = self._obj_devid_next
self._obj_devid_next += n
dev_ids = list(range(base, base + n))
tuples = [(0, sizes[i], dev_ids[i], keys[i]) for i in range(n)]
with self._registered(tuples, "OBJ") as reg:
if reg is None:
yield None
return
yield self.agent.get_xfer_descs(
[(0, sizes[i], dev_ids[i]) for i in range(n)],
self.mem_type,
)
@@ -0,0 +1,29 @@
"""Deterministic path routing for NIXL FILE-backed HiCache storage."""
import hashlib
BUCKET_HEX_CHARS = 2
_BUCKET_MASK = (1 << (4 * BUCKET_HEX_CHARS)) - 1
def stable_key_hash(key: str) -> int:
"""Return a process-stable 64-bit hash for a NIXL storage key."""
return int.from_bytes(
hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest(), "big"
)
def route_key(key: str, num_disks: int) -> tuple[int, str]:
"""Return the storage disk index and bucket directory for a storage key."""
if num_disks <= 0:
raise ValueError("num_disks must be positive")
key_hash = stable_key_hash(key)
return (
(key_hash >> 16) % num_disks,
f"{key_hash & _BUCKET_MASK:0{BUCKET_HEX_CHARS}x}",
)
def route_disk(key: str, num_disks: int) -> int:
"""Return the storage disk index for a storage key."""
return route_key(key, num_disks)[0]
@@ -0,0 +1,315 @@
import logging
import os
from typing import Optional
from sglang.srt.environ import envs
from sglang.srt.mem_cache.storage.nixl.nixl_routing import (
_BUCKET_MASK,
BUCKET_HEX_CHARS,
route_key,
)
logger = logging.getLogger(__name__)
_SGLANG_NIXL_CONFIG_KEYS = {
"use_direct_io",
"l3_cleaner_enabled",
"l3_cleaner_high_watermark",
"l3_cleaner_low_watermark",
}
class NixlBackendConfig:
"""Handles NIXL backend configurations"""
def __init__(self, config: Optional[dict[str, str]] = None):
"""Initialize backend configuration.
Args:
config: configurations in a dictionary. This config comes from --hicache-storage-backend-extra-config
config can be in two forms:
1. fully qualified form (for all plugins, some of them are enabled, others not):
{'plugin': { 'posix': {...}, 'gds': {...}, ...}}
2. flat form (for a specific selected plugin), assuming all params apply to a selected plugin
{'param1': 'value1', 'param2': 'value2', ...}
"""
self.config = config or {}
def get_use_direct_io(self) -> bool:
"""Return True if O_DIRECT should be requested when opening files.
Checks the top-level ``use_direct_io`` key in the long-form JSON config first,
then falls back to the ``SGLANG_HICACHE_NIXL_USE_DIRECT_IO`` environment variable
(default: enabled).
"""
if "use_direct_io" in self.config:
return bool(self.config["use_direct_io"])
return envs.SGLANG_HICACHE_NIXL_USE_DIRECT_IO.get()
def get_l3_cleaner_config(self) -> dict:
"""Return typed NIXL FILE L3 cleaner options from top-level config."""
config = {
"enabled": True,
"high_watermark": 80.0,
"low_watermark": 70.0,
}
if "l3_cleaner_enabled" in self.config:
enabled = self.config["l3_cleaner_enabled"]
if not isinstance(enabled, bool):
raise ValueError("l3_cleaner_enabled must be a boolean")
config["enabled"] = enabled
key_map = {
"l3_cleaner_high_watermark": ("high_watermark", float),
"l3_cleaner_low_watermark": ("low_watermark", float),
}
for raw_key, (cleaner_key, parser) in key_map.items():
if raw_key in self.config:
config[cleaner_key] = parser(self.config[raw_key])
return config
def get_specified_plugin(self) -> str:
"""decide which plugin to use: either config or SGLANG_HICACHE_NIXL_BACKEND_PLUGIN specifies the plugin, if not, use "auto" """
if "plugin" in self.config:
# fully qualified form: {'plugin': { 'posix': {...}, 'gds': {...}, ...}}
# choose the FIRST active plugin
for key, item in self.config["plugin"].items():
if item.get("active", False) in [True, "true", "True"]:
plugin = key.upper()
break
else:
# config is empty, or in flat form {'param1': 'value1', 'param2': 'value2', ...}
plugin = os.getenv("SGLANG_HICACHE_NIXL_BACKEND_PLUGIN", "auto")
return plugin
def get_backend_initparams(self, backend_name) -> dict:
"""Get initialization parameters from config of NIXL backend for backend creation.
Args:
backend_name: a specific backend's name (already converted "auto" into a specific backend name)
"""
initparams = {}
# config can be in two forms:
if "plugin" in self.config:
# fully qualified form: {'plugin': { 'posix': {...}, 'gds': {...}, ...}}
if backend_name.lower() in self.config["plugin"]:
config_data = self.config["plugin"][backend_name.lower()]
else:
logger.debug(
f"No specific config found for plugin {backend_name} in extra_config. Use default init params."
)
config_data = {}
else:
# flat form {'param1': 'value1', 'param2': 'value2', ...}
config_data = self.config
for key, value in config_data.items():
# These keys are consumed by SGLang itself, not by NIXL plugins.
if key in _SGLANG_NIXL_CONFIG_KEYS:
continue
initparams[key] = str(value)
return initparams
class NixlBackendSelection:
"""Handles NIXL backend selection and creation."""
# Priority order for File-based plugins in case of auto selection
FILE_PLUGINS = ["3FS", "POSIX", "GDS_MT", "GDS"]
# Priority order for File-based plugins in case of auto selection (add more as needed)
OBJ_PLUGINS = ["OBJ"] # Based on Amazon S3 SDK
def __init__(
self, plugin: str = "auto", nixlconfig: Optional[NixlBackendConfig] = None
):
"""Initialize backend selection.
Args:
plugin: Plugin to use (default "auto" selects best available).
Can be a file plugin (3FS, POSIX, GDS, GDS_MT) or
an object plugin (OBJ).
"""
self.plugin = plugin
self.backend_name = None
self.mem_type = None
self.nixlconfig = nixlconfig
def create_backend(self, agent) -> bool:
"""Create the appropriate NIXL backend based on configuration."""
try:
plugin_list = agent.get_plugin_list()
logger.debug(f"Available NIXL plugins: {plugin_list}")
# Handle explicit plugin selection or auto priority
if self.plugin == "auto":
# Try all file plugins first
for plugin in self.FILE_PLUGINS:
if plugin in plugin_list:
self.backend_name = plugin
break
# If no file plugin found, try object plugins
if not self.backend_name:
for plugin in self.OBJ_PLUGINS:
if plugin in plugin_list:
self.backend_name = plugin
break
else:
# Use explicitly requested plugin
self.backend_name = self.plugin
if self.backend_name not in plugin_list:
logger.error(
f"Backend {self.backend_name} not available in plugins: {plugin_list}"
)
return False
# obtain initparams for the backend from the NIXL config
initparams = (
self.nixlconfig.get_backend_initparams(self.backend_name)
if self.nixlconfig
else {}
)
# Create backend and set memory type
if self.backend_name in self.OBJ_PLUGINS and "bucket" not in initparams:
bucket = os.environ.get("AWS_DEFAULT_BUCKET")
if not bucket:
logger.error(
"AWS_DEFAULT_BUCKET environment variable must be set for object storage"
)
return False
initparams["bucket"] = bucket
# create backend using initialization parameters
agent.create_backend(self.backend_name, initparams)
logger.info(
f"NixlBackendSelection.create_backend: backend_name {self.backend_name} initparams {initparams} customParams {agent.get_backend_params(self.backend_name)} supported plugins {plugin_list}"
)
self.mem_type = "OBJ" if self.backend_name in self.OBJ_PLUGINS else "FILE"
logger.debug(
f"Created NIXL backend: {self.backend_name} with memory type: {self.mem_type}"
)
return True
except Exception as e:
logger.error(
f"Failed to create NIXL backend: {e}, backend_name {self.backend_name}, supported plugins {plugin_list} initparams {initparams}"
)
return False
class NixlFileManager:
"""Handles file system operations for NIXL."""
def __init__(self, base_dir: "list[str] | str", use_direct_io: bool = True):
"""
Initialize file manager.
Args:
base_dir: Base directory or ordered base directories for tensor files.
use_direct_io: If True, open files with O_DIRECT (bypasses OS page cache).
Falls back to buffered I/O with a warning when O_DIRECT is unavailable.
"""
if isinstance(base_dir, str):
self.base_dirs = [base_dir] if base_dir else []
else:
self.base_dirs = [d for d in base_dir if d]
self.use_direct_io = use_direct_io
self._created_bucket_dirs: set[str] = set()
if not self.base_dirs:
logger.debug(
f"Initialized file manager without a base directory. Direct I/O: {use_direct_io}"
)
else:
for base in self.base_dirs:
os.makedirs(base, exist_ok=True)
self.ensure_all_bucket_dirs()
logger.debug(
f"Initialized file manager with base directories: {self.base_dirs}. Direct I/O: {use_direct_io}"
)
def clear(self) -> None:
"""Clear all files below every configured base directory."""
if not self.base_dirs:
logger.warning("Base directories are empty, skipping clear operation")
return
for base in self.base_dirs:
try:
for root, _dirs, files in os.walk(base):
for file in files:
file_path = os.path.join(root, file)
try:
os.remove(file_path)
except OSError as e:
logger.warning(f"Failed to remove file {file_path}: {e}")
except Exception as e:
logger.error(f"Failed to clear base directory {base}: {e}")
logger.debug(f"Cleared all files in base directories: {self.base_dirs}")
def ensure_all_bucket_dirs(self) -> None:
"""Pre-create every possible bucket directory under each base dir.
Called once when path mode is active so NIXL O_CREAT writes never
fail due to a missing parent directory.
"""
for base in self.base_dirs:
for i in range(_BUCKET_MASK + 1):
os.makedirs(
os.path.join(base, f"{i:0{BUCKET_HEX_CHARS}x}"),
exist_ok=True,
)
def iter_all_base_dirs(self) -> list[str]:
"""Return base directories that may contain NIXL FILE cache entries."""
return list(self.base_dirs)
def get_file_path(self, key: str) -> str:
"""Get full file path for a given key."""
if not self.base_dirs:
return key
disk_idx, bucket = route_key(key, len(self.base_dirs))
return os.path.join(self.base_dirs[disk_idx], bucket, key)
def open_file(self, file_path: str, create: bool = False) -> Optional[int]:
"""Open a file and return its file descriptor.
If ``create`` is True, the file is created if it does not exist
(mode 0o644, no truncation). When ``self.use_direct_io`` is True,
the file is opened with ``O_DIRECT`` (bypasses the OS page cache);
falls back to buffered I/O with a warning if ``O_DIRECT`` is
unavailable on this platform.
"""
flags = os.O_RDWR | os.O_CREAT if create else os.O_RDWR
if self.use_direct_io:
if hasattr(os, "O_DIRECT"):
flags |= os.O_DIRECT
else:
logger.warning(
"use_direct_io is True, but O_DIRECT is not available on "
"this system. Falling back to buffered I/O."
)
try:
if create:
parent = os.path.dirname(file_path)
if parent and parent not in self._created_bucket_dirs:
os.makedirs(parent, exist_ok=True)
self._created_bucket_dirs.add(parent)
return os.open(file_path, flags, 0o644)
except Exception as e:
logger.error(f"Failed to open file {file_path}: {e}")
return None
def close_file(self, fd: int) -> bool:
"""Close a file descriptor."""
try:
os.close(fd)
return True
except Exception as e:
logger.error(f"Failed to close file descriptor {fd}: {e}")
return False
@@ -0,0 +1,121 @@
# SiMM as L3 KV Cache
This document describes how to use SiMM as the L3 KV cache for SGLang.
## About SiMM
SiMM(Scalable In-Memory Middleware) is a distributed, high-performance, elastic cache acceleration layer for all AI workloads.
For more details about SiMM, please refer to [SiMM project](https://github.com/scitix/SiMM) and [SiMM documents](https://github.com/scitix/SiMM/tree/main/docs).
### SiMM & SGLang HiCache
SiMM serves as a high-performance L3 storage backend for SGLang HiCache, enabling distributed KV cache storage across multiple servers with RDMA-baed transport. This integration addresses the capacity limitations of traditional GPU-only or GPU+CPU caching by providing virtually unlimited cache storage through a distributed memory pool.
When a cache miss occurs in L1 and L2, HiCache automatically fetches the required KV cache from SiMM's distributed memory pool. The system uses intelligent prefetching strategies to minimize latency, and utilize RDMA technology and zero-copy technique to ensure high-bandwidth, low-latency data transfer between SGLang instances and SiMM data servers.
## Install SiMM
**from source**
Clone SiMM project:
```bash
git clone https://github.com/scitix/SiMM --recursive
```
Install dependencies:
```bash
cd SiMM
bash configure.sh
```
Build and install SiMM:
```bash
bash build.sh --mode=release --clean
```
For more details, please refer to [SiMM official installation guide](https://github.com/scitix/SiMM/blob/main/README.md).
## Deployment
**SiMM**
Before launch `SGLang server` with SiMM, you should launch SiMM `cluster manager service` and `data server service`.
You can visit [SiMM official deploy guide](https://github.com/scitix/SiMM/blob/main/docs/deploy_guide.md) and deploy SiMM on your K8S cluster with RDMA network.
**Start the `SGLang server` with SiMM enabled:**
There are three ways to configure SiMM:
1. Via extra configuration passed through sglang parameters
2. Using JSON configuration files
3. Using environment variables
SiMM loads configuration in the following priority order:
1. If SiMM-specific options are provided in `--hicache-storage-backend-extra-config`, they are used first.
2. If not, SiMM checks whether the environment variable `DEFAULT_SIMM_CONFIG_PATH_ENV` is set, and loads the JSON config file from that path.
3. If neither of the above is provided, SiMM falls back to environment variables.
**HiCache Related Parameters for SGLang Server**
For a comprehensive overview of HiCache-related parameters, please refer to [this document](https://docs.sglang.io/advanced_features/hicache_design.html#related-parameters).
Note that, for `--hicache-mem-layout {layer_first,page_first,page_first_direct}`, which specifies the memory layout for the host memory pool, `page_first` or `page_first_direct` are required if use SiMM backend.
### Distributed Deployment
**Using extra-config of sglang arguments to configure SiMM**
```bash
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend simm \
--model-path [model_path] \
--hicache-storage-backend-extra-config '{"manager_address": "127.0.0.1:30001"}'
```
**Using JSON file to configure SiMM**
SGLang server can load SiMM config from `SGLANG_HICACHE_SIMM_CONFIG_PATH`.
```bash
export SGLANG_HICACHE_SIMM_CONFIG_PATH=/sgl-workspace/sglang/benchmark/hicache/simm_config.json
echo '{
"manager_address": "127.0.0.1:30001"
}' > ${SGLANG_HICACHE_SIMM_CONFIG_PATH}
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend simm \
--model-path [model_path]
```
**Using env variables to configure SiMM**
```bash
SIMM_CLUSTER_MANAGER="127.0.0.1:30001"
python -m sglang.launch_server \
--enable-hierarchical-cache \
--hicache-storage-backend simm \
--model-path [model_path]
```
## Test SiMM
This test is intended for developers to quickly verify that the SiMM class interfaces are functioning correctly.
First, start the `cluster manager service` and `data server service`. Then run the `test_hicache_simm.py`.
```bash
SIMM_CLUSTER_MANAGER="127.0.0.1:30001" \
python3 [path of test_hicache_simm.py]
```
If all tests pass, the message "✅ All tests passed" will be printed at the end.
@@ -0,0 +1,544 @@
import json
import logging
import os
import re
import time
import uuid
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Dict, List, Optional
import torch
from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorage,
HiCacheStorageConfig,
HiCacheStorageExtraInfo,
)
from sglang.srt.mem_cache.pool_host import HostKVCache
# Third Party
try:
from simm.kv import BlockView, Store, register_mr, set_flag
except ImportError as e:
raise ImportError(
"Please install simm by following the instructions at https://github.com/scitix/SiMM "
"to run SGLang with SimmConnector."
) from e
SGLANG_HICACHE_SIMM_JSON_ENV_VAR = "SGLANG_HICACHE_SIMM_CONFIG_PATH"
logger = logging.getLogger(__name__)
@dataclass
class SiMMConfig:
manager_address: str
clnt_threadpool_size: int
enable_profile: bool
@staticmethod
def from_file() -> "SiMMConfig":
"""Load the config from a JSON file."""
if os.environ.get(SGLANG_HICACHE_SIMM_JSON_ENV_VAR) is None:
raise RuntimeError(
f"Config file path not set. Please set {SGLANG_HICACHE_SIMM_JSON_ENV_VAR}"
)
file_path = os.environ.get(SGLANG_HICACHE_SIMM_JSON_ENV_VAR)
try:
with open(file_path) as fin:
config = json.load(fin)
except Exception as e:
raise RuntimeError(f"Failed to load config from {file_path}: {str(e)}")
if "manager_address" not in config:
raise ValueError("Manager_address is required in config file")
return SiMMConfig(
manager_address=config.get("manager_address"),
clnt_threadpool_size=config.get("clnt_threadpool_size", 10),
enable_profile=config.get("enable_profile", False),
)
@staticmethod
def load_from_extra_config(extra_config: dict) -> "SiMMConfig":
"""Load config from extra_config dictionary."""
if "manager_address" not in extra_config:
raise ValueError("manager_address is required in extra_config")
return SiMMConfig(
manager_address=extra_config.get("manager_address"),
clnt_threadpool_size=extra_config.get("clnt_threadpool_size", 10),
enable_profile=extra_config.get("enable_profile", False),
)
def get_current_process_numa() -> int:
"""
Return value: numa_node of current process, failed return -1
"""
try:
# get current cpu
with open("/proc/self/stat", "r") as f:
stat_data = f.read()
# the 39th field is processor
fields = stat_data.split()
if len(fields) < 39:
return -1
current_cpu = int(fields[38])
numa_path = f"/sys/devices/system/cpu/cpu{current_cpu}/node0"
if os.path.exists(numa_path) and os.path.islink(numa_path):
link_target = os.readlink(numa_path)
# parse numa node from path
match = re.search(r"node(\d+)$", link_target)
if match:
return int(match.group(1))
return -1
except Exception:
return -1
def get_numa_nic_mapping() -> Dict[int, List[str]]:
"""
Return value: Dict[numa_node, List(rdma_device_name)]
"""
ib_root = "/sys/class/infiniband"
device_map = defaultdict(list)
if not os.path.exists(ib_root):
logger.error(f"SiMM ERROR: {ib_root} not found. Are RDMA drivers loaded?")
return []
for device_name in os.listdir(ib_root):
numa_path = os.path.join(ib_root, device_name, "device", "numa_node")
numa_node = -1 # default value, if system is UMA.
try:
if os.path.exists(numa_path):
with open(numa_path, "r") as f:
content = f.read().strip()
numa_node = int(content)
except (IOError, ValueError):
pass
device_map[numa_node].append(device_name)
return device_map
class HiCacheSiMM(HiCacheStorage):
def __init__(
self, storage_config: HiCacheStorageConfig = None, mem_pool: HostKVCache = None
):
try:
extra_config = (
getattr(storage_config, "extra_config", None)
if storage_config
else None
)
# Load configuration with manager_address prioritized from extra_config if available
if (
extra_config is not None
and extra_config.get("manager_address") is not None
):
# Load from extra_config
self.config = SiMMConfig.load_from_extra_config(extra_config)
logger.info("SiMM Configuration loaded from extra_config successfully.")
else:
# Load from config file
self.config = SiMMConfig.from_file()
logger.info("SiMM Configuration loaded from file successfully.")
# Check if extra_backend_tag should be passed to SiMM data server
self.extra_backend_tag = None
if extra_config and "extra_backend_tag" in extra_config:
self.extra_backend_tag = extra_config["extra_backend_tag"]
logger.info(f"Using extra_backend_tag: {self.extra_backend_tag}")
# Set nic device according to current process numa node
nic_mapping = get_numa_nic_mapping()
logger.info(f"SiMM NUMA-awared allocation: {nic_mapping}")
current_numa = get_current_process_numa()
if current_numa >= 0:
rdma_devices = nic_mapping.get(current_numa)
if rdma_devices is not None and len(rdma_devices) > 0:
rdma_device_str = ",".join(rdma_devices)
os.environ["SICL_NET_DEVICES"] = rdma_device_str
logger.info(f"SiMM using rdma {rdma_device_str}")
# Set simm log path: /var/log/simm/{filename_ts}-{pid}/simm_clnt.log
filename_ts = datetime.now().strftime("%Y%m%d-%H%M%S")
log_file_path: str = (
f"/var/log/simm/{filename_ts}-{os.getpid()}/simm_clnt.log"
)
cm_ip = self.config.manager_address.split(":")[0]
cm_port = self.config.manager_address.split(":")[1]
set_flag("cm_primary_node_ip", cm_ip)
set_flag("cm_primary_node_port", cm_port)
set_flag("clnt_log_file", log_file_path)
set_flag("clnt_thread_pool_size", str(self.config.clnt_threadpool_size))
self.store = Store()
logger.info("SiMM store setup successfully.")
self.mr_ext = None
self.warmup()
logger.info("SiMM store warmup successfully.")
if storage_config is not None:
self.model_name = storage_config.model_name
self.is_mla_backend = storage_config.is_mla_model
self.local_rank = storage_config.tp_rank
self.pp_rank = storage_config.pp_rank
self.pp_size = storage_config.pp_size
else:
self.model_name = ""
self.is_mla_backend = False
self.local_rank = 0
self.pp_rank = 0
self.pp_size = 1
self.enable_pp = self.pp_size > 1
if self.enable_pp:
self.mha_suffix = f"{self.local_rank}_{self.pp_rank}"
self.mla_suffix = f"{self.pp_rank}"
else:
self.mha_suffix = f"{self.local_rank}"
self.mla_suffix = ""
except ValueError as e:
logger.error("Configuration loading failed: %s", e)
raise
except Exception as exc:
logger.error("An error occurred while loading the configuration: %s", exc)
raise
def warmup(self):
"""Dryrun a key to warmup SiMM client"""
logger.info("begin warm up SiMM client")
start_time = time.perf_counter_ns()
warmup_key = "sglang_simm_warmup_key" + uuid.uuid4().hex
warmup_tensor = torch.frombuffer(
bytearray(warmup_key.encode()), dtype=torch.uint8
)
warmup_size = 4 * 1024 # 4 KB
block = self.store.allocate(warmup_size)
block_ = block.as_ref()
block_[: len(warmup_key)] = warmup_tensor
if self.store.put(warmup_key, block.view()) != 0:
logger.warning(f"SiMM client warmup put key {warmup_key} failed")
if not self.store.exists(warmup_key):
logger.warning(f"SiMM client warmup key {warmup_key} not exists")
got_block = self.store.allocate(warmup_size)
if self.store.get(warmup_key, got_block.view()) < 0:
logger.warning(f"SiMM client warmup get key {warmup_key} failed")
if not all(got_block.as_ref()[: len(warmup_key)] == warmup_tensor):
logger.warning(f"SiMM client warmup key {warmup_key} data wrong")
logger.info(
f"finish SiMM client warm up, cost {(time.perf_counter_ns() - start_time)/1000:.2f} us"
)
def register_mem_pool_host(self, mem_pool_host: HostKVCache):
super().register_mem_pool_host(mem_pool_host)
assert self.mem_pool_host.layout in [
"page_first",
"page_first_direct",
], "simm storage backend only support page first or page first direct layout"
buffer = self.mem_pool_host.kv_buffer
try:
self.mr_ext = register_mr(buffer)
if self.mr_ext is None:
logger.error(
f"Failed to register buffer, {buffer=}, please check buffer and RDMA network"
)
raise RuntimeError(f"Failed to register buffer to SiMM")
except TypeError as err:
logger.error("Failed to register buffer to SiMM: %s", err)
raise TypeError("SiMM Register Buffer Error.") from err
def _get_mha_buffer_meta(self, keys, indices):
ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices)
key_list = []
for key_ in keys:
key_list.append(f"{key_}_{self.mha_suffix}_k")
key_list.append(f"{key_}_{self.mha_suffix}_v")
if len(key_list) != len(ptr_list):
logger.error(
f"key size {len(key_list)} not equal with incides ptr size {len(ptr_list)}"
)
assert len(key_list) == len(ptr_list)
return key_list, ptr_list, element_size_list
def _get_mla_buffer_meta(self, keys, indices):
ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices)
key_list = []
for key_ in keys:
key_list.append(f"{key_}_{self.mla_suffix}_k")
if len(key_list) != len(ptr_list):
logger.error(
f"key size {len(key_list)} not equal with incides ptr size {len(ptr_list)}"
)
assert len(key_list) == len(ptr_list)
return key_list, ptr_list, element_size_list
def _batch_preprocess(self, keys, host_indices):
assert len(keys) > 0
assert len(keys) == len(host_indices) // self.mem_pool_host.page_size
if self.is_mla_backend:
return self._get_mla_buffer_meta(keys, host_indices)
else:
return self._get_mha_buffer_meta(keys, host_indices)
def _batch_postprocess(self, results: List[int], is_set_operate=False):
"""
for batch_get_into, results is Vector of integers,
where each element is the number of bytes read on success, or a negative value on error
for batch_put_from, results is Vector of integers,
where each element is 0 on success, or a negative value on error
"""
if self.is_mla_backend:
return [k_res == 0 if is_set_operate else k_res > 0 for k_res in results]
else:
kv_pairs = zip(results[::2], results[1::2])
return [
(
(k_res == 0 and v_res == 0)
if is_set_operate
else (k_res > 0 and v_res > 0)
)
for k_res, v_res in kv_pairs
]
def batch_get_v1(
self,
keys: List[str],
host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]:
# Apply extra_backend_tag prefix if available
if self.extra_backend_tag is not None:
prefix = self.extra_backend_tag
keys = [f"{prefix}_{key}" for key in keys]
t1 = time.perf_counter_ns()
key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
get_results = self._get_batch_zero_copy_impl(
key_strs, buffer_ptrs, buffer_sizes
)
t2 = time.perf_counter_ns()
total_size = sum([k_res if k_res > 0 else 0 for k_res in get_results])
if self.config.enable_profile:
logger.info(
f"SiMM batch_get_v1 {len(keys)} keys, total size: {total_size / 1024**2} MiB, \
using {(t2 - t1)/1000} us, Throughput: {total_size / 1024**3 / ((t2 - t1) / 1000**3):.2f} GiB/s"
)
return self._batch_postprocess(get_results, is_set_operate=False)
def batch_set_v1(
self,
keys: List[str],
host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]:
# Apply extra_backend_tag prefix if available
if self.extra_backend_tag is not None:
prefix = self.extra_backend_tag
keys = [f"{prefix}_{key}" for key in keys]
t1 = time.perf_counter_ns()
key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
exist_result = self._batch_exist_impl(key_strs)
t2 = time.perf_counter_ns()
if self.config.enable_profile:
logger.info(
f"SiMM batch exists {len(keys)} keys, using {(t2 - t1)/1000} us"
)
set_keys = []
set_buffer_ptrs = []
set_buffer_sizes = []
set_indices = []
set_results = [-1] * len(key_strs)
total_size = 0
for i in range(len(key_strs)):
if not exist_result[i]:
set_keys.append(key_strs[i])
set_buffer_ptrs.append(buffer_ptrs[i])
set_buffer_sizes.append(buffer_sizes[i])
set_indices.append(i)
total_size += buffer_sizes[i]
else:
set_results[i] = 0
# Only set non-existing keys to storage
if len(set_keys) > 0:
put_results = self._put_batch_zero_copy_impl(
set_keys, set_buffer_ptrs, set_buffer_sizes
)
for i in range(len(set_indices)):
set_results[set_indices[i]] = put_results[i]
t3 = time.perf_counter_ns()
if self.config.enable_profile:
logger.info(
f"SiMM batch_put_v1 {len(keys)} keys, total size: {total_size / 1024**2} MiB, \
using {(t3 - t2)/1000} us, Throughput: {total_size / 1024**3 / ((t3 - t2) / 1000**3):.2f} GiB/s"
)
return self._batch_postprocess(set_results, is_set_operate=True)
def set(
self,
key,
value: Optional[Any] = None,
target_location: Optional[List[int]] = None,
target_sizes: Optional[List[int]] = None,
) -> bool:
# Only support zero copy set for now
assert target_location is not None and target_sizes is not None
exist_result = self._batch_exist_impl([key])
if exist_result[0]:
return True
put_result = self._put_batch_zero_copy_impl(
[key], [target_location], [target_sizes]
)
return put_result[0] == 0
def batch_set(
self,
keys: List[str],
values: Optional[List[torch.Tensor]] = None,
target_locations: Optional[List[int]] = None,
target_sizes: Optional[List[int]] = None,
) -> bool:
# Only support zero copy set for now
assert target_locations is not None and target_sizes is not None
assert len(keys) == len(target_locations) == len(target_sizes)
if len(keys) == 0:
return False
for i in range(len(keys)):
if (
keys[i] is None
or target_locations[i] is None
or target_sizes[i] is None
):
return False
exist_result = self._batch_exist_impl(keys)
set_keys = []
set_target_locations = []
set_target_sizes = []
set_indices = []
for i in range(len(keys)):
if not exist_result[i]:
set_keys.append(keys[i])
set_target_locations.append(target_locations[i])
set_target_sizes.append(target_sizes[i])
set_indices.append(i)
# Only set non-existing keys to storage
put_result = self._put_batch_zero_copy_impl(
set_keys, set_target_locations, set_target_sizes
)
for i in range(len(set_indices)):
if put_result[i] == 0:
exist_result[set_indices[i]] = 1
# return the number of consecutive successful operations from the start.
success_count = 0
for i in range(len(keys)):
if exist_result[i] == 0:
break
success_count += 1
return success_count == len(keys)
def get(
self,
key,
target_location: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> bool:
assert target_location is not None and target_sizes is not None
get_result = self._get_batch_zero_copy_impl(
[key], [target_location], [target_sizes]
)
return get_result[0] >= 0
def batch_get(
self,
keys: List[str],
target_locations: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> int:
assert len(keys) == len(target_locations) == len(target_sizes)
if len(keys) == 0:
return 0
get_result = self._get_batch_zero_copy_impl(
keys, target_locations, target_sizes
)
if self.is_mla_backend:
key_multiplier = 1
else:
key_multiplier = 2
for i in range(len(keys)):
if get_result[i] < 0:
return i // key_multiplier
return len(keys) // key_multiplier
def exists(self, key) -> bool:
exist_result = self._batch_exist_impl([key])
return exist_result[0]
def batch_exists(
self, keys, extra_info: Optional[HiCacheStorageExtraInfo] = None
) -> int:
if self.is_mla_backend:
query_keys = [f"{key}_{self.mla_suffix}_k" for key in keys]
key_multiplier = 1
else:
query_keys = []
for key in keys:
query_keys.append(f"{key}_{self.mha_suffix}_k")
query_keys.append(f"{key}_{self.mha_suffix}_v")
key_multiplier = 2
t1 = time.perf_counter_ns()
exist_result = self._batch_exist_impl(query_keys)
t2 = time.perf_counter_ns()
if self.config.enable_profile:
logger.info(
f"SiMM batch exists {len(keys)} keys, using {(t2 - t1)/1000} us"
)
for i in range(len(query_keys)):
if not exist_result[i]:
return i // key_multiplier
return len(query_keys) // key_multiplier
def _put_batch_zero_copy_impl(
self, key_strs: List[str], buffer_ptrs: List[int], buffer_sizes: List[int]
) -> List[int]:
block_views = []
for i in range(len(buffer_ptrs)):
block_view = BlockView.from_buffer(
buffer_ptrs[i], buffer_sizes[i], self.mr_ext
)
block_views.append(block_view)
return self.store.mput(key_strs, block_views)
def _get_batch_zero_copy_impl(
self, key_strs: List[str], buffer_ptrs: List[int], buffer_sizes: List[int]
) -> List[int]:
block_views = []
for i in range(len(buffer_ptrs)):
block_view = BlockView.from_buffer(
buffer_ptrs[i], buffer_sizes[i], self.mr_ext
)
block_views.append(block_view)
return self.store.mget(key_strs, block_views)
def _batch_exist_impl(self, key_strs: List[str]) -> List[bool]:
return self.store.mexists(key_strs)
@@ -0,0 +1,181 @@
import logging
import uuid
import torch
from python.sglang.srt.mem_cache.storage.simm.hicache_simm import HiCacheSiMM
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
def generate_batch_query_keys(kv_num: int, config: HiCacheStorageConfig):
keys = ["test_" + str(uuid.uuid4()) for _ in range(kv_num)]
set_keys = []
for key in keys:
if config.is_mla_model:
set_keys.append(key + "_k")
else:
set_keys.append(key + f"_{config.tp_rank}_k")
set_keys.append(key + f"_{config.tp_rank}_v")
get_keys = set_keys
exist_keys = keys
return set_keys, get_keys, exist_keys
def create_mock_host_kv_cache(buffer_size, dtype=torch.float32):
"""Create a mock HostKVCache-like object for testing."""
buffer = torch.randn(buffer_size, dtype=dtype)
class MockHostKVCache:
def __init__(self, buffer):
self.kv_buffer = buffer
self.layout = "page_first"
self.page_size = 1 # Simple page size for testing
def get_page_buffer_meta(self, indices):
"""Mock implementation of get_page_buffer_meta."""
ptr_list = []
element_size_list = []
for idx in indices:
# Create mock pointers and sizes for each page
ptr_list.append(idx * self.page_size * self.kv_buffer.element_size())
element_size_list.append(self.page_size * self.kv_buffer.element_size())
return ptr_list, element_size_list
return MockHostKVCache(buffer), buffer
def test_single_operation():
"""Test the set API with a single key-value pair."""
print("=" * 100)
print("Testing single operation")
buffer_size = 1024 * 1024 * 16 # 16MB
value_elements = 1024
store = HiCacheSiMM()
mock_host_kv_cache, buffer = create_mock_host_kv_cache(buffer_size)
# Register the memory pool host - this is the proper workflow
store.register_mem_pool_host(mock_host_kv_cache)
value_size = value_elements * buffer.element_size()
key = str(uuid.uuid4())
set_slice = buffer[:value_elements]
get_slice = buffer[value_elements : 2 * value_elements]
set_location = set_slice.data_ptr()
get_location = get_slice.data_ptr()
# Test set operation
result = store.set(key, target_location=set_location, target_sizes=value_size)
assert result is True, f"❌set operation failed for key: {key}"
# Test exists operation
assert store.exists(key), f"❌key {key} should exist after set operation"
# Test get operation
result = store.get(key, target_location=get_location, target_sizes=value_size)
assert result is True, f"❌get operation failed for key: {key}"
# Compare the data using proper tensor indices
assert torch.allclose(
set_slice, get_slice, atol=1e-6
), f"❌get operation failed for key: {key}"
logger.info(f"✅ Single operation passed")
def test_batch_operation(config: HiCacheStorageConfig):
"""Test the batch set/get APIs with multiple key-value pairs."""
print("=" * 100)
print(f"Testing batch operation with config: {config}")
buffer_size = 1024 * 1024 * 16 # 16MB
value_elements = 256
kv_num = 13
store = HiCacheSiMM(config)
mock_host_kv_cache, buffer = create_mock_host_kv_cache(buffer_size)
store.register_mem_pool_host(mock_host_kv_cache)
value_size = value_elements * buffer.element_size()
set_keys, get_keys, exist_keys = generate_batch_query_keys(kv_num, config)
set_slices = [
buffer[i * value_elements : (i + 1) * value_elements]
for i in range(len(set_keys))
]
set_indices = torch.cat(set_slices)
# Test batch set operation
result = store.batch_set_v1(set_keys, set_indices)
assert all(result), f"❌batch set operation failed"
# Test batch exists operation
assert store.batch_exists(
exist_keys
), f"❌keys should exist after batch set operation"
# Test batch get operation
get_slices = [
buffer[
(len(set_keys) + i)
* value_elements : (len(set_keys) + i + 1)
* value_elements
]
for i in range(len(get_keys))
]
get_indices = torch.cat(get_slices)
result = store.batch_get_v1(get_keys, get_indices)
assert all(result), f"❌batch get operation failed"
for i in range(len(get_keys)):
assert torch.allclose(
set_slices[i], get_slices[i], atol=1e-6
), f"❌batch get operation failed for key: {get_keys[i]}"
logger.info(f"✅ Batch operation passed")
if __name__ == "__main__":
test_single_operation()
test_batch_operation(
HiCacheStorageConfig(
is_mla_model=False,
tp_rank=0,
tp_size=1,
model_name=None,
is_page_first_layout=True,
)
)
test_batch_operation(
HiCacheStorageConfig(
is_mla_model=True,
tp_rank=0,
tp_size=1,
model_name=None,
is_page_first_layout=True,
)
)
test_batch_operation(
HiCacheStorageConfig(
is_mla_model=False,
tp_rank=1,
tp_size=4,
model_name=None,
is_page_first_layout=True,
)
)
test_batch_operation(
HiCacheStorageConfig(
is_mla_model=True,
tp_rank=3,
tp_size=8,
model_name=None,
is_page_first_layout=True,
)
)
logger.info(f"✅ All tests passed")
@@ -0,0 +1,142 @@
import ctypes
import logging
import math
import os
from typing import Any, Dict
import torch
from sglang.srt.mem_cache.pool_host.common import HostTensorAllocator
logger = logging.getLogger(__name__)
def _bool_env(name: str, default: bool) -> bool:
raw = os.getenv(name)
if raw is None:
return default
return raw.strip().lower() in ("1", "true", "yes", "on")
def _int_env(name: str, default: int) -> int:
raw = os.getenv(name)
return int(raw) if raw is not None and raw != "" else default
class UMBPHostTensorAllocator(HostTensorAllocator):
"""Allocate the HiCache L2 host tensor from mori's UMBPHostMemAllocator."""
def __init__(self) -> None:
super().__init__()
try:
import mori.umbp as umbp_mod
except ImportError as exc:
raise RuntimeError(
"mori.umbp is not available. Build mori with BUILD_UMBP=ON "
"or fall back to the default torch host allocator."
) from exc
self._mod = umbp_mod
self._allocator = umbp_mod.UMBPHostMemAllocator()
self._use_hugepage = _bool_env("SGLANG_HICACHE_HOST_HUGEPAGE", True)
self._hugepage_size = _int_env(
"SGLANG_HICACHE_HOST_HUGEPAGE_SIZE", 2 * 1024 * 1024
)
self._numa_node = _int_env("SGLANG_HICACHE_HOST_NUMA_NODE", -1)
self._prefault = _bool_env("SGLANG_HICACHE_HOST_PREFAULT", True)
self._handles: Dict[int, Any] = {}
def allocate(
self, dims: tuple, dtype: torch.dtype, device: str = "cpu"
) -> torch.Tensor:
if device != "cpu":
raise ValueError(
"UMBPHostTensorAllocator only supports CPU host memory, "
f"got device={device}"
)
self.dims = dims
self.dtype = dtype
element_size = torch.empty((), dtype=dtype).element_size()
nbytes = math.prod(int(dim) for dim in dims) * element_size
requested_backing = (
self._mod.UMBPHostBufferBacking.AnonymousHugetlb
if self._use_hugepage
else self._mod.UMBPHostBufferBacking.Anonymous
)
handle = self._allocator.alloc(
nbytes,
requested_backing,
self._hugepage_size,
self._numa_node,
self._prefault,
)
if not handle:
raise RuntimeError(
f"UMBPHostMemAllocator.alloc({nbytes} bytes) failed "
f"(requested_backing={requested_backing}, "
f"numa_node={self._numa_node})."
)
self._handles[int(handle.ptr)] = handle
c_array = (ctypes.c_byte * nbytes).from_address(handle.ptr)
tensor = torch.frombuffer(c_array, dtype=torch.uint8, count=nbytes)
if dtype != torch.uint8:
tensor = tensor.view(dtype)
logger.info(
"UMBPHostTensorAllocator: allocated %.2f GB at 0x%x "
"requested_backing=%s actual_backing=%s actual_alignment=%d "
"mapped_size=%d numa_node=%d",
nbytes / 1e9,
handle.ptr,
requested_backing,
handle.actual_backing,
handle.actual_alignment,
handle.mapped_size,
self._numa_node,
)
if (
self._use_hugepage
and handle.actual_backing == self._mod.UMBPHostBufferBacking.Anonymous
):
logger.warning(
"UMBPHostTensorAllocator: requested AnonymousHugetlb backing "
"but kernel demoted to Anonymous (4 KiB pages). Check "
"vm.nr_hugepages and HugePages_Free in /proc/meminfo. "
"Performance and AINIC MR-size benefits will not apply."
)
return tensor.view(dims)
def mapped_size_for(self, ptr: int) -> int:
"""Actual mmap size for the allocation whose base address is *ptr*."""
handles = getattr(self, "_handles", None)
if handles is None:
return 0
h = handles.get(ptr)
return int(h.mapped_size) if h is not None else 0
@property
def mapped_size(self) -> int:
"""Largest mapped_size across all live allocations, or 0."""
handles = getattr(self, "_handles", None)
if not handles:
return 0
return max(int(h.mapped_size) for h in handles.values())
def __del__(self) -> None:
try:
handles = getattr(self, "_handles", None)
allocator = getattr(self, "_allocator", None)
if handles and allocator is not None:
for h in handles.values():
allocator.free(h)
self._handles.clear()
except Exception:
pass
File diff suppressed because it is too large Load Diff