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
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:
@@ -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. ~30–46 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`. ~7–12 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 system’s 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
Reference in New Issue
Block a user