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

[UnifiedRadixTree]: Support L3 HiStorage framework (#26062)

Browse Source
This commit is contained in:
Zhangheng
2026-05-26 22:38:00 +08:00
committed by GitHub
parent c47f0e7cdd
commit 38f32c38ab
12 changed files with 1148 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/base_prefix_cache.py
+4
View File
@@ -75,6 +75,7 @@ class InsertResult:
prefix_len: int
mamba_exist: bool = False
inserted_host_node: Any = None
@dataclasses.dataclass
@@ -101,6 +102,7 @@ class IncLockRefResult:
delta: Optional[int] = None
swa_uuid_for_lock: Optional[int] = None
swa_uuid_for_host_lock: Optional[int] = None
# Component nodes that were tombstones at acquire time. Replaying this set
# at release prevents a short-lived lock from consuming a later load-back or
# request lock after that tombstone becomes a valid device value.
@@ -112,6 +114,7 @@ class IncLockRefResult:
"""Convert to the corresponding DecLockRefParams for dec_lock_ref."""
return DecLockRefParams(
swa_uuid_for_lock=self.swa_uuid_for_lock,
swa_uuid_for_host_lock=self.swa_uuid_for_host_lock,
skip_lock_node_ids={
component_type: set(node_ids)
for component_type, node_ids in self.skip_lock_node_ids.items()
@@ -124,6 +127,7 @@ class DecLockRefParams:
"""Parameters for dec_lock_ref operation."""
swa_uuid_for_lock: Optional[int] = None
swa_uuid_for_host_lock: Optional[int] = None
skip_lock_node_ids: dict[ComponentType, set[int]] = dataclasses.field(
default_factory=dict
)
python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py
+131
View File
@@ -1,8 +1,11 @@
from __future__ import annotations
import json
import logging
import os
import threading
import time
from queue import Queue
from typing import TYPE_CHECKING, Any, Callable, List, Optional
import torch
@@ -171,6 +174,7 @@ class HybridCacheController(BaseHiCacheController):
enable_storage_metrics: bool = False,
):
startup_storage_backend = storage_backend
self.extra_host_mem_release_queues: dict[PoolName, Queue] = {}
super().__init__(
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
mem_pool_host=mem_pool_host,
@@ -204,6 +208,10 @@ class HybridCacheController(BaseHiCacheController):
host_pools=getattr(mem_pool_host, "entries", None),
)
def _start_storage_threads(self):
super()._start_storage_threads()
self._init_extra_host_mem_release_queues()
def attach_storage_backend(
self,
storage_backend: str,
@@ -222,10 +230,133 @@ class HybridCacheController(BaseHiCacheController):
for entry in host_pools or []:
self.storage_backend.register_mem_host_pool_v2(entry.host_pool, entry.name)
@staticmethod
def parse_storage_backend_extra_config(
storage_backend_extra_config: Optional[str],
) -> tuple[dict, int, float, float, bool]:
extra_config = {}
if storage_backend_extra_config:
if storage_backend_extra_config.startswith("@"):
path = storage_backend_extra_config[1:]
ext = os.path.splitext(path)[1].lower()
with open(path, "rb" if ext == ".toml" else "r") as f:
if ext == ".json":
extra_config = json.load(f)
elif ext == ".toml":
import tomllib
extra_config = tomllib.load(f)
elif ext in (".yaml", ".yml"):
import yaml
extra_config = yaml.safe_load(f)
else:
raise ValueError(
f"Unsupported config file {path} (config format: {ext})"
)
else:
extra_config = json.loads(storage_backend_extra_config)
prefetch_threshold = extra_config.pop("prefetch_threshold", 256)
prefetch_timeout_base = extra_config.pop("prefetch_timeout_base", 1)
prefetch_timeout_per_ki_token = extra_config.pop(
"prefetch_timeout_per_ki_token", 0.25
)
hicache_storage_pass_prefix_keys = extra_config.pop(
"hicache_storage_pass_prefix_keys", False
)
if not isinstance(prefetch_threshold, int):
raise ValueError(
f"prefetch_threshold must be int, got {type(prefetch_threshold).__name__}"
)
if not isinstance(prefetch_timeout_base, (int, float)):
raise ValueError(
f"prefetch_timeout_base must be number, got {type(prefetch_timeout_base).__name__}"
)
if not isinstance(prefetch_timeout_per_ki_token, (int, float)):
raise ValueError(
"prefetch_timeout_per_ki_token must be number, got "
f"{type(prefetch_timeout_per_ki_token).__name__}"
)
if not isinstance(hicache_storage_pass_prefix_keys, bool):
raise ValueError(
"hicache_storage_pass_prefix_keys must be bool, got "
f"{type(hicache_storage_pass_prefix_keys).__name__}"
)
return (
extra_config,
prefetch_threshold,
float(prefetch_timeout_base),
float(prefetch_timeout_per_ki_token),
hicache_storage_pass_prefix_keys,
)
def clear_storage_backend(self) -> bool:
if not self.enable_storage:
logger.warning("Hierarchical cache storage backend is not enabled.")
return False
if not hasattr(self.storage_backend, "clear"):
logger.warning(
"Storage backend %s does not support clear operation.",
type(self.storage_backend).__name__,
)
return False
self.storage_backend.clear()
return True
def _init_extra_host_mem_release_queues(self) -> None:
self.extra_host_mem_release_queues = {}
entries = getattr(self.mem_pool_host, "entries", None) or []
anchor_entry = getattr(self.mem_pool_host, "anchor_entry", None)
for entry in entries:
if entry is anchor_entry or entry.is_primary_index_anchor:
continue
self.extra_host_mem_release_queues[entry.name] = Queue()
def _append_host_mem_release_pages(
self, release_queue: Queue, host_indices: torch.Tensor, page_size: int
) -> None:
if host_indices.numel() == 0:
return
for page in host_indices.split(page_size):
release_queue.put(page)
def append_host_mem_release(
self,
host_indices: Optional[torch.Tensor] = None,
extra_pools: Optional[list[PoolTransfer]] = None,
):
if host_indices is not None:
self._append_host_mem_release_pages(
self.host_mem_release_queue,
host_indices,
self.mem_pool_host.page_size,
)
for transfer in extra_pools or []:
if transfer.host_indices is None or transfer.host_indices.numel() == 0:
continue
entry = self.mem_pool_host.entry_map.get(transfer.name)
if (
entry is None
or entry.is_primary_index_anchor
or transfer.indices_from_pool is not None
):
continue
release_queue = self.extra_host_mem_release_queues.get(transfer.name)
if release_queue is None:
continue
self._append_host_mem_release_pages(
release_queue, transfer.host_indices, entry.host_pool.page_size
)
def reset(self):
super().reset()
if self.enable_storage:
self.host_mem_release_queue.queue.clear()
for release_queue in self.extra_host_mem_release_queues.values():
release_queue.queue.clear()
self.prefetch_tokens_occupied = 0
def write(
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
+63 -5
View File
@@ -679,6 +679,11 @@ class StackStrategy:
load_cache_event,
attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
storage_backend: Optional[str] = None,
storage_backend_extra_config: Optional[dict] = None,
prefetch_threshold: int = 256,
model_name: Optional[str] = None,
enable_storage_metrics: bool = False,
) -> StackBuildResult:
raise NotImplementedError
@@ -704,6 +709,11 @@ class _DeepSeekV4Strategy(StackStrategy):
load_cache_event,
attn_cp_group=None,
attn_tp_group=None,
storage_backend=None,
storage_backend_extra_config=None,
prefetch_threshold=256,
model_name=None,
enable_storage_metrics=False,
):
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
@@ -716,11 +726,15 @@ class _DeepSeekV4Strategy(StackStrategy):
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
storage_backend=None,
storage_backend=storage_backend,
host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA),
device_swa_evict_fn=lambda n: cache.evict(EvictParams(swa_num_tokens=n)),
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
pp_rank=params.pp_rank,
pp_size=params.pp_size,
enable_storage_metrics=enable_storage_metrics,
)
sidecars = [
SidecarPoolSpec(pool_name=name, indices_from_pool=src)
@@ -766,6 +780,11 @@ class _MambaStrategy(StackStrategy):
load_cache_event,
attn_cp_group=None,
attn_tp_group=None,
storage_backend=None,
storage_backend_extra_config=None,
prefetch_threshold=256,
model_name=None,
enable_storage_metrics=False,
):
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
@@ -783,12 +802,16 @@ class _MambaStrategy(StackStrategy):
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
storage_backend=None,
storage_backend=storage_backend,
use_mla=kvcache.use_mla,
host_mamba_evict_fn=lambda n: cache.evict_host(n, ComponentType.MAMBA),
device_mamba_evict_fn=lambda n: cache.evict(EvictParams(mamba_num=n)),
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
pp_rank=params.pp_rank,
pp_size=params.pp_size,
enable_storage_metrics=enable_storage_metrics,
)
return StackBuildResult(
host_pool_group=host_pool_group,
@@ -834,6 +857,11 @@ class _SwaStrategy(StackStrategy):
load_cache_event,
attn_cp_group=None,
attn_tp_group=None,
storage_backend=None,
storage_backend_extra_config=None,
prefetch_threshold=256,
model_name=None,
enable_storage_metrics=False,
):
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
@@ -850,12 +878,16 @@ class _SwaStrategy(StackStrategy):
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
storage_backend=None,
storage_backend=storage_backend,
use_mla=False,
host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA),
device_swa_evict_fn=lambda n: cache.evict(EvictParams(swa_num_tokens=n)),
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
pp_rank=params.pp_rank,
pp_size=params.pp_size,
enable_storage_metrics=enable_storage_metrics,
)
return StackBuildResult(
host_pool_group=host_pool_group,
@@ -887,6 +919,11 @@ class _DsaStrategy(StackStrategy):
load_cache_event,
attn_cp_group=None,
attn_tp_group=None,
storage_backend=None,
storage_backend_extra_config=None,
prefetch_threshold=256,
model_name=None,
enable_storage_metrics=False,
):
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
@@ -904,7 +941,7 @@ class _DsaStrategy(StackStrategy):
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
storage_backend=None,
storage_backend=storage_backend,
use_mla=use_mla,
override_kv_cache_dim=full_kv_pool.kv_cache_dim,
sidecar_host_pool_factory=lambda kv_host_pool: DSAIndexerPoolHost(
@@ -913,8 +950,12 @@ class _DsaStrategy(StackStrategy):
server_args.hicache_mem_layout,
allocator_type=server_args.hicache_storage_backend,
),
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
pp_rank=params.pp_rank,
pp_size=params.pp_size,
enable_storage_metrics=enable_storage_metrics,
)
return StackBuildResult(
host_pool_group=host_pool_group,
@@ -961,6 +1002,11 @@ class _PlainKvStrategy(StackStrategy):
load_cache_event,
attn_cp_group=None,
attn_tp_group=None,
storage_backend=None,
storage_backend_extra_config=None,
prefetch_threshold=256,
model_name=None,
enable_storage_metrics=False,
):
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
@@ -977,10 +1023,14 @@ class _PlainKvStrategy(StackStrategy):
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
storage_backend=None,
storage_backend=storage_backend,
use_mla=use_mla,
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
pp_rank=params.pp_rank,
pp_size=params.pp_size,
enable_storage_metrics=enable_storage_metrics,
)
return StackBuildResult(
host_pool_group=host_pool_group,
@@ -1057,6 +1107,9 @@ def attach_hybrid_pool_to_unified_cache(
load_cache_event,
attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
storage_backend: Optional[str] = None,
storage_extra_config: Optional[dict] = None,
storage_prefetch_threshold: int = 256,
) -> None:
"""Attach HostPoolGroup + HybridCacheController to UnifiedRadixCache."""
try:
@@ -1071,6 +1124,11 @@ def attach_hybrid_pool_to_unified_cache(
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
storage_backend=storage_backend,
storage_backend_extra_config=storage_extra_config,
prefetch_threshold=storage_prefetch_threshold,
model_name=server_args.served_model_name,
enable_storage_metrics=cache._enable_metrics_flag,
)
_apply_stack_result(cache, kvcache, params, result)
except Exception:
python/sglang/srt/mem_cache/unified_cache_components/full_component.py
+28 -3
View File
@@ -65,7 +65,7 @@ class FullComponent(TreeComponent):
# last_device_node, summing host_value lengths of evicted nodes.
ct = self.component_type
kv_host_hit = 0
node = result.last_host_node
node = result.best_match_node
root_node = self.cache.root_node
while node is not result.last_device_node and node is not root_node:
full_host = node.component_data[ct].host_value
@@ -155,9 +155,22 @@ class FullComponent(TreeComponent):
heapq.heappush(heap, (x.parent.last_access_time, x.parent))
def acquire_component_lock(
self, node: UnifiedTreeNode, result: IncLockRefResult
self,
node: UnifiedTreeNode,
result: IncLockRefResult,
lock_host: bool = False,
) -> IncLockRefResult:
ct = self.component_type
# Only the last host node needs to be protected.
if lock_host:
cd = node.component_data[ct]
if cd.host_value is None:
return result
cd.host_lock_ref += 1
self.cache._update_evictable_leaf_sets(node)
return result
root = self.cache.root_node
cur = node
@@ -185,9 +198,20 @@ class FullComponent(TreeComponent):
return result
def release_component_lock(
self, node: UnifiedTreeNode, params: Optional[DecLockRefParams]
self,
node: UnifiedTreeNode,
params: Optional[DecLockRefParams],
lock_host: bool = False,
) -> None:
ct = self.component_type
if lock_host:
cd = node.component_data[ct]
if cd.host_value is None or cd.host_lock_ref == 0:
return
cd.host_lock_ref -= 1
self.cache._update_evictable_leaf_sets(node)
return
root = self.cache.root_node
skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
cur = node
@@ -255,6 +279,7 @@ class FullComponent(TreeComponent):
node: UnifiedTreeNode,
phase: CacheTransferPhase,
transfers: list[PoolTransfer] = (),
**kw,
) -> None:
ct = self.component_type
python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py
+102 -11
View File
@@ -13,7 +13,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchPrefixParams,
MatchResult,
)
from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
CacheTransferPhase,
ComponentType,
@@ -213,34 +213,59 @@ class MambaComponent(TreeComponent):
x = x_next
def acquire_component_lock(
self, node: UnifiedTreeNode, result: IncLockRefResult
self,
node: UnifiedTreeNode,
result: IncLockRefResult,
lock_host: bool = False,
) -> IncLockRefResult:
ct = self.component_type
if node is self.cache.root_node:
return result
cd = node.component_data[ct]
value = cd.value
value = cd.host_value if lock_host else cd.value
# A node in skip_lock_node_ids was a tombstone when this lock was acquired.
if value is None:
result.skip_lock_node_ids.setdefault(ct, set()).add(node.id)
return result
if cd.lock_ref == 0:
vlen = len(value)
self.cache.component_evictable_size_[ct] -= vlen
self.cache.component_protected_size_[ct] += vlen
cd.lock_ref += 1
if lock_host:
if cd.host_lock_ref == 0:
host_lru = self.cache.host_lru_lists[ct]
if host_lru.in_list(node):
host_lru.remove_node(node)
cd.host_lock_ref += 1
else:
if cd.lock_ref == 0:
vlen = len(value)
self.cache.component_evictable_size_[ct] -= vlen
self.cache.component_protected_size_[ct] += vlen
cd.lock_ref += 1
return result
def release_component_lock(
self, node: UnifiedTreeNode, params: Optional[DecLockRefParams]
self,
node: UnifiedTreeNode,
params: Optional[DecLockRefParams],
lock_host: bool = False,
) -> None:
ct = self.component_type
if node is self.cache.root_node:
return
cd = node.component_data[ct]
skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
if node.id in skip_lock_node_ids:
return
value = cd.value
if value is not None and cd.lock_ref > 0:
value = cd.host_value if lock_host else cd.value
if lock_host:
cd.host_lock_ref -= 1
if cd.host_lock_ref == 0 and cd.value is None and cd.host_value is not None:
host_lru = self.cache.host_lru_lists[ct]
if not host_lru.in_list(node):
host_lru.insert_mru(node)
return
if cd.lock_ref > 0:
if cd.lock_ref == 1:
vlen = len(value)
self.cache.component_evictable_size_[ct] += vlen
@@ -392,6 +417,35 @@ class MambaComponent(TreeComponent):
return transfers if transfers else None
if phase == CacheTransferPhase.BACKUP_STORAGE:
cd = node.component_data[ct]
if cd.host_value is None or not node.hash_value:
return None
return [
PoolTransfer(
name=PoolName.MAMBA,
host_indices=cd.host_value,
keys=[node.hash_value[-1]],
hit_policy=PoolHitPolicy.TRAILING_PAGES,
)
]
if phase == CacheTransferPhase.PREFETCH:
host_indices = self._mamba_pool_host.alloc(1)
if host_indices is None:
self.cache.evict_host(1, ComponentType.MAMBA)
host_indices = self._mamba_pool_host.alloc(1)
if host_indices is None:
return []
return [
PoolTransfer(
name=PoolName.MAMBA,
host_indices=host_indices,
keys=["__placeholder__"],
hit_policy=PoolHitPolicy.TRAILING_PAGES,
)
]
return None
def commit_hicache_transfer(
@@ -399,6 +453,7 @@ class MambaComponent(TreeComponent):
node: UnifiedTreeNode,
phase: CacheTransferPhase,
transfers: list[PoolTransfer] = (),
**kw,
) -> None:
ct = self.component_type
@@ -423,6 +478,41 @@ class MambaComponent(TreeComponent):
self.cache.lru_lists[ct].insert_mru(node)
self.cache.component_evictable_size_[ct] += count
elif phase == CacheTransferPhase.PREFETCH:
if not transfers:
return
transfer = transfers[0]
host_indices = transfer.host_indices
insert_result = kw.get("insert_result")
pool_storage_result = kw.get("pool_storage_result")
loaded = (
pool_storage_result is not None
and pool_storage_result.extra_pool_hit_pages.get(PoolName.MAMBA, 0) >= 1
)
target_node = (
insert_result.inserted_host_node if insert_result is not None else None
)
if (
host_indices is None
or target_node is None
or not loaded
or target_node.component_data[ct].host_value is not None
):
self.cache.cache_controller.append_host_mem_release(
extra_pools=[transfer]
)
if insert_result is not None:
insert_result.mamba_exist = True
return
target_node.component_data[ct].host_value = host_indices.clone()
if target_node.component_data[ct].value is None:
host_lru = self.cache.host_lru_lists[ct]
if not host_lru.in_list(target_node):
host_lru.insert_mru(target_node)
if insert_result is not None:
insert_result.mamba_exist = False
def drive_host_eviction(
self, num_tokens: int, tracker: dict[ComponentType, int]
) -> None:
@@ -445,4 +535,5 @@ class MambaComponent(TreeComponent):
x, self, target=EvictLayer.HOST, tracker=tracker
)
self.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST)
self.cache._update_evictable_leaf_sets(x)
x = x_next
python/sglang/srt/mem_cache/unified_cache_components/swa_component.py
+9 -2
View File
@@ -350,7 +350,10 @@ class SWAComponent(TreeComponent):
x = x_next
def acquire_component_lock(
self, node: UnifiedTreeNode, result: IncLockRefResult
self,
node: UnifiedTreeNode,
result: IncLockRefResult,
lock_host: bool = False,
) -> IncLockRefResult:
ct = self.component_type
root = self.cache.root_node
@@ -384,7 +387,10 @@ class SWAComponent(TreeComponent):
return result
def release_component_lock(
self, node: UnifiedTreeNode, params: Optional[DecLockRefParams]
self,
node: UnifiedTreeNode,
params: Optional[DecLockRefParams],
lock_host: bool = False,
) -> None:
ct = self.component_type
root = self.cache.root_node
@@ -484,6 +490,7 @@ class SWAComponent(TreeComponent):
node: UnifiedTreeNode,
phase: CacheTransferPhase,
transfers: list[PoolTransfer] = (),
**kw,
) -> None:
ct = self.component_type
python/sglang/srt/mem_cache/unified_cache_components/tree_component.py
+20 -6
View File
@@ -276,9 +276,12 @@ class TreeComponent(ABC):
@abstractmethod
def acquire_component_lock(
self, node: UnifiedTreeNode, result: IncLockRefResult
self,
node: UnifiedTreeNode,
result: IncLockRefResult,
lock_host: bool = False,
) -> IncLockRefResult:
"""Increment lock_ref for this component, protecting nodes from
"""Increment component lock refs, protecting nodes from
eviction. Updates evictable → protected size on first lock.
- Full: path-lock — walks from node up to root, incrementing
lock_ref on every ancestor.
@@ -286,21 +289,31 @@ class TreeComponent(ABC):
sliding window is filled; records a component_uuid at the
boundary for release_component_lock to know where to stop.
- Mamba: single-node lock — only increments lock_ref on the
node itself (mamba state is per-leaf, not per-path)."""
node itself (mamba state is per-leaf, not per-path).
When ``lock_host`` is True, the lock applies to host-side state:
- Full: single-node host lock.
- SWA: host window-lock with a dedicated host UUID boundary.
- Mamba: single-node host lock with host LRU detach."""
...
@abstractmethod
def release_component_lock(
self, node: UnifiedTreeNode, params: Optional[DecLockRefParams]
self,
node: UnifiedTreeNode,
params: Optional[DecLockRefParams],
lock_host: bool = False,
) -> None:
"""Decrement lock_ref for this component, un-protecting nodes.
"""Decrement component lock refs, un-protecting nodes.
Updates protected → evictable size when lock_ref drops to 0.
- Full: path-unlock — walks from node up to root, decrementing
lock_ref on every ancestor.
- SWA: path-unlock — walks upward, stopping at the node whose
component_uuid matches the one recorded during acquire.
- Mamba: single-node unlock — only decrements lock_ref on the
node itself."""
node itself.
When ``lock_host`` is True, the inverse host-side semantics apply."""
...
def prepare_for_caching_req(
@@ -351,6 +364,7 @@ class TreeComponent(ABC):
node: UnifiedTreeNode,
phase: CacheTransferPhase,
transfers: list[PoolTransfer] = (),
**kw,
) -> None:
"""Post-transfer bookkeeping: store host indices, update LRU, etc."""
pass
python/sglang/srt/mem_cache/unified_radix_cache.py
+671 -19
View File
@@ -5,7 +5,8 @@ import threading
import time
from array import array
from collections import defaultdict
from functools import partial
from functools import lru_cache, partial
from queue import Empty
from typing import TYPE_CHECKING, Any, Optional
import torch
@@ -28,6 +29,9 @@ from sglang.srt.mem_cache.hicache_storage import (
PoolTransfer,
SidecarPoolSpec,
)
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
HybridCacheController,
)
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.unified_cache_components import (
_NUM_COMPONENT_TYPES,
@@ -42,6 +46,8 @@ from sglang.srt.mem_cache.unified_cache_components import (
TreeComponent,
get_and_increase_time_counter,
)
from sglang.srt.mem_cache.utils import compute_node_hash_values, split_node_hash_value
from sglang.srt.observability.metrics_collector import StorageMetricsCollector
from sglang.srt.session.streaming_session import StreamingSession
if TYPE_CHECKING:
@@ -93,6 +99,18 @@ class UnifiedTreeNode:
def __lt__(self, other: UnifiedTreeNode):
return self.last_access_time < other.last_access_time
def get_last_hash_value(self) -> Optional[str]:
if self.hash_value is None or len(self.hash_value) == 0:
return None
return self.hash_value[-1]
@lru_cache(maxsize=1)
def get_prefix_hash_values(self, node: UnifiedTreeNode) -> list[str]:
if node is None or node.hash_value is None:
return []
return node.get_prefix_hash_values(node.parent) + node.hash_value
class UnifiedLRUList:
def __init__(
@@ -220,6 +238,10 @@ class UnifiedRadixCache(BasePrefixCache):
if params.enable_metrics:
self.init_metrics_collector()
self._enable_metrics_flag = params.enable_metrics
self.enable_storage_metrics = False
self.storage_metrics_collector: Optional[StorageMetricsCollector] = None
self.extra_metric_labels = None
assert params.tree_components is not None
self.tree_components = tuple(params.tree_components)
@@ -248,6 +270,11 @@ class UnifiedRadixCache(BasePrefixCache):
# HiCache D↔H defaults (overridden by init_hicache)
self.cache_controller = None
self.write_through_threshold = 256
self.prefetch_stop_policy = "best_effort"
self.prefetch_threshold = 256
self.prefetch_timeout_base = 1.0
self.prefetch_timeout_per_page = 0.25
self.hicache_storage_pass_prefix_keys = False
self.reset()
logger.info(f"Init Unified RadixTree with components {self.tree_components}")
@@ -260,6 +287,7 @@ class UnifiedRadixCache(BasePrefixCache):
self.root_node = UnifiedTreeNode(self.tree_components)
self.root_node.key = RadixKey(array("q"), None)
self.root_node.component_data[BASE_COMPONENT_TYPE].value = []
self.root_node.hash_value = []
for ct in self.tree_components:
self.root_node.component_data[ct].lock_ref = 1
self.component_evictable_size_ = {ct: 0 for ct in self.tree_components}
@@ -281,12 +309,14 @@ class UnifiedRadixCache(BasePrefixCache):
] = {}
self.ongoing_load_back: dict[int, tuple[UnifiedTreeNode, DecLockRefParams]] = {}
self.enable_storage = False
self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {}
self.ongoing_prefetch: dict = {}
self.ongoing_backup: dict = {}
if self.cache_controller is not None:
self.cache_controller.reset()
self.cache_controller.mem_pool_host.clear()
self.enable_storage = self.cache_controller.enable_storage
self._empty_match_result = MatchResult(
device_indices=torch.empty(
@@ -316,6 +346,26 @@ class UnifiedRadixCache(BasePrefixCache):
self.load_cache_event = threading.Event()
self.sidecar_pool_specs.clear()
self.extra_metric_labels = server_args.extra_metric_labels
# Parse storage config once, share with assembler and tree
storage_backend = server_args.hicache_storage_backend
storage_extra_config = None
storage_prefetch_threshold = 256
prefetch_timeout_base = 1.0
prefetch_timeout_per_ki_token = 0.25
hicache_storage_pass_prefix_keys = False
if storage_backend is not None:
(
storage_extra_config,
storage_prefetch_threshold,
prefetch_timeout_base,
prefetch_timeout_per_ki_token,
hicache_storage_pass_prefix_keys,
) = HybridCacheController.parse_storage_backend_extra_config(
server_args.hicache_storage_backend_extra_config
)
attach_hybrid_pool_to_unified_cache(
self,
params,
@@ -323,6 +373,9 @@ class UnifiedRadixCache(BasePrefixCache):
load_cache_event=self.load_cache_event,
attn_cp_group=params.attn_cp_cache_group,
attn_tp_group=params.attn_tp_cache_group,
storage_backend=storage_backend,
storage_extra_config=storage_extra_config,
storage_prefetch_threshold=storage_prefetch_threshold,
)
# State initialization
@@ -330,14 +383,19 @@ class UnifiedRadixCache(BasePrefixCache):
1 if server_args.hicache_write_policy == "write_through" else 2
)
self.load_back_threshold = 256
self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy
logger.info(
f"HiCache D\u2194H initialized: "
f"host_pool_size={self.host_pool_group.size}, "
f"write_policy={server_args.hicache_write_policy}, "
f"tp_world_size={self.tp_world_size}, "
f"transfer_layer_num={self.cache_controller.layer_num}"
)
if storage_backend is not None:
self._apply_storage_runtime_config(
storage_backend=storage_backend,
prefetch_threshold=storage_prefetch_threshold,
prefetch_timeout_base=prefetch_timeout_base,
prefetch_timeout_per_ki_token=prefetch_timeout_per_ki_token,
hicache_storage_pass_prefix_keys=hicache_storage_pass_prefix_keys,
enable_storage=self.cache_controller.enable_storage,
enable_storage_metrics=self._enable_metrics_flag,
extra_metric_labels=self.extra_metric_labels,
)
def register_sidecar_pool(self, spec: SidecarPoolSpec) -> None:
self.sidecar_pool_specs.append(spec)
@@ -435,6 +493,29 @@ class UnifiedRadixCache(BasePrefixCache):
# TODO: delta is not aggregated from components; no caller uses it yet.
return DecLockRefResult()
def inc_host_lock_ref(self, node: Any) -> IncLockRefResult:
if self.disable:
return IncLockRefResult()
result = IncLockRefResult()
for component in self._components_tuple:
result = component.acquire_component_lock(
node=node, result=result, lock_host=True
)
self._update_evictable_leaf_sets(node)
return result
def dec_host_lock_ref(
self, node: Any, params: Optional[DecLockRefParams] = None
) -> DecLockRefResult:
if self.disable:
return DecLockRefResult()
for component in self._components_tuple:
component.release_component_lock(node=node, params=params, lock_host=True)
self._update_evictable_leaf_sets(node)
return DecLockRefResult()
def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs) -> None:
if self.session.try_cache_finished_req(req, is_insert=is_insert, **kwargs):
return
@@ -703,13 +784,15 @@ class UnifiedRadixCache(BasePrefixCache):
cur_time -= 0.00001
node_update = node_update.parent
# Walk up to find last_host_node for full component.
if self.cache_controller is None:
last_host_node = best_match_device_node
else:
last_host_node = best_match_node
while last_host_node is not self.root_node and not last_host_node.backuped:
last_host_node = last_host_node.parent
# last_host_node will be used as the starting node for the subsequent
# `prefetch_from_storage` flow. We directly use best_match_node here,
# because best_match_node represents the node where all components
# have reached consensus on both device & host availability.
last_host_node = (
best_match_node
if self.cache_controller is not None
else best_match_device_node
)
if best_match_device_value_len > 0:
device_indices = torch.cat(value[:best_match_device_value_len])
@@ -744,6 +827,9 @@ class UnifiedRadixCache(BasePrefixCache):
child.parent = new_node
child.key = child.key[split_len:]
new_node.hash_value, child.hash_value = split_node_hash_value(
child.hash_value, split_len, self.page_size
)
for component in self._components_tuple:
component.redistribute_on_node_split(new_parent=new_node, child=child)
@@ -778,6 +864,8 @@ class UnifiedRadixCache(BasePrefixCache):
new_node.component_data[BASE_COMPONENT_TYPE].value = value.clone()
parent.children[key.child_key(self.page_size)] = new_node
self.component_evictable_size_[BASE_COMPONENT_TYPE] += len(value)
if self.enable_storage:
new_node.hash_value = compute_node_hash_values(new_node, self.page_size)
self._update_evictable_leaf_sets(new_node)
self._update_evictable_leaf_sets(parent)
@@ -894,6 +982,58 @@ class UnifiedRadixCache(BasePrefixCache):
self._inc_hit_count(target_node, params.chunked)
return result
def _insert_helper_host(
self,
node: UnifiedTreeNode,
key: RadixKey,
host_value: torch.Tensor,
hash_value: list[str],
) -> InsertResult:
total_len = len(key)
self._touch_node(node)
if total_len == 0:
return InsertResult(prefix_len=0, mamba_exist=True)
child_key = key.child_key(self.page_size)
matched_length = 0
while len(key) > 0 and child_key in node.children:
node = node.children[child_key]
self._touch_node(node)
prefix_len = node.key.match(key, page_size=self.page_size)
key = key[prefix_len:]
host_value = host_value[prefix_len:]
hash_value = hash_value[prefix_len // self.page_size :]
matched_length += prefix_len
if prefix_len < len(node.key):
node = self._split_node(node.key, node, prefix_len)
if len(key):
child_key = key.child_key(self.page_size)
result = InsertResult(
prefix_len=matched_length,
)
if len(key) == 0:
if (
node is not self.root_node
and node.component_data[BASE_COMPONENT_TYPE].host_value is not None
):
result.inserted_host_node = node
return result
new_node = UnifiedTreeNode(self.tree_components)
new_node.parent = node
new_node.key = key
new_node.hash_value = hash_value
new_node.component_data[BASE_COMPONENT_TYPE].host_value = host_value.clone()
node.children[child_key] = new_node
self._update_evictable_leaf_sets(new_node)
self._update_evictable_leaf_sets(node)
result.inserted_host_node = new_node
return result
# ---- Evict Helpers ----
def _cascade_evict(
@@ -1126,6 +1266,7 @@ class UnifiedRadixCache(BasePrefixCache):
and self.cache_controller.write_policy == "write_back"
):
self.write_backup(node, write_back=True)
self.writing_check(write_back=True)
self._evict_to_host(node, tracker)
return
else:
@@ -1368,6 +1509,511 @@ class UnifiedRadixCache(BasePrefixCache):
if not node.backuped and node.hit_count >= self.write_through_threshold:
self.write_backup(node)
def write_backup_storage(self, node: UnifiedTreeNode) -> None:
if (
not self.enable_storage
or self.cache_controller is None
or not node.backuped
):
return
prefix_keys = None
if self.hicache_storage_pass_prefix_keys:
prefix_keys = node.get_prefix_hash_values(node.parent)
comp_xfers: dict[ComponentType, list[PoolTransfer]] = {}
for comp in self._components_tuple:
if comp.component_type == BASE_COMPONENT_TYPE:
continue
transfers = comp.build_hicache_transfers(
node,
CacheTransferPhase.BACKUP_STORAGE,
)
if transfers:
comp_xfers[comp.component_type] = transfers
kv_xfer = PoolTransfer(
name=PoolName.KV,
host_indices=node.component_data[BASE_COMPONENT_TYPE].host_value,
keys=node.hash_value,
)
sidecar_xfers = self._build_sidecar_transfers(
CacheTransferPhase.BACKUP_STORAGE, kv_xfer, comp_xfers
)
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
aux_xfers.extend(sidecar_xfers)
operation_id = self.cache_controller.write_storage(
node.component_data[BASE_COMPONENT_TYPE].host_value,
node.key.token_ids,
node.hash_value,
prefix_keys,
extra_pools=aux_xfers or None,
)
self.ongoing_backup[operation_id] = (
node,
self.inc_host_lock_ref(node).to_dec_params(),
)
def prefetch_from_storage(
self,
req_id: str,
last_host_node: UnifiedTreeNode,
new_input_tokens: list[int],
last_hash: Optional[str] = None,
prefix_keys: Optional[list[str]] = None,
) -> None:
if not self.enable_storage or self.cache_controller is None:
return
extra_key = last_host_node.key.extra_key if last_host_node.key else None
prefetch_key = RadixKey(
new_input_tokens,
extra_key=extra_key,
is_bigram=self.is_eagle,
).page_aligned(self.page_size)
prefetch_length = len(prefetch_key)
if (
prefetch_length < self.prefetch_threshold
or self.cache_controller.prefetch_rate_limited()
):
return
anchor_lock_params = self.inc_host_lock_ref(last_host_node).to_dec_params()
host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
if host_indices is None:
self.evict_host(prefetch_length)
host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
if host_indices is None:
available_size = self.cache_controller.mem_pool_host.available_size()
prefetch_length = available_size - (available_size % self.page_size)
if prefetch_length >= self.prefetch_threshold:
prefetch_key = prefetch_key[:prefetch_length]
host_indices = self.cache_controller.mem_pool_host.alloc(
prefetch_length
)
else:
self.dec_host_lock_ref(last_host_node, anchor_lock_params)
return
if host_indices is None:
self.dec_host_lock_ref(last_host_node, anchor_lock_params)
return
comp_xfers: dict[ComponentType, list[PoolTransfer]] = {}
alloc_failed = False
for comp in self._components_tuple:
if comp.component_type == BASE_COMPONENT_TYPE:
continue
transfers = comp.build_hicache_transfers(
last_host_node,
CacheTransferPhase.PREFETCH,
token_ids=prefetch_key.token_ids,
prefetch_tokens=len(prefetch_key),
last_hash=last_hash,
)
if transfers == []:
alloc_failed = True
break
if transfers:
comp_xfers[comp.component_type] = transfers
kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=host_indices)
sidecar_xfers = self._build_sidecar_transfers(
CacheTransferPhase.PREFETCH, kv_xfer, comp_xfers
)
if alloc_failed:
self.cache_controller.append_host_mem_release(
host_indices=host_indices,
extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
)
self.dec_host_lock_ref(last_host_node, anchor_lock_params)
return
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
aux_xfers.extend(sidecar_xfers)
operation = self.cache_controller.prefetch(
req_id,
host_indices,
prefetch_key.token_ids,
last_hash,
prefix_keys,
extra_pools=aux_xfers or None,
)
self.ongoing_prefetch[req_id] = (
last_host_node,
prefetch_key,
host_indices,
operation,
anchor_lock_params,
comp_xfers,
)
self.cache_controller.prefetch_tokens_occupied += len(prefetch_key)
def _prefetch_timeout_check_linear_func(self, operation) -> bool:
return (
time.monotonic() - operation.start_time
> self.prefetch_timeout_base
+ len(operation.hash_value) * self.prefetch_timeout_per_page
)
def can_terminate_prefetch(self, operation) -> bool:
if self.prefetch_stop_policy == "best_effort":
return True
if len(operation.hash_value) == 0:
completed = False
else:
completed = (
operation.completed_tokens == len(operation.hash_value) * self.page_size
)
if self.prefetch_stop_policy == "wait_complete":
can_terminate = completed
elif self.prefetch_stop_policy == "timeout":
can_terminate = completed or self._prefetch_timeout_check_linear_func(
operation
)
else:
return True
operation_terminated = operation.is_terminated()
states = torch.tensor(
[1 - int(can_terminate), int(operation_terminated)],
dtype=torch.int,
)
if self.tp_world_size > 1:
torch.distributed.all_reduce(
states, op=torch.distributed.ReduceOp.MAX, group=self.tp_group
)
can_terminate = states[0].item() == 0
operation_terminated = states[1].item() == 1
return can_terminate or operation_terminated
def check_prefetch_progress(self, req_id: str) -> bool:
if req_id not in self.ongoing_prefetch:
return True
(
last_host_node,
prefetch_key,
host_indices,
operation,
anchor_lock_params,
comp_xfers,
) = self.ongoing_prefetch[req_id]
if operation.host_indices is None:
return True
if not self.can_terminate_prefetch(operation):
return False
completed_tokens, hash_value = self.cache_controller.terminate_prefetch(
operation
)
min_completed_tokens = completed_tokens
if self.tp_world_size > 1:
completed_tokens_tensor = torch.tensor(
min_completed_tokens, dtype=torch.int
)
torch.distributed.all_reduce(
completed_tokens_tensor,
op=torch.distributed.ReduceOp.MIN,
group=self.tp_group,
)
min_completed_tokens = int(completed_tokens_tensor.item())
fetched_key = prefetch_key[:min_completed_tokens]
insert_result = self._insert_helper_host(
last_host_node,
fetched_key,
host_indices[:min_completed_tokens],
hash_value[: min_completed_tokens // self.page_size],
)
for ct, xfers in comp_xfers.items():
self.components[ct].commit_hicache_transfer(
last_host_node,
CacheTransferPhase.PREFETCH,
xfers,
insert_result=insert_result,
pool_storage_result=operation.pool_storage_result,
)
self.cache_controller.mem_pool_host.free(
host_indices[: insert_result.prefix_len]
)
self.cache_controller.append_host_mem_release(
host_indices[min_completed_tokens:completed_tokens]
)
self.dec_host_lock_ref(last_host_node, anchor_lock_params)
del self.ongoing_prefetch[req_id]
self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key)
loaded_from_storage = min_completed_tokens - insert_result.prefix_len
self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage
logger.info(
"HiCache prefetch success req=%s completed_local=%d completed_synced=%d matched=%d loaded=%d tail_release=%d occupied=%d",
req_id,
completed_tokens,
min_completed_tokens,
insert_result.prefix_len,
loaded_from_storage,
completed_tokens - min_completed_tokens,
self.cache_controller.prefetch_tokens_occupied,
)
if self.enable_storage_metrics and self.storage_metrics_collector is not None:
self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage)
return True
def terminate_prefetch(self, req_id: str) -> None:
if req_id not in self.ongoing_prefetch:
return
_, _, _, operation, _, _ = self.ongoing_prefetch[req_id]
if operation.host_indices is None:
return
operation.mark_terminate()
def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)
def release_aborted_request(self, rid: str) -> None:
self.prefetch_loaded_tokens_by_reqid.pop(rid, None)
if rid not in self.ongoing_prefetch:
return
(
last_host_node,
prefetch_key,
host_indices,
operation,
anchor_lock_params,
comp_xfers,
) = self.ongoing_prefetch[rid]
if operation.host_indices is None:
return
completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
if self.tp_world_size > 1:
torch.distributed.barrier(group=self.tp_group)
self.dec_host_lock_ref(last_host_node, anchor_lock_params)
del self.ongoing_prefetch[rid]
self.cache_controller.append_host_mem_release(
host_indices=host_indices[:completed_tokens],
extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
)
self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key)
def _drain_storage_control_queues_impl(
self,
n_revoke: Optional[int],
n_backup: Optional[int],
n_release: Optional[int],
extra_release_counts: Optional[dict[PoolName, int]],
log_metrics: bool,
) -> None:
cc = self.cache_controller
def _drain_queue(q, limit: Optional[int]):
drained = 0
while limit is None or drained < limit:
try:
item = q.get_nowait()
except Empty:
break
drained += 1
yield item
def _drain_revoke():
drained = 0
for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke):
info = self.ongoing_prefetch.pop(req_id, None)
if info is None:
continue
drained += 1
(
last_host_node,
prefetch_key,
_host_indices,
_operation,
anchor_lock_params,
comp_xfers,
) = info
cc.append_host_mem_release(
extra_pools=[x for xfers in comp_xfers.values() for x in xfers]
)
self.dec_host_lock_ref(last_host_node, anchor_lock_params)
cc.prefetch_tokens_occupied -= len(prefetch_key)
if cc.prefetch_tokens_occupied < 0:
cc.prefetch_tokens_occupied = 0
return drained
def _drain_backup():
drained = 0
for operation in _drain_queue(cc.ack_backup_queue, n_backup):
drained += 1
entry = self.ongoing_backup.pop(operation.id, None)
if entry is not None:
node, lock_params = entry
self.dec_host_lock_ref(node, lock_params)
if (
log_metrics
and self.enable_storage_metrics
and self.storage_metrics_collector is not None
):
self.storage_metrics_collector.log_backuped_tokens(
operation.completed_tokens
)
return drained
def _drain_release():
host_indices_list = []
released_tokens = 0
for host_indices in _drain_queue(cc.host_mem_release_queue, n_release):
host_indices_list.append(host_indices)
released_tokens += len(host_indices)
if host_indices_list:
cc.mem_pool_host.free(torch.cat(host_indices_list, dim=0))
return len(host_indices_list), released_tokens
def _drain_extra_release():
drained: dict[PoolName, tuple[int, int]] = {}
if not extra_release_counts:
return drained
for pool_name, limit in extra_release_counts.items():
release_queue = cc.extra_host_mem_release_queues.get(pool_name)
if release_queue is None:
continue
host_indices_list = []
released_tokens = 0
for host_indices in _drain_queue(release_queue, limit):
host_indices_list.append(host_indices)
released_tokens += len(host_indices)
if host_indices_list:
entry = cc.mem_pool_host.entry_map.get(pool_name)
if entry is not None:
entry.host_pool.free(torch.cat(host_indices_list, dim=0))
drained[pool_name] = (len(host_indices_list), released_tokens)
return drained
_drain_revoke()
_drain_backup()
_drain_release()
_drain_extra_release()
def drain_storage_control_queues(self) -> None:
cc = self.cache_controller
extra_release_queues = getattr(cc, "extra_host_mem_release_queues", {})
extra_pool_names = list(extra_release_queues)
local_qsize_list = [
cc.prefetch_revoke_queue.qsize(),
cc.ack_backup_queue.qsize(),
cc.host_mem_release_queue.qsize(),
*[
extra_release_queues[pool_name].qsize()
for pool_name in extra_pool_names
],
]
qsizes = torch.tensor(
local_qsize_list,
dtype=torch.int,
)
if self.tp_world_size > 1:
torch.distributed.all_reduce(
qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
)
qsize_list = list(map(int, qsizes.tolist()))
n_revoke, n_backup, n_release = qsize_list[:3]
extra_release_counts = {
pool_name: count
for pool_name, count in zip(extra_pool_names, qsize_list[3:])
}
self._drain_storage_control_queues_impl(
n_revoke=n_revoke,
n_backup=n_backup,
n_release=n_release,
extra_release_counts=extra_release_counts,
log_metrics=True,
)
def _apply_storage_runtime_config(
self,
*,
storage_backend: Optional[str],
prefetch_threshold: int,
prefetch_timeout_base: float,
prefetch_timeout_per_ki_token: float,
hicache_storage_pass_prefix_keys: bool,
enable_storage: bool,
enable_storage_metrics: bool,
extra_metric_labels: Optional[dict[str, str]],
) -> None:
self.enable_storage = enable_storage
self.prefetch_threshold = prefetch_threshold
self.prefetch_timeout_base = prefetch_timeout_base
self.prefetch_timeout_per_page = (
self.page_size / 1024 * prefetch_timeout_per_ki_token
)
self.hicache_storage_pass_prefix_keys = hicache_storage_pass_prefix_keys
self.enable_storage_metrics = enable_storage_metrics
if self.enable_storage_metrics:
attn_cp_rank, attn_cp_size = (
self.cache_controller.get_attn_cp_rank_and_size()
)
labels = {
"storage_backend": storage_backend,
"tp_rank": self.cache_controller.tp_rank,
"dp_rank": self.cache_controller.dp_rank,
"pp_rank": self.cache_controller.pp_rank,
"pp_size": self.cache_controller.pp_size,
"attn_cp_rank": attn_cp_rank,
"attn_cp_size": attn_cp_size,
}
if extra_metric_labels:
labels.update(extra_metric_labels)
existing_collector = self.storage_metrics_collector
if existing_collector is None:
self.storage_metrics_collector = StorageMetricsCollector(labels=labels)
elif set(existing_collector.labels.keys()) == set(labels.keys()):
existing_collector.labels = labels
else:
logger.warning(
"Storage metrics labels changed (%s -> %s). Keep existing labels to avoid duplicate metric registration.",
sorted(existing_collector.labels.keys()),
sorted(labels.keys()),
)
else:
self.storage_metrics_collector = None
def attach_storage_backend(
self,
storage_backend: str,
storage_backend_extra_config_json: Optional[str] = None,
served_model_name: Optional[str] = None,
hicache_storage_prefetch_policy: Optional[str] = None,
hicache_write_policy: Optional[str] = None,
) -> tuple[bool, str]:
return (
False,
"UnifiedRadixCache does not support runtime HiCache storage attach yet. "
"Configure hicache_storage_backend at startup instead.",
)
def detach_storage_backend(self) -> tuple[bool, str]:
return (
False,
"UnifiedRadixCache does not support runtime HiCache storage detach yet. "
"Restart without hicache_storage_backend to disable it.",
)
def clear_storage_backend(self) -> bool:
try:
ok = self.cache_controller.clear_storage_backend()
except Exception as e:
logger.error("Failed to clear hierarchical cache storage backend: %s", e)
return False
if ok:
logger.info("Hierarchical cache storage backend cleared successfully!")
return ok
# ---- HiCache: Async Event Management ----
def writing_check(self, write_back: bool = False) -> None:
@@ -1387,6 +2033,8 @@ class UnifiedRadixCache(BasePrefixCache):
node, params = entry
if params is not None:
self.dec_lock_ref(node, params)
if self.enable_storage:
self.write_backup_storage(node)
cc.ack_write_queue.clear()
assert len(self.ongoing_write_through) == 0
return
@@ -1415,6 +2063,8 @@ class UnifiedRadixCache(BasePrefixCache):
for ack_id in ack_list:
node, params = self.ongoing_write_through.pop(ack_id)
self.dec_lock_ref(node, params)
if self.enable_storage:
self.write_backup_storage(node)
finish_count -= 1
def loading_check(self) -> None:
@@ -1484,6 +2134,12 @@ class UnifiedRadixCache(BasePrefixCache):
"""Called per scheduler step to poll async HiCache events."""
self.writing_check()
self.loading_check()
if self.enable_storage:
self.drain_storage_control_queues()
if self.enable_storage_metrics and self.storage_metrics_collector is not None:
self.storage_metrics_collector.log_storage_metrics(
self.cache_controller.storage_backend.get_stats()
)
def flush_write_through_acks(self) -> None:
"""Flush pending write-through acknowledgements."""
@@ -1883,10 +2539,6 @@ class UnifiedRadixCache(BasePrefixCache):
logger.error(msg)
self.pretty_print()
raise AssertionError(msg)
logger.debug(
f"Sanity check PASSED: {len(all_nodes)} nodes, "
f"{len(self.tree_components)} components"
)
def _check_lru_linked_list(
self,