293 lines
7.9 KiB
Python
293 lines
7.9 KiB
Python
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import time
|
|
from abc import ABC, abstractmethod
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
NamedTuple,
|
|
Optional,
|
|
Protocol,
|
|
Tuple,
|
|
runtime_checkable,
|
|
)
|
|
|
|
import torch
|
|
|
|
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
|
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
|
|
from sglang.srt.observability.metrics_collector import RadixCacheMetricsCollector
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.managers.schedule_batch import Req
|
|
from sglang.srt.mem_cache.radix_cache import RadixKey
|
|
|
|
|
|
@runtime_checkable
|
|
class PrefixCacheTrait(Protocol):
|
|
req_to_token_pool: ReqToTokenPool
|
|
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
|
|
page_size: int
|
|
disable: bool
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class MatchPrefixParams:
|
|
"""Unified parameters for match_prefix across different cache types"""
|
|
|
|
key: RadixKey
|
|
|
|
# Mamba specific
|
|
cow_mamba: bool = False
|
|
req: Optional[Req] = None
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class InsertParams:
|
|
"""Unified parameters for insert across different cache types"""
|
|
|
|
key: Optional[RadixKey] = None
|
|
value: Optional[torch.Tensor] = None
|
|
|
|
# Mamba specific
|
|
mamba_value: Optional[torch.Tensor] = None
|
|
|
|
# SWA specific
|
|
prev_prefix_len: int = 0
|
|
swa_evicted_seqlen: int = 0
|
|
|
|
# General
|
|
chunked: bool = False
|
|
priority: int = 0
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class InsertResult:
|
|
"""Result of an insert operation"""
|
|
|
|
prefix_len: int
|
|
mamba_exist: bool = False
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class EvictParams:
|
|
"""Unified parameters for evict across different cache types"""
|
|
|
|
num_tokens: int
|
|
swa_num_tokens: int = 0
|
|
mamba_num: int = 0
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class EvictResult:
|
|
"""Result of an evict operation"""
|
|
|
|
num_tokens_evicted: int = 0
|
|
swa_num_tokens_evicted: int = 0
|
|
mamba_num_evicted: int = 0
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class IncLockRefResult:
|
|
"""Result of an inc_lock_ref operation."""
|
|
|
|
delta: Optional[int] = None
|
|
swa_uuid_for_lock: Optional[int] = None
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class DecLockRefParams:
|
|
"""Parameters for dec_lock_ref operation."""
|
|
|
|
swa_uuid_for_lock: Optional[int] = None
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class DecLockRefResult:
|
|
"""Result of an dec_lock_ref operation."""
|
|
|
|
delta: Optional[int] = None
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class InitLoadBackParams:
|
|
"""Unified parameters for init_load_back across different cache types"""
|
|
|
|
last_host_node: Any
|
|
host_hit_length: int
|
|
mem_quota: Optional[int] = None
|
|
req: Optional[Req] = None
|
|
|
|
|
|
class MatchResult(NamedTuple):
|
|
"""Result of a prefix match operation.
|
|
|
|
Attributes:
|
|
device_indices : Indices of the KV cache on the device matched by common prefix.
|
|
last_device_node: The last TreeNode on the device that was matched.
|
|
last_host_node : The last TreeNode on the host that was matched.
|
|
Note that if HiCache is not enabled,
|
|
this **must** be the same as `last_device_node`.
|
|
host_hit_length : Length of the host cache hit. For pure-KV caches this is the
|
|
number of evicted KV tokens on CPU. For hybrid Mamba models this
|
|
is max(kv_host_tokens, 1-if-mamba-on-host) so that a mamba-only
|
|
host hit still triggers load-back without adding a separate field.
|
|
0 if HiCache is not enabled.
|
|
mamba_branching_seqlen: The mamba radix cache branching point, which is the longest
|
|
page-aligned position that could've been cache hit if there
|
|
exists a mamba state.
|
|
"""
|
|
|
|
device_indices: torch.Tensor
|
|
last_device_node: Any
|
|
last_host_node: Any
|
|
host_hit_length: int = 0
|
|
mamba_branching_seqlen: Optional[int] = None
|
|
cache_protected_len: Optional[int] = None
|
|
|
|
|
|
class BasePrefixCache(ABC, PrefixCacheTrait):
|
|
"""Cache can be indexed by either rid or key."""
|
|
|
|
metrics_collector: Optional[RadixCacheMetricsCollector] = (
|
|
None # metrics collector for the cache
|
|
)
|
|
|
|
def init_metrics_collector(self):
|
|
from sglang.srt.server_args import get_global_server_args
|
|
|
|
server_args = get_global_server_args()
|
|
labels = {"cache_type": self.__class__.__name__}
|
|
if server_args.extra_metric_labels:
|
|
labels.update(server_args.extra_metric_labels)
|
|
self.metrics_collector = RadixCacheMetricsCollector(labels=labels)
|
|
|
|
def update_eviction_metrics(self, num_evicted: int, start_time: float):
|
|
if self.metrics_collector is not None and num_evicted > 0:
|
|
self.metrics_collector.observe_eviction_duration(
|
|
time.perf_counter() - start_time
|
|
)
|
|
self.metrics_collector.increment_eviction_num_tokens(num_evicted)
|
|
|
|
@abstractmethod
|
|
def reset(self):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def cache_unfinished_req(self, req: Req, **kwargs):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def evict(self, params: EvictParams) -> EvictResult:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def inc_lock_ref(self, node: Any) -> IncLockRefResult:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def dec_lock_ref(
|
|
self, node: Any, params: Optional[DecLockRefParams] = None
|
|
) -> DecLockRefResult:
|
|
pass
|
|
|
|
def evictable_size(self):
|
|
return 0
|
|
|
|
def full_evictable_size(self):
|
|
return 0
|
|
|
|
def swa_evictable_size(self):
|
|
return 0
|
|
|
|
def protected_size(self):
|
|
return 0
|
|
|
|
def full_protected_size(self):
|
|
return 0
|
|
|
|
def swa_protected_size(self):
|
|
return 0
|
|
|
|
def total_size(self):
|
|
raise NotImplementedError()
|
|
|
|
def pretty_print(self):
|
|
raise NotImplementedError()
|
|
|
|
def init_load_back(
|
|
self,
|
|
params: InitLoadBackParams,
|
|
) -> Tuple[torch.Tensor, Any]:
|
|
"""
|
|
Preparing KV cache loading from host to device.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def ready_to_load_host_cache(self) -> Any:
|
|
"""
|
|
Notify the cache controller to start the KV cache loading
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def flush_write_through_acks(self) -> None:
|
|
"""Release lock_ref on radix-tree nodes whose write-through has completed.
|
|
|
|
Lightweight operation that only processes finished write acks.
|
|
No-op for caches without hierarchical write-through support.
|
|
"""
|
|
pass
|
|
|
|
def check_hicache_events(self) -> Any:
|
|
"""
|
|
Check HiCache related activities to update radix tree and synchronize across TP workers if needed
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def take_events(self):
|
|
return []
|
|
|
|
def supports_swa(self) -> bool:
|
|
return False
|
|
|
|
def supports_mamba(self) -> bool:
|
|
return False
|
|
|
|
def supports_streaming_session(self) -> bool:
|
|
return False
|
|
|
|
def release_session(self, session_id: str) -> None:
|
|
pass
|
|
|
|
def session_held_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
|
|
return 0
|
|
|
|
def session_held_full_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
|
|
return 0
|
|
|
|
def session_held_swa_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
|
|
return 0
|
|
|
|
def session_held_req_count(self, active_pool_idxs: Optional[set] = None) -> int:
|
|
return 0
|
|
|
|
def is_chunk_cache(self) -> bool:
|
|
return False
|
|
|
|
def is_tree_cache(self) -> bool:
|
|
return not self.is_chunk_cache()
|
|
|
|
def available_and_evictable_str(self) -> str:
|
|
available_size = self.token_to_kv_pool_allocator.available_size()
|
|
evictable_size = self.evictable_size()
|
|
return f"Available tokens: {available_size + evictable_size} ({available_size=} + {evictable_size=})\n"
|