1212 lines
45 KiB
Python
1212 lines
45 KiB
Python
from __future__ import annotations
|
|
|
|
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
|
from sglang.srt.mem_cache.utils import convert_to_bigram_key
|
|
|
|
"""
|
|
Copyright 2023-2024 SGLang Team
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
"""
|
|
The radix tree data structure for managing the KV cache.
|
|
"""
|
|
|
|
import heapq
|
|
import logging
|
|
import sys
|
|
import time
|
|
from functools import lru_cache, partial
|
|
from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
from sglang.srt.disaggregation.kv_events import (
|
|
MEDIUM_GPU,
|
|
AllBlocksCleared,
|
|
BlockRemoved,
|
|
BlockStored,
|
|
)
|
|
from sglang.srt.mem_cache.base_prefix_cache import (
|
|
BasePrefixCache,
|
|
DecLockRefParams,
|
|
DecLockRefResult,
|
|
EvictParams,
|
|
EvictResult,
|
|
IncLockRefResult,
|
|
InsertParams,
|
|
InsertResult,
|
|
MatchPrefixParams,
|
|
MatchResult,
|
|
)
|
|
from sglang.srt.mem_cache.evict_policy import (
|
|
EvictionStrategy,
|
|
FIFOStrategy,
|
|
FILOStrategy,
|
|
LFUStrategy,
|
|
LRUStrategy,
|
|
MRUStrategy,
|
|
PriorityStrategy,
|
|
SLRUStrategy,
|
|
)
|
|
from sglang.srt.mem_cache.hicache_storage import get_hash_str, hash_str_to_int64
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.managers.schedule_batch import Req
|
|
|
|
|
|
class RadixKey:
|
|
"""A sequence of token ids used as a radix tree key.
|
|
|
|
Slicing returns an O(1) view sharing the underlying token list instead of
|
|
copying it. The tree walks in ``match_prefix``/``insert`` re-slice the
|
|
remaining key at every node hop, so copying slices makes a single match
|
|
O(len * hops) — quadratic for long cached prefixes. Views keep it linear.
|
|
|
|
Views keep the underlying list alive, so keys stored in long-lived tree
|
|
nodes must be compacted first (see ``compacted()``); only transient keys
|
|
inside a tree walk should remain views.
|
|
"""
|
|
|
|
__slots__ = ("_tokens", "_start", "_stop", "extra_key", "is_bigram")
|
|
|
|
def __init__(
|
|
self,
|
|
token_ids: List[int],
|
|
extra_key: Optional[str] = None,
|
|
is_bigram: bool = False,
|
|
):
|
|
# token ids sequence (full backing list; this key covers
|
|
# [_start, _stop) of it)
|
|
self._tokens = token_ids
|
|
self._start = 0
|
|
self._stop = len(token_ids)
|
|
# extra key (e.g. lora_id, cache_salt)
|
|
self.extra_key = extra_key
|
|
# is bigram key
|
|
self.is_bigram = is_bigram
|
|
|
|
@classmethod
|
|
def _view(cls, base: "RadixKey", start: int, stop: int) -> "RadixKey":
|
|
obj = cls.__new__(cls)
|
|
obj._tokens = base._tokens
|
|
obj._start = start
|
|
obj._stop = stop
|
|
obj.extra_key = base.extra_key
|
|
obj.is_bigram = base.is_bigram
|
|
return obj
|
|
|
|
@property
|
|
def token_ids(self) -> List[int]:
|
|
if self._start == 0 and self._stop == len(self._tokens):
|
|
return self._tokens
|
|
return self._tokens[self._start : self._stop]
|
|
|
|
@token_ids.setter
|
|
def token_ids(self, value: List[int]) -> None:
|
|
self._tokens = value
|
|
self._start = 0
|
|
self._stop = len(value)
|
|
|
|
def compacted(self) -> "RadixKey":
|
|
"""Return an equivalent key that does not pin a larger backing list."""
|
|
if self._start == 0 and self._stop == len(self._tokens):
|
|
return self
|
|
return RadixKey(
|
|
self._tokens[self._start : self._stop], self.extra_key, self.is_bigram
|
|
)
|
|
|
|
def __len__(self) -> int:
|
|
return self._stop - self._start
|
|
|
|
def __iter__(self) -> Iterator[int]:
|
|
# NOTE: map over an index range instead of islice — islice skips the
|
|
# first `_start` elements one by one, which is O(_start) per call.
|
|
return map(self._tokens.__getitem__, range(self._start, self._stop))
|
|
|
|
def __getitem__(self, idx: Union[int, slice]) -> "RadixKey":
|
|
if isinstance(idx, slice):
|
|
start, stop, step = idx.indices(len(self))
|
|
if step != 1:
|
|
raise ValueError(f"RadixKey does not support strided slicing: {idx}")
|
|
return RadixKey._view(self, self._start + start, self._start + stop)
|
|
if idx < 0:
|
|
idx += len(self)
|
|
return RadixKey([self._tokens[self._start + idx]], self.extra_key)
|
|
|
|
def __repr__(self) -> str:
|
|
preview = self.token_ids[:10]
|
|
return f"RadixKey(extra_key={self.extra_key!r}, token_ids={preview}{'...' if len(self) > 10 else ''})"
|
|
|
|
|
|
def maybe_bigram_convert(
|
|
is_eagle: bool,
|
|
key: RadixKey,
|
|
value: Optional[torch.Tensor] = None,
|
|
) -> Tuple[RadixKey, Optional[torch.Tensor]]:
|
|
if is_eagle and not key.is_bigram:
|
|
key.token_ids = convert_to_bigram_key(key.token_ids)
|
|
key.is_bigram = True
|
|
if value is not None:
|
|
value = value[: len(key)]
|
|
return key, value
|
|
|
|
|
|
def page_align_keys(key: list, page_size) -> list:
|
|
if page_size == 1:
|
|
return key
|
|
page_aligned_len = len(key) // page_size * page_size
|
|
return key[:page_aligned_len]
|
|
|
|
|
|
def ceil_to_page_len(length: int, page_size: int) -> int:
|
|
if page_size <= 1:
|
|
return length
|
|
return ((length + page_size - 1) // page_size) * page_size
|
|
|
|
|
|
def floor_to_page_len(length: int, page_size: int) -> int:
|
|
if page_size <= 1:
|
|
return length
|
|
return length // page_size * page_size
|
|
|
|
|
|
class TreeNode:
|
|
|
|
counter = 0
|
|
# Replicated logical access clock (mirrors mamba_radix_cache's
|
|
# get_last_access_time pattern). A process-global monotonically-increasing
|
|
# counter used as the source of ``last_access_time``. Under CP shared-KV
|
|
# every rank advances it through the IDENTICAL sequence of match/insert
|
|
# events (requests are broadcast from rank 0 and the radix tree is
|
|
# replicated), so ``last_access_time`` is rank-replicated and eviction
|
|
# victim order is collective-free. The old per-rank ``time.monotonic()``
|
|
# diverged victim order across ranks and could deadlock the collective-
|
|
# coupled writeback/load-back. Unique integers also give a strict total
|
|
# order (no tie ambiguity). Reset alongside ``counter`` in HiRadixCache.reset.
|
|
last_access_counter = 0
|
|
|
|
@staticmethod
|
|
def next_access_time() -> int:
|
|
v = TreeNode.last_access_counter
|
|
TreeNode.last_access_counter += 1
|
|
return v
|
|
|
|
def __init__(self, id: Optional[int] = None, priority: int = 0):
|
|
self.children = {}
|
|
self.parent: TreeNode = None
|
|
self.key: RadixKey = None
|
|
self.value: Optional[torch.Tensor] = None
|
|
self.lock_ref = 0
|
|
self.pin_expiry: float = (
|
|
0.0 # absolute expiry time (time.monotonic()), 0 = not pinned
|
|
)
|
|
self.pin_ttl: int = 0 # original TTL in seconds, for refresh-on-hit
|
|
self.last_access_time = TreeNode.next_access_time()
|
|
# creation_time stays wall-clock: it is only read by FIFO/FILO eviction
|
|
# (not the CP/default LRU/SLRU path). FIFO/FILO are not CP-safe today and
|
|
# are not the default; if ever enabled under CP, route this to the
|
|
# logical clock too.
|
|
self.creation_time = time.monotonic()
|
|
|
|
self.hit_count = 0
|
|
# indicating the node is locked to protect from eviction
|
|
# incremented when the node is referenced by a storage operation
|
|
self.host_ref_counter = 0
|
|
# store the host indices of KV cache
|
|
self.host_value: Optional[torch.Tensor] = None
|
|
self.host_len: int = 0
|
|
self.cp_hicache = None
|
|
# store hash values of each pages
|
|
self.hash_value: Optional[List[str]] = None
|
|
# priority for priority-aware eviction
|
|
self.priority = priority
|
|
|
|
self.id = TreeNode.counter if id is None else id
|
|
TreeNode.counter += 1
|
|
|
|
@property
|
|
def evicted(self):
|
|
return self.value is None
|
|
|
|
@property
|
|
def backuped(self):
|
|
return self.host_value is not None
|
|
|
|
def protect_host(self):
|
|
"""Protect the host value from eviction."""
|
|
self.host_ref_counter += 1
|
|
|
|
def release_host(self):
|
|
"""Release the host value, allowing it to be evicted."""
|
|
if self.host_ref_counter > 0:
|
|
self.host_ref_counter -= 1
|
|
else:
|
|
raise RuntimeError("Host reference counter is already zero.")
|
|
|
|
def get_last_hash_value(self) -> Optional[str]:
|
|
"""Returns the hash value of the last page in this node."""
|
|
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: TreeNode) -> List[str]:
|
|
if node is None or node.hash_value is None:
|
|
return []
|
|
|
|
return node.get_prefix_hash_values(node.parent) + node.hash_value
|
|
|
|
def __lt__(self, other: "TreeNode"):
|
|
return self.last_access_time < other.last_access_time
|
|
|
|
|
|
def _check_extra_key(key0: RadixKey, key1: RadixKey):
|
|
if key0.extra_key != key1.extra_key:
|
|
raise ValueError(
|
|
f"_key_match should be run on the same extra key, but got key0.extra_key={key0.extra_key} != key1.extra_key={key1.extra_key}"
|
|
)
|
|
|
|
|
|
def _key_match_page_size1(key0: RadixKey, key1: RadixKey):
|
|
_check_extra_key(key0, key1)
|
|
# Index the backing lists directly so that view keys (nonzero _start) are
|
|
# matched without materializing a copy of their suffix.
|
|
t0, s0 = key0._tokens, key0._start
|
|
t1, s1 = key1._tokens, key1._start
|
|
n = min(len(key0), len(key1))
|
|
i = 0
|
|
while i < n and t0[s0 + i] == t1[s1 + i]:
|
|
i += 1
|
|
return i
|
|
|
|
|
|
def _key_match_paged(key0: RadixKey, key1: RadixKey, page_size: int):
|
|
_check_extra_key(key0, key1)
|
|
t0, s0 = key0._tokens, key0._start
|
|
t1, s1 = key1._tokens, key1._start
|
|
min_len = min(len(key0), len(key1))
|
|
|
|
i = 0
|
|
while i < min_len:
|
|
step = min(page_size, min_len - i)
|
|
if t0[s0 + i : s0 + i + step] != t1[s1 + i : s1 + i + step]:
|
|
break
|
|
i += step
|
|
|
|
return i
|
|
|
|
|
|
def get_child_key(key: RadixKey, page_size: int = 1):
|
|
if page_size == 1:
|
|
plain_key = key._tokens[key._start]
|
|
else:
|
|
plain_key = tuple(key._tokens[key._start : key._start + page_size])
|
|
if key.extra_key is None:
|
|
return plain_key
|
|
else:
|
|
return (key.extra_key, plain_key)
|
|
|
|
|
|
def compute_node_hash_values(node: "TreeNode", page_size: int) -> List[str]:
|
|
"""Compute SHA256-based hash values for position-aware identification.
|
|
|
|
Args:
|
|
node: The TreeNode to compute hash values for
|
|
page_size: The page size for chunking tokens
|
|
|
|
Returns:
|
|
List of SHA256 hex strings, one per page
|
|
"""
|
|
hash_values = []
|
|
|
|
# Get parent's last hash value if parent exists
|
|
parent_hash = None
|
|
if node.parent is not None and node.parent.hash_value is not None:
|
|
# Check if parent is root by checking if it has empty key
|
|
if len(node.parent.key) > 0 and len(node.parent.hash_value) > 0:
|
|
parent_hash = node.parent.hash_value[-1]
|
|
|
|
# Iterate through node's pages
|
|
for start in range(0, len(node.key), page_size):
|
|
page_tokens = node.key.token_ids[start : start + page_size]
|
|
if not page_tokens:
|
|
continue
|
|
|
|
# Use SHA256-based chaining via get_hash_str
|
|
hash_val = get_hash_str(page_tokens, prior_hash=parent_hash)
|
|
hash_values.append(hash_val)
|
|
parent_hash = hash_val
|
|
|
|
return hash_values
|
|
|
|
|
|
def split_node_hash_value(
|
|
child_hash_value: Optional[List[str]], split_len: int, page_size: int
|
|
) -> tuple[Optional[List[str]], Optional[List[str]]]:
|
|
"""Split hash_value between parent and child nodes during node splitting.
|
|
|
|
Args:
|
|
child_hash_value: The hash_value list from the child node being split
|
|
split_len: The length at which to split (in tokens)
|
|
page_size: The page size for calculating number of pages
|
|
|
|
Returns:
|
|
Tuple of (new_node_hash_value, updated_child_hash_value)
|
|
"""
|
|
if child_hash_value is None:
|
|
return None, None
|
|
|
|
if page_size == 1:
|
|
split_pages = split_len
|
|
else:
|
|
split_pages = split_len // page_size
|
|
|
|
new_node_hash = child_hash_value[:split_pages]
|
|
child_hash = child_hash_value[split_pages:]
|
|
|
|
return new_node_hash, child_hash
|
|
|
|
|
|
class RadixCache(BasePrefixCache):
|
|
def __init__(self, params: CacheInitParams):
|
|
self.disable = params.disable
|
|
self.req_to_token_pool = params.req_to_token_pool
|
|
self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
|
|
self.page_size = params.page_size
|
|
self.enable_kv_cache_events = params.enable_kv_cache_events
|
|
self.is_eagle = params.is_eagle
|
|
self.disable_finished_insert = params.disable_finished_insert
|
|
self.eviction_policy = params.eviction_policy.lower()
|
|
|
|
self.kv_event_queue = []
|
|
|
|
if params.enable_metrics:
|
|
self.init_metrics_collector()
|
|
|
|
if self.token_to_kv_pool_allocator:
|
|
self.device = self.token_to_kv_pool_allocator.device
|
|
else:
|
|
self.device = torch.device("cpu")
|
|
|
|
if self.page_size == 1:
|
|
self.key_match_fn = _key_match_page_size1
|
|
self.get_child_key_fn = get_child_key
|
|
else:
|
|
self.key_match_fn = partial(_key_match_paged, page_size=self.page_size)
|
|
self.get_child_key_fn = partial(get_child_key, page_size=self.page_size)
|
|
|
|
if self.eviction_policy == "lru":
|
|
self.eviction_strategy: EvictionStrategy = LRUStrategy()
|
|
elif self.eviction_policy == "lfu":
|
|
self.eviction_strategy: EvictionStrategy = LFUStrategy()
|
|
elif self.eviction_policy == "fifo":
|
|
self.eviction_strategy: EvictionStrategy = FIFOStrategy()
|
|
elif self.eviction_policy == "mru":
|
|
self.eviction_strategy: EvictionStrategy = MRUStrategy()
|
|
elif self.eviction_policy == "filo":
|
|
self.eviction_strategy: EvictionStrategy = FILOStrategy()
|
|
elif self.eviction_policy == "priority":
|
|
self.eviction_strategy: EvictionStrategy = PriorityStrategy()
|
|
elif self.eviction_policy == "slru":
|
|
self.eviction_strategy: EvictionStrategy = SLRUStrategy()
|
|
|
|
else:
|
|
raise ValueError(
|
|
f"Unknown eviction policy: {self.eviction_policy}. Supported policies: 'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority', 'slru'."
|
|
)
|
|
|
|
self.evictable_leaves = set()
|
|
self.reset()
|
|
|
|
@classmethod
|
|
def create_simulated(
|
|
self,
|
|
disable: bool = False,
|
|
mock_allocator: Optional[Any] = None,
|
|
page_size: int = 1,
|
|
enable_kv_cache_events: bool = False,
|
|
) -> RadixCache:
|
|
"""Init a radix cache without memory pools for simulation purpose."""
|
|
params = CacheInitParams(
|
|
disable=disable,
|
|
req_to_token_pool=None,
|
|
token_to_kv_pool_allocator=mock_allocator,
|
|
page_size=page_size,
|
|
enable_kv_cache_events=enable_kv_cache_events,
|
|
)
|
|
return RadixCache(params)
|
|
|
|
##### Public API #####
|
|
|
|
def reset(self):
|
|
# Initialize root with minimum priority so any real priority overrides it
|
|
self.root_node = TreeNode(priority=-sys.maxsize)
|
|
self.root_node.key = RadixKey(token_ids=[], extra_key=None)
|
|
self.root_node.value = []
|
|
self.root_node.host_value = []
|
|
self.root_node.lock_ref = 1
|
|
self.root_node.hash_value = []
|
|
self.evictable_size_ = 0
|
|
self.protected_size_ = 0
|
|
self.evictable_leaves.clear()
|
|
self._record_all_cleared_event()
|
|
|
|
def maybe_bigram_convert(
|
|
self, key: RadixKey, value: Optional[torch.Tensor] = None
|
|
) -> Tuple[RadixKey, Optional[torch.Tensor]]:
|
|
return maybe_bigram_convert(self.is_eagle, key, value)
|
|
|
|
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
|
|
"""Find the longest cached prefix of ``key`` in the radix tree.
|
|
|
|
The logical namespace for prefix matching is determined by both the
|
|
token id sequence and the optional ``extra_key`` carried by ``RadixKey``.
|
|
Entries that share identical leading token ids but have *different*
|
|
``extra_key`` values are intentionally kept disjoint and never share
|
|
prefix nodes. This is useful to:
|
|
|
|
* Isolate KV cache lines for different LoRA / adapter IDs.
|
|
* Separate requests that intentionally should not share state (e.g.,
|
|
different sampling salt, cache version, or retrieval augmentation
|
|
context) by supplying a distinct ``extra_key``.
|
|
|
|
Args:
|
|
params (MatchPrefixParams): Parameters containing the lookup key
|
|
with a list of token ids and an optional ``extra_key`` namespace tag.
|
|
If ``page_size > 1`` the length is internally truncated to a multiple
|
|
of ``page_size`` before matching. Passing an empty key returns an
|
|
empty result with the root as the last node.
|
|
|
|
Returns:
|
|
MatchResult: ``device_indices`` is a 1-D ``torch.int64`` tensor of
|
|
the concatenated KV cache indices corresponding to the longest
|
|
cached prefix (may be length 0). ``last_device_node`` and
|
|
``last_host_node`` (currently the same) are the tree node objects
|
|
representing the terminal node of the matched prefix. This method
|
|
may mutate internal structure by splitting an existing node if the
|
|
match ends inside a stored segment.
|
|
|
|
Internal updates:
|
|
* Refreshes access metadata (timestamps) used by the
|
|
configured eviction strategy.
|
|
* If the lookup ends inside a stored segment the node is split once
|
|
to expose a precise boundary; this structural refinement improves
|
|
subsequent match efficiency and does not duplicate data.
|
|
"""
|
|
key = params.key
|
|
key, _ = self.maybe_bigram_convert(key)
|
|
|
|
def empty_match_result():
|
|
return MatchResult(
|
|
device_indices=torch.empty(
|
|
(0,),
|
|
dtype=torch.int64,
|
|
device=self.device,
|
|
),
|
|
last_device_node=self.root_node,
|
|
last_host_node=self.root_node,
|
|
)
|
|
|
|
if self.disable or len(key) == 0:
|
|
return empty_match_result()
|
|
|
|
if self.page_size != 1:
|
|
page_aligned_len = len(key) // self.page_size * self.page_size
|
|
key = key[:page_aligned_len]
|
|
|
|
if len(key) == 0:
|
|
return empty_match_result()
|
|
|
|
value, last_node = self._match_prefix_helper(self.root_node, key)
|
|
if value:
|
|
value = torch.cat(value)
|
|
else:
|
|
value = torch.empty((0,), dtype=torch.int64, device=self.device)
|
|
return MatchResult(
|
|
device_indices=value,
|
|
last_device_node=last_node,
|
|
last_host_node=last_node,
|
|
)
|
|
|
|
def insert(self, params: InsertParams) -> InsertResult:
|
|
if self.disable:
|
|
return InsertResult(prefix_len=0)
|
|
|
|
key = params.key
|
|
value = params.value
|
|
priority = params.priority
|
|
chunked = params.chunked
|
|
|
|
if value is None:
|
|
value = torch.tensor(key.token_ids, dtype=torch.int64)
|
|
|
|
key, value = self.maybe_bigram_convert(key, value)
|
|
|
|
prefix_len = self._insert_helper(self.root_node, key, value, priority, chunked)
|
|
return InsertResult(prefix_len=prefix_len)
|
|
|
|
def _free_kv_indices_range(
|
|
self,
|
|
kv_indices: torch.Tensor,
|
|
start: int,
|
|
end: int,
|
|
*,
|
|
include_partial_tail: bool = False,
|
|
):
|
|
if getattr(self, "_uses_cp_hicache", False):
|
|
start = ceil_to_page_len(start, self.page_size)
|
|
if not include_partial_tail:
|
|
end = floor_to_page_len(end, self.page_size)
|
|
if end < start:
|
|
end = start
|
|
self.token_to_kv_pool_allocator.free(kv_indices[start:end])
|
|
|
|
def _should_free_cp_exact_match_tail(
|
|
self, *, start: int, end: int, key_len: int
|
|
) -> bool:
|
|
"""Return whether duplicate-prefix freeing must include a CP tail page.
|
|
|
|
CP HiCache keeps radix keys at scheduler-visible valid lengths while the
|
|
allocator owns whole physical pages. For an exact duplicate hit whose
|
|
valid key ends inside a page (typical tiny EAGLE/bigram requests), the
|
|
duplicate request page is not retained by a new radix node. Flooring
|
|
``end`` to the previous page boundary would therefore leak that newly
|
|
allocated page. Partial/new-node insertions still use the conservative
|
|
page-floored free path because their tail page is represented by the
|
|
radix node being inserted.
|
|
"""
|
|
|
|
return (
|
|
getattr(self, "_uses_cp_hicache", False)
|
|
and key_len > 0
|
|
and end == key_len
|
|
and end > start
|
|
)
|
|
|
|
def cache_finished_req(self, req: Req, is_insert: bool = True):
|
|
"""Cache request when it finishes."""
|
|
# In deterministic mode, disable finished request insertion to radix cache
|
|
if self.disable_finished_insert:
|
|
is_insert = False
|
|
|
|
kv_committed_len = req.pop_committed_kv_cache()
|
|
if self.disable:
|
|
kv_indices = self.req_to_token_pool.req_to_token[
|
|
req.req_pool_idx, :kv_committed_len
|
|
]
|
|
self.token_to_kv_pool_allocator.free(kv_indices)
|
|
return
|
|
|
|
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, : len(token_ids)
|
|
]
|
|
|
|
# Maybe convert to bigram keys for EAGLE
|
|
keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
|
|
if not getattr(self, "_uses_cp_hicache", False):
|
|
keys = page_align_keys(keys, self.page_size)
|
|
values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
|
|
radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
|
|
prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None)
|
|
|
|
# Radix Cache takes one ref in memory pool
|
|
if is_insert:
|
|
priority = getattr(req, "priority", 0) or 0
|
|
result = self.insert(
|
|
InsertParams(
|
|
key=radix_key,
|
|
value=values,
|
|
priority=priority,
|
|
cp_hicache_prepared_backup=prepared_cp_backup,
|
|
)
|
|
)
|
|
if getattr(result, "pending_backup_deferred_node", None) is not None:
|
|
if (
|
|
prepared_cp_backup is not None
|
|
and not getattr(prepared_cp_backup, "attached", False)
|
|
and hasattr(self, "_rollback_prepared_cp_backup")
|
|
):
|
|
self._rollback_prepared_cp_backup(
|
|
prepared_cp_backup, "pending_backup_split_deferred"
|
|
)
|
|
self._free_kv_indices_range(
|
|
kv_indices,
|
|
req.cache_protected_len,
|
|
len(keys),
|
|
include_partial_tail=True,
|
|
)
|
|
else:
|
|
if (
|
|
prepared_cp_backup is not None
|
|
and not getattr(prepared_cp_backup, "attached", False)
|
|
and hasattr(self, "_rollback_prepared_cp_backup")
|
|
):
|
|
self._rollback_prepared_cp_backup(
|
|
prepared_cp_backup, "insert_miss"
|
|
)
|
|
new_prefix_len = result.prefix_len
|
|
# Free the duplicates that were already in the tree
|
|
self._free_kv_indices_range(
|
|
kv_indices,
|
|
req.cache_protected_len,
|
|
new_prefix_len,
|
|
include_partial_tail=self._should_free_cp_exact_match_tail(
|
|
start=req.cache_protected_len,
|
|
end=new_prefix_len,
|
|
key_len=len(keys),
|
|
),
|
|
)
|
|
else:
|
|
if prepared_cp_backup is not None and hasattr(
|
|
self, "_rollback_prepared_cp_backup"
|
|
):
|
|
self._rollback_prepared_cp_backup(prepared_cp_backup, "no_insert")
|
|
self._free_kv_indices_range(
|
|
kv_indices,
|
|
req.cache_protected_len,
|
|
len(keys),
|
|
include_partial_tail=True,
|
|
)
|
|
if prepared_cp_backup is not None:
|
|
req.cp_hicache_prepared_backup = None
|
|
|
|
# Free the unaligned tail.
|
|
#
|
|
# CP HiCache radix keys are scheduler-visible valid lengths, while the
|
|
# backing allocator is page-granular. Freeing a loc inside the retained
|
|
# tail page releases the whole page and leaves radix accounting pointing
|
|
# at freed memory (observed as available_size=max with evictable_size>0
|
|
# after EAGLE warmup). Therefore CP starts tail free at the next page
|
|
# boundary; non-CP keeps the historical token boundary because keys were
|
|
# already floored by page_align_keys above.
|
|
tail_free_start = len(keys)
|
|
if getattr(self, "_uses_cp_hicache", False):
|
|
tail_free_start = ceil_to_page_len(tail_free_start, self.page_size)
|
|
self.token_to_kv_pool_allocator.free(kv_indices[tail_free_start:])
|
|
|
|
# Remove req slot release the cache lock
|
|
self.dec_lock_ref(req.last_node)
|
|
|
|
def cache_unfinished_req(self, req: Req, chunked=False):
|
|
"""Cache request when it is unfinished."""
|
|
if self.disable:
|
|
return
|
|
|
|
token_ids = req.fill_ids
|
|
kv_indices = self.req_to_token_pool.req_to_token[
|
|
req.req_pool_idx, : len(token_ids)
|
|
]
|
|
|
|
# Maybe convert to bigram keys for EAGLE
|
|
keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
|
|
if not getattr(self, "_uses_cp_hicache", False):
|
|
keys = page_align_keys(keys, self.page_size)
|
|
elif chunked and self.page_size > 1:
|
|
# CP HiCache owns KV at page granularity. For an unfinished
|
|
# chunked request, cache only completed pages and keep the current
|
|
# sub-page tail in the request-owned KV slots until a later chunk
|
|
# or the final insert can make it page-complete. This preserves
|
|
# intermediate chunk caching without repeatedly creating/pruning
|
|
# stale valid-tail radix nodes.
|
|
completed_len = floor_to_page_len(len(keys), self.page_size)
|
|
if completed_len <= req.cache_protected_len:
|
|
req.prefix_indices = kv_indices
|
|
return
|
|
keys = keys[:completed_len]
|
|
values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
|
|
radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
|
|
prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None)
|
|
|
|
# Radix Cache takes one ref in memory pool
|
|
result = self.insert(
|
|
InsertParams(
|
|
key=radix_key,
|
|
value=values,
|
|
chunked=chunked,
|
|
priority=getattr(req, "priority", 0) or 0,
|
|
cp_hicache_prepared_backup=prepared_cp_backup,
|
|
)
|
|
)
|
|
if getattr(result, "pending_backup_deferred_node", None) is not None:
|
|
if (
|
|
prepared_cp_backup is not None
|
|
and not getattr(prepared_cp_backup, "attached", False)
|
|
and hasattr(self, "_rollback_prepared_cp_backup")
|
|
and chunked
|
|
):
|
|
# A chunked middle insert can be followed by another forward for
|
|
# the same request, so an unattached prepared reservation would
|
|
# cover the wrong suffix and block the next chunk from preparing
|
|
# its own backup. Non-chunked disagg prefill still has a later
|
|
# cache_finished_req() pass after KV transfer; keep the prepared
|
|
# backup alive so that final insert can attach it instead of
|
|
# falling back to post-forward all-layer catch-up backup.
|
|
self._rollback_prepared_cp_backup(
|
|
prepared_cp_backup, "unfinished_pending_backup_split_deferred"
|
|
)
|
|
req.cp_hicache_prepared_backup = None
|
|
if len(values) < len(kv_indices):
|
|
req.prefix_indices = torch.cat([values, kv_indices[len(values) :]])
|
|
else:
|
|
req.prefix_indices = values
|
|
return
|
|
if (
|
|
prepared_cp_backup is not None
|
|
and not getattr(prepared_cp_backup, "attached", False)
|
|
and hasattr(self, "_rollback_prepared_cp_backup")
|
|
):
|
|
self._rollback_prepared_cp_backup(prepared_cp_backup, "insert_miss")
|
|
if prepared_cp_backup is not None:
|
|
req.cp_hicache_prepared_backup = None
|
|
new_prefix_len = result.prefix_len
|
|
|
|
self._free_kv_indices_range(
|
|
kv_indices,
|
|
req.cache_protected_len,
|
|
new_prefix_len,
|
|
include_partial_tail=self._should_free_cp_exact_match_tail(
|
|
start=req.cache_protected_len,
|
|
end=new_prefix_len,
|
|
key_len=len(keys),
|
|
),
|
|
)
|
|
|
|
# The prefix indices could be updated, reuse it
|
|
match_result = self.match_prefix(
|
|
MatchPrefixParams(key=radix_key, cp_floor_exact=False)
|
|
)
|
|
new_indices, new_last_node = (
|
|
match_result.device_indices,
|
|
match_result.last_device_node,
|
|
)
|
|
assert len(new_indices) == len(keys), f"{len(new_indices)=}, {len(keys)=}"
|
|
|
|
self.req_to_token_pool.write(
|
|
(req.req_pool_idx, slice(req.cache_protected_len, len(new_indices))),
|
|
new_indices[req.cache_protected_len :],
|
|
)
|
|
|
|
# The cache_protected_len is not always equal to len(req.prefix_indices)
|
|
# since for page_size > 1, the partial part is added to req.prefix_indices, but that part of kv indices is not added to the tree.
|
|
# It should be freed in the next cache_unfinished_req and final cache_finished_req to avoid memory leak.
|
|
# So we introduce this `cache_protected_len` field to make sure the partial part can be freed correctly.
|
|
req.cache_protected_len = len(new_indices)
|
|
|
|
self.dec_lock_ref(req.last_node)
|
|
self.inc_lock_ref(new_last_node)
|
|
|
|
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
|
|
# - page_size != 1: there is a partial page at the end, keep the full kv_indices
|
|
# - eagle case: bigram keys will only cache len - 1 kv indices
|
|
if len(new_indices) < len(kv_indices):
|
|
req.prefix_indices = torch.cat(
|
|
[new_indices, kv_indices[len(new_indices) :]]
|
|
)
|
|
else:
|
|
req.prefix_indices = new_indices
|
|
|
|
req.last_node = new_last_node
|
|
|
|
def pretty_print(self):
|
|
self._print_helper(self.root_node, 0)
|
|
print(f"#tokens: {self.total_size()}")
|
|
|
|
def total_size(self):
|
|
return self._total_size_helper()
|
|
|
|
def evict(self, params: EvictParams) -> EvictResult:
|
|
if self.disable:
|
|
return EvictResult()
|
|
|
|
start_time = time.perf_counter()
|
|
num_tokens = params.num_tokens
|
|
leaves = list(self.evictable_leaves)
|
|
eviction_heap = [
|
|
(self.eviction_strategy.get_priority(node), node) for node in leaves
|
|
]
|
|
heapq.heapify(eviction_heap)
|
|
|
|
num_evicted = 0
|
|
while num_evicted < num_tokens and len(eviction_heap):
|
|
_priority, x = heapq.heappop(eviction_heap)
|
|
|
|
self.token_to_kv_pool_allocator.free(x.value)
|
|
num_evicted += len(x.value)
|
|
self._delete_leaf(x)
|
|
|
|
if len(x.parent.children) == 0 and x.parent.lock_ref == 0:
|
|
new_priority = self.eviction_strategy.get_priority(x.parent)
|
|
heapq.heappush(eviction_heap, (new_priority, x.parent))
|
|
|
|
self._record_remove_event(x)
|
|
|
|
self.update_eviction_metrics(num_evicted, start_time)
|
|
return EvictResult(num_tokens_evicted=num_evicted)
|
|
|
|
def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult:
|
|
if self.disable:
|
|
return IncLockRefResult(delta=0)
|
|
|
|
delta = 0
|
|
while node != self.root_node:
|
|
if node.lock_ref == 0:
|
|
self.evictable_size_ -= len(node.key)
|
|
self.protected_size_ += len(node.key)
|
|
delta -= len(node.key)
|
|
node.lock_ref += 1
|
|
self._update_leaf_status(node)
|
|
node = node.parent
|
|
return IncLockRefResult(delta=delta)
|
|
|
|
def dec_lock_ref(
|
|
self, node: TreeNode, params: Optional[DecLockRefParams] = None
|
|
) -> DecLockRefResult:
|
|
if self.disable:
|
|
return DecLockRefResult(delta=0)
|
|
|
|
delta = 0
|
|
while node != self.root_node:
|
|
if node.lock_ref == 1:
|
|
self.evictable_size_ += len(node.key)
|
|
self.protected_size_ -= len(node.key)
|
|
delta += len(node.key)
|
|
node.lock_ref -= 1
|
|
self._update_leaf_status(node)
|
|
if node.parent is None:
|
|
assert (
|
|
node is self.root_node
|
|
), f"This request holds the node from another tree"
|
|
node = node.parent
|
|
return DecLockRefResult(delta=delta)
|
|
|
|
def inc_node_lock_ref(self, node: TreeNode):
|
|
"""Increment lock_ref on a single node without walking to root.
|
|
|
|
Used by hicache write-through to protect only the specific node being
|
|
backed up, avoiding unnecessary locking of ancestor nodes. Ancestors
|
|
cannot be evicted anyway while this node has a device value because
|
|
_collect_leaves_device() skips parents with non-evicted children.
|
|
"""
|
|
if self.disable:
|
|
return
|
|
if node == self.root_node:
|
|
return
|
|
if node.lock_ref == 0:
|
|
self.evictable_size_ -= len(node.key)
|
|
self.protected_size_ += len(node.key)
|
|
node.lock_ref += 1
|
|
self._update_leaf_status(node)
|
|
|
|
def dec_node_lock_ref(self, node: TreeNode):
|
|
"""Decrement lock_ref on a single node without walking to root.
|
|
|
|
Counterpart of inc_node_lock_ref for releasing write-through locks.
|
|
"""
|
|
if self.disable:
|
|
return
|
|
if node == self.root_node:
|
|
return
|
|
if node.lock_ref == 1:
|
|
self.evictable_size_ += len(node.key)
|
|
self.protected_size_ -= len(node.key)
|
|
node.lock_ref -= 1
|
|
self._update_leaf_status(node)
|
|
|
|
def evictable_size(self):
|
|
return self.evictable_size_
|
|
|
|
def protected_size(self):
|
|
# protected size refers to the size of the cache that is locked
|
|
return self.protected_size_
|
|
|
|
def all_values_flatten(self):
|
|
values = []
|
|
|
|
def _dfs_helper(node: TreeNode):
|
|
for _, child in node.children.items():
|
|
values.append(child.value)
|
|
_dfs_helper(child)
|
|
|
|
_dfs_helper(self.root_node)
|
|
return torch.cat(values)
|
|
|
|
##### Internal Helper Functions #####
|
|
|
|
def _match_prefix_helper(self, node: TreeNode, key: RadixKey):
|
|
access_time = TreeNode.next_access_time()
|
|
node.last_access_time = access_time
|
|
|
|
child_key = self.get_child_key_fn(key)
|
|
|
|
value = []
|
|
while len(key) > 0 and child_key in node.children.keys():
|
|
child = node.children[child_key]
|
|
child.last_access_time = access_time
|
|
prefix_len = self.key_match_fn(child.key, key)
|
|
if prefix_len < len(child.key):
|
|
new_node = self._split_node(child.key, child, prefix_len)
|
|
value.append(new_node.value)
|
|
node = new_node
|
|
break
|
|
else:
|
|
value.append(child.value)
|
|
node = child
|
|
key = key[prefix_len:]
|
|
|
|
if len(key):
|
|
child_key = self.get_child_key_fn(key)
|
|
|
|
return value, node
|
|
|
|
def _split_node(self, key: RadixKey, child: TreeNode, split_len: int):
|
|
# new_node -> child
|
|
# New node inherits child's priority (represents shared prefix)
|
|
new_node = TreeNode(priority=child.priority)
|
|
new_node.hit_count = child.hit_count
|
|
new_node.children = {self.get_child_key_fn(key[split_len:]): child}
|
|
new_node.parent = child.parent
|
|
new_node.lock_ref = child.lock_ref
|
|
# Node keys must be compacted: a view would pin the backing list of the
|
|
# (possibly transient) key it was sliced from for the node's lifetime.
|
|
new_node.key = child.key[:split_len].compacted()
|
|
new_node.value = child.value[:split_len].clone()
|
|
child.parent = new_node
|
|
child.key = child.key[split_len:].compacted()
|
|
child.value = child.value[split_len:].clone()
|
|
new_node.parent.children[self.get_child_key_fn(key)] = new_node
|
|
|
|
# Split hash_value if it was already computed, otherwise leave as None
|
|
new_node.hash_value, child.hash_value = split_node_hash_value(
|
|
child.hash_value, split_len, self.page_size
|
|
)
|
|
|
|
return new_node
|
|
|
|
def _inc_hit_count(self, node: TreeNode, chunked: bool = False):
|
|
# Skip the hit count update for chunked requests to avoid self-referencing
|
|
# inflation where a chunked request increments hit_count on nodes it created
|
|
# in previous chunks.
|
|
if chunked:
|
|
return
|
|
node.hit_count += 1
|
|
|
|
def _insert_helper(
|
|
self,
|
|
node: TreeNode,
|
|
key: RadixKey,
|
|
value,
|
|
priority: int = 0,
|
|
chunked: bool = False,
|
|
):
|
|
# Convert None priority to 0
|
|
if priority is None:
|
|
priority = 0
|
|
access_time = TreeNode.next_access_time()
|
|
node.last_access_time = access_time
|
|
# Update priority along the path (take max to propagate higher priority)
|
|
node.priority = max(node.priority, priority)
|
|
if len(key) == 0:
|
|
return 0
|
|
|
|
child_key = self.get_child_key_fn(key)
|
|
|
|
total_prefix_length = 0
|
|
while len(key) > 0 and child_key in node.children.keys():
|
|
node = node.children[child_key]
|
|
node.last_access_time = access_time
|
|
prefix_len = self.key_match_fn(node.key, key)
|
|
total_prefix_length += prefix_len
|
|
key = key[prefix_len:]
|
|
value = value[prefix_len:]
|
|
|
|
if prefix_len < len(node.key):
|
|
new_node = self._split_node(node.key, node, prefix_len)
|
|
new_node.priority = max(new_node.priority, priority)
|
|
self._inc_hit_count(new_node, chunked)
|
|
node = new_node
|
|
else:
|
|
node.priority = max(node.priority, priority)
|
|
self._inc_hit_count(node, chunked)
|
|
if len(key):
|
|
child_key = self.get_child_key_fn(key)
|
|
|
|
if len(key):
|
|
new_node = TreeNode(priority=priority)
|
|
new_node.parent = node
|
|
new_node.key = key.compacted()
|
|
new_node.value = value.clone()
|
|
self._inc_hit_count(new_node, chunked)
|
|
node.children[child_key] = new_node
|
|
self.evictable_size_ += len(key)
|
|
self._update_leaf_status(node)
|
|
self._update_leaf_status(new_node)
|
|
# Hash will be computed lazily during event emission
|
|
self._record_store_event(new_node)
|
|
return total_prefix_length
|
|
|
|
def _print_helper(self, node: TreeNode, indent: int):
|
|
"""Prints the radix tree in a human-readable format."""
|
|
stack = [(node, indent)]
|
|
while stack:
|
|
current_node, current_indent = stack.pop()
|
|
print(
|
|
" " * current_indent,
|
|
len(current_node.key),
|
|
current_node.key.token_ids[:10],
|
|
f"r={current_node.lock_ref}",
|
|
)
|
|
for key, child in current_node.children.items():
|
|
stack.append((child, current_indent + 2))
|
|
|
|
assert key == self.get_child_key_fn(
|
|
child.key
|
|
), f"{key=}, {self.get_child_key_fn(child.key)=}"
|
|
|
|
def _delete_leaf(self, node):
|
|
key = self.get_child_key_fn(node.key)
|
|
v = node.parent.children.pop(key, None)
|
|
assert v == node, f"parent does not have child key, {key}"
|
|
|
|
self.evictable_size_ -= len(node.key)
|
|
if node in self.evictable_leaves:
|
|
self.evictable_leaves.remove(node)
|
|
self._update_leaf_status(node.parent)
|
|
|
|
def _update_leaf_status(self, node: TreeNode):
|
|
if node.evicted or node.lock_ref > 0:
|
|
if node in self.evictable_leaves:
|
|
self.evictable_leaves.remove(node)
|
|
return
|
|
|
|
for child in node.children.values():
|
|
if not child.evicted:
|
|
if node in self.evictable_leaves:
|
|
self.evictable_leaves.remove(node)
|
|
return
|
|
|
|
if node not in self.evictable_leaves:
|
|
self.evictable_leaves.add(node)
|
|
|
|
def _total_size_helper(self):
|
|
total_size = 0
|
|
stack = [self.root_node]
|
|
while stack:
|
|
current_node = stack.pop()
|
|
total_size += len(current_node.value)
|
|
for child in current_node.children.values():
|
|
if child.evicted:
|
|
continue
|
|
stack.append(child)
|
|
return total_size
|
|
|
|
def _record_store_event(self, node: TreeNode):
|
|
# One BlockStored per ``page_size`` chunk.
|
|
if self.enable_kv_cache_events:
|
|
# Compute hash_value lazily if not already set
|
|
if node.hash_value is None:
|
|
node.hash_value = compute_node_hash_values(node, self.page_size)
|
|
|
|
# Get parent's last hash value for first page
|
|
parent_block_hash = None
|
|
if node.parent is not None and node.parent != self.root_node:
|
|
if (
|
|
node.parent.hash_value is not None
|
|
and len(node.parent.hash_value) > 0
|
|
):
|
|
parent_block_hash = hash_str_to_int64(node.parent.hash_value[-1])
|
|
|
|
page_index = 0
|
|
for start in range(0, len(node.key), self.page_size):
|
|
page_tokens = node.key.token_ids[start : start + self.page_size]
|
|
if not page_tokens:
|
|
continue
|
|
|
|
block_hash = hash_str_to_int64(node.hash_value[page_index])
|
|
|
|
self.kv_event_queue.append(
|
|
BlockStored(
|
|
block_hashes=[block_hash],
|
|
parent_block_hash=parent_block_hash,
|
|
token_ids=page_tokens,
|
|
block_size=len(page_tokens),
|
|
lora_id=None,
|
|
medium=MEDIUM_GPU,
|
|
)
|
|
)
|
|
|
|
parent_block_hash = block_hash
|
|
page_index += 1
|
|
|
|
def _record_remove_event(self, node: TreeNode):
|
|
# One BlockRemoved per chunk.
|
|
if self.enable_kv_cache_events:
|
|
# Compute hash_value lazily if not already set (must match what was stored)
|
|
if node.hash_value is None:
|
|
node.hash_value = compute_node_hash_values(node, self.page_size)
|
|
|
|
page_index = 0
|
|
for start in range(0, len(node.key), self.page_size):
|
|
page_tokens = node.key.token_ids[start : start + self.page_size]
|
|
if not page_tokens:
|
|
continue
|
|
|
|
block_hash = hash_str_to_int64(node.hash_value[page_index])
|
|
|
|
self.kv_event_queue.append(
|
|
BlockRemoved(block_hashes=[block_hash], medium=MEDIUM_GPU)
|
|
)
|
|
|
|
page_index += 1
|
|
|
|
def _record_all_cleared_event(self):
|
|
if self.enable_kv_cache_events:
|
|
self.kv_event_queue.append(AllBlocksCleared())
|
|
|
|
def take_events(self):
|
|
"""Atomically takes all events and clears the queue.
|
|
|
|
Returns:
|
|
A list of KV cache events.
|
|
"""
|
|
if not self.enable_kv_cache_events:
|
|
return []
|
|
events = self.kv_event_queue
|
|
self.kv_event_queue = []
|
|
return events
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tree = RadixCache.create_simulated()
|
|
|
|
# Example token id sequences (as lists of ints)
|
|
tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3], extra_key=None)))
|
|
tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3], extra_key=None)))
|
|
tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 4, 5], extra_key=None)))
|
|
tree.insert(
|
|
InsertParams(key=RadixKey(token_ids=[1, 2, 4, 5, 6, 7], extra_key=None))
|
|
)
|
|
tree.insert(
|
|
InsertParams(key=RadixKey(token_ids=[8, 9, 10, 11, 12], extra_key=None))
|
|
)
|
|
tree.pretty_print()
|
|
|
|
print(
|
|
tree.match_prefix(
|
|
MatchPrefixParams(key=RadixKey(token_ids=[1, 2, 3, 13, 14], extra_key=None))
|
|
)
|
|
)
|