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,716 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Any, List, Optional, Set
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.mem_cache.pool_host import HostKVCache
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Max pages per batched storage IO call.
|
||||
STORAGE_BATCH_SIZE = 128
|
||||
|
||||
|
||||
@dataclass
|
||||
class HiCacheStorageConfig:
|
||||
tp_rank: int
|
||||
tp_size: int
|
||||
pp_rank: int
|
||||
pp_size: int
|
||||
attn_cp_rank: int
|
||||
attn_cp_size: int
|
||||
is_mla_model: bool
|
||||
enable_storage_metrics: bool
|
||||
is_page_first_layout: bool
|
||||
model_name: Optional[str]
|
||||
tp_lcm_size: Optional[int] = None
|
||||
should_split_heads: bool = False
|
||||
extra_config: Optional[dict] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class HiCacheStorageExtraInfo:
|
||||
prefix_keys: Optional[List[str]] = None
|
||||
extra_info: Optional[dict] = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PrefetchTimeoutConfig:
|
||||
"""Knobs for the linear prefetch-timeout policy used by HiCache."""
|
||||
|
||||
base: float = 2.0 # seconds, fixed overhead unrelated to token count
|
||||
per_ki_token: float = 0.1 # seconds per 1024 tokens
|
||||
max: float = 30.0 # seconds, upper bound for the linear timeout
|
||||
|
||||
|
||||
class PoolName(str, Enum):
|
||||
"""Well-known pool names used as PoolTransfer/PoolEntry identifiers."""
|
||||
|
||||
KV = "kv"
|
||||
MAMBA = "mamba"
|
||||
SWA = "swa"
|
||||
INDEXER = "indexer"
|
||||
# TODO(hzh0425): Current DeepSeek V4 pool naming is verbose; will be normalized to
|
||||
# 'COMPRESSED_KV / COMPRESSED_INDEXER / COMPRESSED_STATE' in the next PR.
|
||||
DEEPSEEK_V4_C4 = "deepseek_v4_c4"
|
||||
DEEPSEEK_V4_C4_INDEXER = "deepseek_v4_c4_indexer"
|
||||
DEEPSEEK_V4_C128 = "deepseek_v4_c128"
|
||||
DEEPSEEK_V4_C4_STATE = "deepseek_v4_c4_state"
|
||||
DEEPSEEK_V4_C4_INDEXER_STATE = "deepseek_v4_c4_indexer_state"
|
||||
DEEPSEEK_V4_C128_STATE = "deepseek_v4_c128_state"
|
||||
|
||||
# Draft KV pool
|
||||
DRAFT = "draft"
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.value
|
||||
|
||||
|
||||
class PoolHitPolicy(str, Enum):
|
||||
"""Hit policy for batch_exists_v2 per-pool prefix matching.
|
||||
|
||||
ALL_PAGES : every page in [0, kv_hit) must exist (e.g. DSA).
|
||||
TRAILING_PAGES : only the last N pages must exist (e.g. Mamba/SWA states).
|
||||
"""
|
||||
|
||||
ALL_PAGES = "all_pages"
|
||||
TRAILING_PAGES = "trailing_pages"
|
||||
|
||||
|
||||
@dataclass
|
||||
class PoolTransfer:
|
||||
"""Unified per-pool transfer descriptor for batch v2 interface.
|
||||
|
||||
device<->host path : host_indices + device_indices
|
||||
host<->storage path: host_indices + keys
|
||||
nodes_to_load : evicted nodes this transfer covers
|
||||
"""
|
||||
|
||||
name: PoolName
|
||||
host_indices: Optional[torch.Tensor] = None
|
||||
device_indices: Optional[torch.Tensor] = None
|
||||
keys: Optional[List[str]] = None
|
||||
hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES
|
||||
nodes_to_load: Optional[List[Any]] = None
|
||||
indices_from_pool: Optional[PoolName] = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SidecarPoolSpec:
|
||||
"""Pool whose transfer indices are reused from one real source pool."""
|
||||
|
||||
pool_name: PoolName
|
||||
indices_from_pool: PoolName
|
||||
hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES
|
||||
|
||||
|
||||
@dataclass
|
||||
class PoolTransferResult:
|
||||
"""Tracks how many pages were successfully processed per pool."""
|
||||
|
||||
kv_hit_pages: int
|
||||
extra_pool_hit_pages: dict[str, int]
|
||||
|
||||
@classmethod
|
||||
def empty(cls) -> PoolTransferResult:
|
||||
return cls(0, {})
|
||||
|
||||
def update_kv_hit_pages(self, kv_hit_pages: int) -> None:
|
||||
"""Accumulate kv_hit_pages across batches (max = last successful batch)."""
|
||||
self.kv_hit_pages = max(self.kv_hit_pages, kv_hit_pages)
|
||||
|
||||
def update_extra_pool_hit_pages(self, results: dict[str, List[bool]]) -> None:
|
||||
"""Record actual load/write success counts per extra pool."""
|
||||
self.extra_pool_hit_pages.update(
|
||||
{name: sum(rs) for name, rs in results.items()}
|
||||
)
|
||||
|
||||
|
||||
class HiCacheStorage(ABC):
|
||||
"""
|
||||
HiCacheStorage is a class that provides a generic key-value interface for storing and retrieving KV cache.
|
||||
It abstracts the underlying storage mechanism, allowing different implementations to be used.
|
||||
"""
|
||||
|
||||
# todo, the page size of storage backend does not have to be the same as the same as host memory pool
|
||||
def register_mem_pool_host(self, mem_pool_host: HostKVCache):
|
||||
self.mem_pool_host = mem_pool_host
|
||||
|
||||
def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name):
|
||||
if not hasattr(self, "registered_pools"):
|
||||
self.registered_pools = {}
|
||||
self.registered_pools[host_pool_name] = host_pool
|
||||
|
||||
def batch_exists_v2(
|
||||
self,
|
||||
keys: List[str],
|
||||
pool_transfers: Optional[List[PoolTransfer]] = None,
|
||||
extra_info: Optional[HiCacheStorageExtraInfo] = None,
|
||||
) -> PoolTransferResult:
|
||||
"""Check which cache pages exist in storage, respecting per-pool hit policies.
|
||||
|
||||
Longest-prefix semantics
|
||||
Extra-pool hit policies (``PoolTransfer.hit_policy``)
|
||||
------------------------------------------------------
|
||||
Each ``PoolTransfer`` in ``pool_transfers`` describes a secondary
|
||||
cache pool (e.g. Mamba SSM states) that must be co-present with the
|
||||
KV pages. The final ``final_pages`` is the minimum across all pools,
|
||||
so a missing auxiliary page shrinks the usable prefix.
|
||||
|
||||
- ``"all_pages"`` (default): every page in [0, kv_hit) must exist
|
||||
for this pool. Used for pools that are required for every token
|
||||
in the prefix (e.g. DeepSeek DSA pool).
|
||||
|
||||
- ``"trailing_pages"``: only the *last* ``len(transfer.keys)`` pages
|
||||
of the KV prefix need to exist. Used for pools whose data covers
|
||||
only the tail of a prefix (e.g. Mamba/SWA Pool).
|
||||
|
||||
Returns
|
||||
-------
|
||||
PoolTransferResult
|
||||
``kv_hit_pages`` = length of the usable KV prefix.
|
||||
``extra_pool_hit_pages`` maps each pool name to the number of pages
|
||||
that were found.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def batch_get_v2(
|
||||
self,
|
||||
transfers: List[PoolTransfer],
|
||||
extra_info: Optional[HiCacheStorageExtraInfo] = None,
|
||||
) -> dict[str, List[bool]]:
|
||||
"""Read data from storage into host memory for each PoolTransfer.
|
||||
|
||||
Returns a dict mapping pool name to a per-entry success list.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def batch_set_v2(
|
||||
self,
|
||||
transfers: List[PoolTransfer],
|
||||
extra_info: Optional[HiCacheStorageExtraInfo] = None,
|
||||
) -> dict[str, List[bool]]:
|
||||
"""Write data from host memory to storage for each PoolTransfer.
|
||||
|
||||
Returns a dict mapping pool name to a per-entry success list.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def batch_get_v1(
|
||||
self,
|
||||
keys: List[str],
|
||||
host_indices: torch.Tensor,
|
||||
extra_info: Optional[HiCacheStorageExtraInfo] = None,
|
||||
) -> List[bool]:
|
||||
"""
|
||||
Retrieve values for multiple keys.
|
||||
Returns a list of booleans indicating success for each key.
|
||||
"""
|
||||
pass
|
||||
|
||||
def batch_set_v1(
|
||||
self,
|
||||
keys: List[str],
|
||||
host_indices: torch.Tensor,
|
||||
extra_info: Optional[HiCacheStorageExtraInfo] = None,
|
||||
) -> List[bool]:
|
||||
"""
|
||||
Store multiple key-value pairs.
|
||||
Returns a list of booleans indicating success for each key.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get(
|
||||
self,
|
||||
key: str,
|
||||
target_location: Optional[Any] = None,
|
||||
target_sizes: Optional[Any] = None,
|
||||
) -> torch.Tensor | None:
|
||||
"""
|
||||
Retrieve the value associated with the given key.
|
||||
Returns None if the key does not exist.
|
||||
"""
|
||||
pass
|
||||
|
||||
# TODO: Deprecate
|
||||
@abstractmethod
|
||||
def batch_get(
|
||||
self,
|
||||
keys: List[str],
|
||||
target_locations: Optional[Any] = None,
|
||||
target_sizes: Optional[Any] = None,
|
||||
) -> List[torch.Tensor | None] | int:
|
||||
"""
|
||||
Retrieve values for multiple keys.
|
||||
Returns a list of tensors or None for each key.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def set(
|
||||
self,
|
||||
key: str,
|
||||
value: Optional[Any] = None,
|
||||
target_location: Optional[Any] = None,
|
||||
target_sizes: Optional[Any] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Store the value associated with the given key.
|
||||
Returns True if the operation was successful, False otherwise.
|
||||
"""
|
||||
pass
|
||||
|
||||
# TODO: Deprecate
|
||||
@abstractmethod
|
||||
def batch_set(
|
||||
self,
|
||||
keys: List[str],
|
||||
values: Optional[Any] = None,
|
||||
target_locations: Optional[Any] = None,
|
||||
target_sizes: Optional[Any] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Store multiple key-value pairs.
|
||||
Returns True if all operations were successful, False otherwise.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def exists(self, key: str) -> bool:
|
||||
"""
|
||||
Check if the key exists in the storage.
|
||||
Returns True if the key exists, False otherwise.
|
||||
"""
|
||||
pass
|
||||
|
||||
# TODO: Use a finer-grained return type (e.g., List[bool])
|
||||
def batch_exists(
|
||||
self, keys: List[str], extra_info: Optional[HiCacheStorageExtraInfo] = None
|
||||
) -> int:
|
||||
"""
|
||||
Check if the keys exist in the storage.
|
||||
return the number of consecutive existing keys from the start.
|
||||
Can be overridden by subclasses for more efficient implementation.
|
||||
"""
|
||||
for i in range(len(keys)):
|
||||
if not self.exists(keys[i]):
|
||||
return i
|
||||
return len(keys)
|
||||
|
||||
def clear(self) -> None:
|
||||
pass
|
||||
|
||||
def get_stats(self):
|
||||
return None
|
||||
|
||||
|
||||
class MetadataCache:
|
||||
def __init__(self, ttl_seconds: float):
|
||||
self.ttl_seconds = ttl_seconds
|
||||
# key -> monotonic timestamp
|
||||
self.cache: dict[str, float] = {}
|
||||
self.lock = threading.Lock()
|
||||
|
||||
def add(self, key: str):
|
||||
with self.lock:
|
||||
if key not in self.cache:
|
||||
self.cache[key] = time.monotonic()
|
||||
|
||||
def remove(self, key: str):
|
||||
with self.lock:
|
||||
self.cache.pop(key, None)
|
||||
|
||||
def contains(self, key: str) -> bool:
|
||||
with self.lock:
|
||||
if key not in self.cache:
|
||||
return False
|
||||
if self.ttl_seconds == -1.0:
|
||||
return True
|
||||
if time.monotonic() - self.cache[key] > self.ttl_seconds:
|
||||
del self.cache[key]
|
||||
return False
|
||||
return True
|
||||
|
||||
def clear(self):
|
||||
with self.lock:
|
||||
self.cache.clear()
|
||||
|
||||
|
||||
class HiCacheFile(HiCacheStorage):
|
||||
|
||||
def __init__(
|
||||
self, storage_config: HiCacheStorageConfig, file_path: str = "/tmp/hicache"
|
||||
):
|
||||
self.file_path = envs.SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR.get() or file_path
|
||||
|
||||
tp_rank, tp_size, pp_rank, pp_size, model_name, is_mla_model = (
|
||||
storage_config.tp_rank,
|
||||
storage_config.tp_size,
|
||||
storage_config.pp_rank,
|
||||
storage_config.pp_size,
|
||||
storage_config.model_name,
|
||||
storage_config.is_mla_model,
|
||||
)
|
||||
attn_cp_rank = storage_config.attn_cp_rank
|
||||
attn_cp_size = storage_config.attn_cp_size
|
||||
model_name = "-".join(model_name.split("/")) if model_name else ""
|
||||
enable_pp = pp_size > 1
|
||||
self.config_suffix = f"_{model_name}"
|
||||
if not is_mla_model:
|
||||
self.config_suffix += f"_{tp_rank}_{tp_size}"
|
||||
if enable_pp:
|
||||
self.config_suffix += f"_{pp_size}_{pp_rank}"
|
||||
# Under NSA context parallel each CP rank holds a disjoint slice of every
|
||||
# page, so give each rank its own file key to avoid a cross-rank write race.
|
||||
if attn_cp_size > 1:
|
||||
self.config_suffix += f"_cp{attn_cp_rank}_{attn_cp_size}"
|
||||
|
||||
if not os.path.exists(self.file_path) and tp_rank == 0 and attn_cp_rank == 0:
|
||||
os.makedirs(self.file_path)
|
||||
logger.info(f"Created HiCacheFile storage directory at {self.file_path}")
|
||||
|
||||
# Metadata cache positive lookup toggle & TTL
|
||||
enable_cache_raw = None
|
||||
if storage_config.extra_config:
|
||||
enable_cache_raw = storage_config.extra_config.get("enable_metadata_cache")
|
||||
if enable_cache_raw is None:
|
||||
enable_cache_raw = (
|
||||
envs.SGLANG_HICACHE_FILE_BACKEND_ENABLE_METADATA_CACHE.get()
|
||||
)
|
||||
|
||||
self.enable_metadata_cache = bool(enable_cache_raw)
|
||||
|
||||
if self.enable_metadata_cache:
|
||||
ttl_raw = None
|
||||
if storage_config.extra_config:
|
||||
ttl_raw = storage_config.extra_config.get("metadata_ttl")
|
||||
if ttl_raw is None:
|
||||
ttl_raw = envs.SGLANG_HICACHE_FILE_BACKEND_METADATA_TTL.get()
|
||||
|
||||
self.metadata_ttl = float(ttl_raw) if ttl_raw is not None else 5.0
|
||||
self.metadata_cache = MetadataCache(self.metadata_ttl)
|
||||
self._scan_existing_files_to_metadata_cache()
|
||||
else:
|
||||
self.metadata_cache = None
|
||||
|
||||
# All LRU / size accounting and disk eviction lives in the evictor so
|
||||
# this backend stays a thin raw-bytes store. Imported lazily: the storage
|
||||
# package __init__ pulls in the backend factory, which imports this
|
||||
# module, so a top-level import here would be circular.
|
||||
from sglang.srt.mem_cache.storage.file.lru_file_evictor import LRUFileEvictor
|
||||
|
||||
self._evictor = LRUFileEvictor(
|
||||
self.file_path,
|
||||
self.config_suffix,
|
||||
tp_rank=tp_rank,
|
||||
is_mla_model=is_mla_model,
|
||||
extra_config=storage_config.extra_config,
|
||||
on_evict=(
|
||||
self.metadata_cache.remove if self.metadata_cache is not None else None
|
||||
),
|
||||
)
|
||||
|
||||
def _get_suffixed_key(self, key: str) -> str:
|
||||
return key + self.config_suffix
|
||||
|
||||
def _get_component_key(self, key: str, component_name: Optional[str] = None) -> str:
|
||||
if component_name is None or component_name in ("__default__", PoolName.KV):
|
||||
return self._get_suffixed_key(key)
|
||||
return self._get_suffixed_key(f"{key}.{component_name}")
|
||||
|
||||
def _get_component_path(
|
||||
self, key: str, component_name: Optional[str] = None
|
||||
) -> str:
|
||||
return os.path.join(
|
||||
self.file_path, f"{self._get_component_key(key, component_name)}.bin"
|
||||
)
|
||||
|
||||
def _scan_existing_files_to_metadata_cache(self) -> None:
|
||||
try:
|
||||
names = os.listdir(self.file_path)
|
||||
except FileNotFoundError:
|
||||
return
|
||||
for fn in names:
|
||||
if not fn.endswith(".bin"):
|
||||
continue
|
||||
stem = fn[:-4]
|
||||
# Only files belonging to this rank/model.
|
||||
if stem.endswith(self.config_suffix):
|
||||
self.metadata_cache.add(stem)
|
||||
|
||||
def get(
|
||||
self,
|
||||
key: str,
|
||||
target_location: torch.Tensor,
|
||||
target_sizes: Optional[Any] = None,
|
||||
) -> torch.Tensor | None:
|
||||
suffixed = self._get_suffixed_key(key)
|
||||
tensor_path = os.path.join(self.file_path, f"{suffixed}.bin")
|
||||
try:
|
||||
expected = target_location.numel() * target_location.element_size()
|
||||
with open(tensor_path, "rb", buffering=0) as f:
|
||||
buf = memoryview(target_location.view(torch.uint8).contiguous().numpy())
|
||||
if f.readinto(buf) != expected:
|
||||
raise IOError(f"Short read for {suffixed}")
|
||||
self._evictor.touch(suffixed, tensor_path)
|
||||
if self.metadata_cache is not None:
|
||||
self.metadata_cache.add(suffixed)
|
||||
return target_location
|
||||
except FileNotFoundError:
|
||||
if self.metadata_cache is not None:
|
||||
self.metadata_cache.remove(suffixed)
|
||||
logger.warning(f"Failed to fetch {key} from HiCacheFile storage.")
|
||||
return None
|
||||
|
||||
def batch_get(
|
||||
self,
|
||||
keys: List[str],
|
||||
target_locations: List[torch.Tensor],
|
||||
target_sizes: Optional[Any] = None,
|
||||
) -> List[torch.Tensor | None]:
|
||||
return [
|
||||
self.get(key, target_location)
|
||||
for key, target_location in zip(
|
||||
keys, target_locations or [None] * len(keys)
|
||||
)
|
||||
]
|
||||
|
||||
def set(
|
||||
self,
|
||||
key: str,
|
||||
value: Optional[Any] = None,
|
||||
target_location: Optional[Any] = None,
|
||||
target_sizes: Optional[Any] = None,
|
||||
) -> bool:
|
||||
suffixed = self._get_suffixed_key(key)
|
||||
tensor_path = os.path.join(self.file_path, f"{suffixed}.bin")
|
||||
|
||||
# Fast path: same key already on disk. Refresh recency and skip rewrite.
|
||||
if self.exists(key):
|
||||
logger.debug(f"Key {key} already exists. Skipped.")
|
||||
self._evictor.touch(suffixed, tensor_path)
|
||||
return True
|
||||
|
||||
tmp_path = None
|
||||
reserved = False
|
||||
try:
|
||||
value_bytes = value.numel() * value.element_size()
|
||||
# Ask the evictor to admit + reserve disk space (evicting if needed).
|
||||
if not self._evictor.reserve(suffixed, value_bytes, key=key):
|
||||
return False
|
||||
reserved = True
|
||||
|
||||
tmp_path = (
|
||||
f"{tensor_path}.tmp."
|
||||
f"{os.getpid()}.{threading.get_ident()}.{uuid.uuid4().hex}"
|
||||
)
|
||||
value.contiguous().view(dtype=torch.uint8).numpy().tofile(tmp_path)
|
||||
os.replace(tmp_path, tensor_path)
|
||||
self._evictor.commit(suffixed)
|
||||
if self.metadata_cache is not None:
|
||||
self.metadata_cache.add(suffixed)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to save tensor {key}: {e}")
|
||||
# Roll back the reservation and clean up any half-written file.
|
||||
if reserved:
|
||||
self._evictor.abort(suffixed)
|
||||
if tmp_path is not None:
|
||||
try:
|
||||
os.remove(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
if self.metadata_cache is not None:
|
||||
self.metadata_cache.remove(suffixed)
|
||||
return False
|
||||
|
||||
def batch_set(
|
||||
self,
|
||||
keys: List[str],
|
||||
values: Optional[Any] = None,
|
||||
target_locations: Optional[Any] = None,
|
||||
target_sizes: Optional[Any] = None,
|
||||
) -> bool:
|
||||
for key, value in zip(keys, values):
|
||||
if not self.set(key, value):
|
||||
return False
|
||||
return True
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
key = self._get_suffixed_key(key)
|
||||
if self.metadata_cache is not None and self.metadata_cache.contains(key):
|
||||
return True
|
||||
tensor_path = os.path.join(self.file_path, f"{key}.bin")
|
||||
if os.path.exists(tensor_path):
|
||||
if self.metadata_cache is not None:
|
||||
self.metadata_cache.add(key)
|
||||
return True
|
||||
return False
|
||||
|
||||
def _collect_existing_component_keys(
|
||||
self,
|
||||
keys: List[str],
|
||||
pool_transfers: Optional[List[PoolTransfer]] = None,
|
||||
) -> Set[str]:
|
||||
target_files = {f"{self._get_component_key(key)}.bin" for key in keys}
|
||||
for transfer in pool_transfers or []:
|
||||
for key in keys:
|
||||
target_files.add(f"{self._get_component_key(key, transfer.name)}.bin")
|
||||
|
||||
if self.metadata_cache is None:
|
||||
existing_files = set()
|
||||
with os.scandir(self.file_path) as entries:
|
||||
for entry in entries:
|
||||
if entry.is_file() and entry.name in target_files:
|
||||
existing_files.add(entry.name)
|
||||
return existing_files
|
||||
|
||||
existing_files = set()
|
||||
for filename in target_files:
|
||||
stem = filename[:-4]
|
||||
if self.metadata_cache.contains(stem):
|
||||
existing_files.add(filename)
|
||||
else:
|
||||
path = os.path.join(self.file_path, filename)
|
||||
if os.path.exists(path):
|
||||
self.metadata_cache.add(stem)
|
||||
existing_files.add(filename)
|
||||
return existing_files
|
||||
|
||||
def batch_exists_v2(
|
||||
self,
|
||||
keys: List[str],
|
||||
pool_transfers: Optional[List[PoolTransfer]] = None,
|
||||
extra_info: Optional[HiCacheStorageExtraInfo] = None,
|
||||
) -> PoolTransferResult:
|
||||
existing_files = self._collect_existing_component_keys(keys, pool_transfers)
|
||||
|
||||
def has_component(page_idx: int, name: str) -> bool:
|
||||
return (
|
||||
f"{self._get_component_key(keys[page_idx], name)}.bin" in existing_files
|
||||
)
|
||||
|
||||
# Longest contiguous KV prefix present in storage.
|
||||
kv_pages = next(
|
||||
(
|
||||
i
|
||||
for i in range(len(keys))
|
||||
if f"{self._get_component_key(keys[i])}.bin" not in existing_files
|
||||
),
|
||||
len(keys),
|
||||
)
|
||||
|
||||
hit_count: dict[str, int] = {PoolName.KV: kv_pages} if kv_pages else {}
|
||||
final_pages = kv_pages
|
||||
|
||||
for transfer in pool_transfers or []:
|
||||
if final_pages == 0:
|
||||
break
|
||||
name = transfer.name
|
||||
if transfer.hit_policy == PoolHitPolicy.ALL_PAGES:
|
||||
boundary = next(
|
||||
(i for i in range(kv_pages) if not has_component(i, name)), kv_pages
|
||||
)
|
||||
else: # trailing_pages
|
||||
trailing = max(1, len(transfer.keys) if transfer.keys else 1)
|
||||
boundary = 0
|
||||
for prefix_len in range(kv_pages, 0, -1):
|
||||
if all(
|
||||
has_component(i, name)
|
||||
for i in range(max(0, prefix_len - trailing), prefix_len)
|
||||
):
|
||||
boundary = prefix_len
|
||||
break
|
||||
if boundary:
|
||||
hit_count[name] = boundary
|
||||
final_pages = min(final_pages, boundary)
|
||||
|
||||
return PoolTransferResult(final_pages, hit_count)
|
||||
|
||||
def _log_key(self, pool_name: str, key: str) -> str:
|
||||
return key if pool_name == PoolName.KV else f"{key}.{pool_name}"
|
||||
|
||||
def _read_page(self, pool_name: str, key: str, host_pool, page_offset: int) -> bool:
|
||||
"""Read one page from storage into host_pool at page_offset."""
|
||||
storage_key = self._log_key(pool_name, key)
|
||||
data_page = self.get(storage_key, host_pool.get_dummy_flat_data_page())
|
||||
if data_page is None:
|
||||
return False
|
||||
host_pool.set_from_flat_data_page(page_offset, data_page)
|
||||
return True
|
||||
|
||||
def _write_page(
|
||||
self, pool_name: str, key: str, host_pool, page_offset: int
|
||||
) -> bool:
|
||||
"""Write one page from host_pool at page_offset to storage as raw bytes."""
|
||||
storage_key = self._log_key(pool_name, key)
|
||||
data_page = host_pool.get_data_page(page_offset, flat=True)
|
||||
return self.set(storage_key, data_page)
|
||||
|
||||
def _batch_io_v2(self, transfers: List[PoolTransfer], op_fn):
|
||||
results: dict[str, List[bool]] = {}
|
||||
for transfer in transfers:
|
||||
host_pool = self.registered_pools[transfer.name]
|
||||
keys = transfer.keys or []
|
||||
page_size = getattr(host_pool, "page_size", 1) or 1
|
||||
expected = len(keys) * page_size
|
||||
host_indices = transfer.host_indices
|
||||
|
||||
if host_indices is None or host_indices.numel() != expected:
|
||||
logger.error(
|
||||
"%s indices length mismatch for %s: expected %s, got %s",
|
||||
op_fn.__name__,
|
||||
transfer.name,
|
||||
expected,
|
||||
host_indices.numel() if host_indices is not None else 0,
|
||||
)
|
||||
results[transfer.name] = [False] * len(keys)
|
||||
continue
|
||||
|
||||
results[transfer.name] = [
|
||||
op_fn(transfer.name, key, host_pool, host_indices[i * page_size].item())
|
||||
for i, key in enumerate(keys)
|
||||
]
|
||||
return results
|
||||
|
||||
def batch_get_v2(
|
||||
self,
|
||||
transfers: List[PoolTransfer],
|
||||
extra_info: Optional[HiCacheStorageExtraInfo] = None,
|
||||
) -> dict[str, List[bool]]:
|
||||
return self._batch_io_v2(transfers, self._read_page)
|
||||
|
||||
def batch_set_v2(
|
||||
self,
|
||||
transfers: List[PoolTransfer],
|
||||
extra_info: Optional[HiCacheStorageExtraInfo] = None,
|
||||
) -> dict[str, List[bool]]:
|
||||
return self._batch_io_v2(transfers, self._write_page)
|
||||
|
||||
def clear(self) -> bool:
|
||||
try:
|
||||
for filename in os.listdir(self.file_path):
|
||||
file_path = os.path.join(self.file_path, filename)
|
||||
if os.path.isfile(file_path):
|
||||
os.remove(file_path)
|
||||
self._evictor.clear()
|
||||
if self.metadata_cache is not None:
|
||||
self.metadata_cache.clear()
|
||||
logger.info("Cleared all entries in HiCacheFile storage.")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to clear HiCacheFile storage: {e}")
|
||||
return False
|
||||
Reference in New Issue
Block a user