Unverified Commit 585e1223 authored by Teng Ma's avatar Teng Ma Committed by GitHub
Browse files

[HiCache] feat: add more eviction policy (#11506)

parent a7043c6f
......@@ -21,3 +21,18 @@ class LRUStrategy(EvictionStrategy):
class LFUStrategy(EvictionStrategy):
def get_priority(self, node: "TreeNode") -> Tuple[int, float]:
return (node.hit_count, node.last_access_time)
class FIFOStrategy(EvictionStrategy):
def get_priority(self, node: "TreeNode") -> float:
return node.creation_time
class MRUStrategy(EvictionStrategy):
def get_priority(self, node: "TreeNode") -> float:
return -node.last_access_time
class FILOStrategy(EvictionStrategy):
def get_priority(self, node: "TreeNode") -> float:
return -node.creation_time
......@@ -34,7 +34,14 @@ from sglang.srt.disaggregation.kv_events import (
)
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult
from sglang.srt.mem_cache.evict_policy import EvictionStrategy, LFUStrategy, LRUStrategy
from sglang.srt.mem_cache.evict_policy import (
EvictionStrategy,
FIFOStrategy,
FILOStrategy,
LFUStrategy,
LRUStrategy,
MRUStrategy,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
if TYPE_CHECKING:
......@@ -76,6 +83,7 @@ class TreeNode:
self.value: Optional[torch.Tensor] = None
self.lock_ref = 0
self.last_access_time = time.monotonic()
self.creation_time = time.monotonic()
self.hit_count = 0
# indicating the node is locked to protect from eviction
......@@ -216,9 +224,15 @@ class RadixCache(BasePrefixCache):
self.eviction_strategy: EvictionStrategy = LRUStrategy()
elif eviction_policy.lower() == "lfu":
self.eviction_strategy: EvictionStrategy = LFUStrategy()
elif eviction_policy.lower() == "fifo":
self.eviction_strategy: EvictionStrategy = FIFOStrategy()
elif eviction_policy.lower() == "mru":
self.eviction_strategy: EvictionStrategy = MRUStrategy()
elif eviction_policy.lower() == "filo":
self.eviction_strategy: EvictionStrategy = FILOStrategy()
else:
raise ValueError(
f"Unknown eviction policy: {eviction_policy}. Supported policies: 'lru', 'lfu'."
f"Unknown eviction policy: {eviction_policy}. Supported policies: 'lru', 'lfu', 'fifo', 'mru', 'filo'."
)
self.reset()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment