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

Refactor kv cache event mixin into a recorder (#35164)

Browse Source
This commit is contained in:
Ke Bao
2026-08-19 01:20:17 +08:00
committed by GitHub
parent 83d7d45330
commit 480033def0
12 changed files with 207 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/base_prefix_cache.py
+4 -1
View File
@@ -17,6 +17,7 @@ from typing import (
import torch import torch
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.events import KVCacheEventRecorder
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.observability.metrics_collector import ( from sglang.srt.observability.metrics_collector import (
STAT_LOGGER_ROLE_RADIX_CACHE, STAT_LOGGER_ROLE_RADIX_CACHE,
@@ -238,6 +239,8 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
None # metrics collector for the cache None # metrics collector for the cache
) )
cache_controller: Optional[HiCacheController] = None cache_controller: Optional[HiCacheController] = None
# Set by caches that publish KV placement events; None means they don't.
kv_events: Optional[KVCacheEventRecorder] = None
def init_metrics_collector(self): def init_metrics_collector(self):
from sglang.srt.runtime_context import get_server_args from sglang.srt.runtime_context import get_server_args
@@ -377,7 +380,7 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
raise NotImplementedError() raise NotImplementedError()
def take_events(self): def take_events(self):
return [] return [] if self.kv_events is None else self.kv_events.take()
def supports_swa(self) -> bool: def supports_swa(self) -> bool:
return False return False
python/sglang/srt/mem_cache/events.py
+68 -51
View File
@@ -11,13 +11,14 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
# ============================================================================== # ==============================================================================
"""KV cache placement event emission mixin. """KV cache placement event recording.
The mixin produces the ``BlockStored`` / ``BlockRemoved`` / ``AllBlocksCleared`` Produces the ``BlockStored`` / ``BlockRemoved`` / ``AllBlocksCleared`` events
events consumed by KV-aware routers (e.g. dynamo). consumed by KV-aware routers (e.g. dynamo). A cache holds one recorder and calls
it; the recorder owns the queue and needs nothing back from its owner.
""" """
from typing import Any from typing import Any, Optional
from sglang.srt.disaggregation.kv_events import ( from sglang.srt.disaggregation.kv_events import (
AllBlocksCleared, AllBlocksCleared,
@@ -34,16 +35,27 @@ from sglang.srt.mem_cache.utils import (
) )
class KVCacheEventMixin: class KVCacheEventRecorder:
def _enqueue_kv_event(self, event): """Collects KV placement events for one cache.
``enabled=False`` makes every ``record_*`` call a no-op and ``take`` return an
empty list, so callers never have to guard.
"""
def __init__(self, *, enabled: bool, page_size: int):
self.enabled = enabled
self.page_size = page_size
self._queue: list = []
def enqueue(self, event) -> None:
"""Append an event, coalescing it with a compatible queue tail. """Append an event, coalescing it with a compatible queue tail.
KV event batches already support multiple block hashes. Combining them KV event batches already support multiple block hashes. Combining them
here avoids emitting one event per page while preserving ordering and here avoids emitting one event per page while preserving ordering and
the parent-linked store chains consumers use to rebuild the cache tree. the parent-linked store chains consumers use to rebuild the cache tree.
""" """
if self.kv_event_queue: if self._queue:
tail = self.kv_event_queue[-1] tail = self._queue[-1]
if isinstance(tail, BlockRemoved) and isinstance(event, BlockRemoved): if isinstance(tail, BlockRemoved) and isinstance(event, BlockRemoved):
if tail.medium == event.medium: if tail.medium == event.medium:
@@ -71,35 +83,46 @@ class KVCacheEventMixin:
tail.token_ids.extend(event.token_ids) tail.token_ids.extend(event.token_ids)
return return
self.kv_event_queue.append(event) self._queue.append(event)
def _record_store_event(self, node: Any, medium=None): def _node_event_hash_values(self, node: Any) -> list:
"""Hash values to publish for ``node``, computing them if not yet set."""
if node.hash_value is None:
node.hash_value = compute_node_hash_values(node, self.page_size)
if node.key.cache_salt is not None:
return compute_node_event_hash_values(node, self.page_size)
return node.hash_value
def _parent_block_hash(self, node: Any) -> Optional[int]:
"""The hash the first page of ``node`` links back to.
``None`` when the parent is the tree root: a root carries an empty
``hash_value`` and no event hash, so it contributes no link. Every other
node on the path has a parent, which is what distinguishes the two.
"""
parent = node.parent
if parent is None or parent.parent is None:
return None
if node.key.cache_salt is not None:
parent_hash_values = parent.event_hash_value
assert parent_hash_values is not None
else:
parent_hash_values = parent.hash_value
if not parent_hash_values:
return None
return hash_str_to_int64(parent_hash_values[-1])
def record_store(self, node: Any, medium=None) -> None:
# One BlockStored per ``page_size`` chunk. # One BlockStored per ``page_size`` chunk.
# ``medium`` defaults to StorageMedium.GPU but callers may override # ``medium`` defaults to StorageMedium.GPU but callers may override
# for lower-tier insertions (e.g. StorageMedium.CPU for host/L2 cache). # for lower-tier insertions (e.g. StorageMedium.CPU for host/L2 cache).
if self.enable_kv_cache_events: if not self.enabled:
return
if medium is None: if medium is None:
medium = StorageMedium.GPU medium = StorageMedium.GPU
# Compute hash_value lazily if not already set event_hash_values = self._node_event_hash_values(node)
if node.hash_value is None: parent_block_hash = self._parent_block_hash(node)
node.hash_value = compute_node_hash_values(node, self.page_size)
event_hash_values = (
compute_node_event_hash_values(node, self.page_size)
if node.key.cache_salt is not None
else node.hash_value
)
# 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.key.cache_salt is not None:
parent_hash_values = node.parent.event_hash_value
assert parent_hash_values is not None
else:
parent_hash_values = node.parent.hash_value
if parent_hash_values:
parent_block_hash = hash_str_to_int64(parent_hash_values[-1])
page_index = 0 page_index = 0
logical_len = len(node.key) logical_len = len(node.key)
@@ -132,27 +155,22 @@ class KVCacheEventMixin:
**event_args, **event_args,
metadata=BlockStoredMetadata(cache_salt=node.key.cache_salt), metadata=BlockStoredMetadata(cache_salt=node.key.cache_salt),
) )
self._enqueue_kv_event(event) self.enqueue(event)
parent_block_hash = block_hash parent_block_hash = block_hash
page_index += 1 page_index += 1
def _record_remove_event(self, node: Any, medium=None): def record_remove(self, node: Any, medium=None) -> None:
# One BlockRemoved per radix node. # One BlockRemoved per radix node.
# ``medium`` defaults to StorageMedium.GPU but callers may override for # ``medium`` defaults to StorageMedium.GPU but callers may override for
# lower-tier removals (e.g. StorageMedium.CPU when evicting from host). # lower-tier removals (e.g. StorageMedium.CPU when evicting from host).
if self.enable_kv_cache_events: if not self.enabled:
return
if medium is None: if medium is None:
medium = StorageMedium.GPU medium = StorageMedium.GPU
# Compute hash_value lazily if not already set (must match what was stored) # Hash values must match what was stored.
if node.hash_value is None: event_hash_values = self._node_event_hash_values(node)
node.hash_value = compute_node_hash_values(node, self.page_size)
event_hash_values = (
compute_node_event_hash_values(node, self.page_size)
if node.key.cache_salt is not None
else node.hash_value
)
block_hashes = [] block_hashes = []
logical_len = len(node.key) logical_len = len(node.key)
@@ -166,22 +184,21 @@ class KVCacheEventMixin:
page_index += 1 page_index += 1
if block_hashes: if block_hashes:
self._enqueue_kv_event( self.enqueue(BlockRemoved(block_hashes=block_hashes, medium=medium))
BlockRemoved(block_hashes=block_hashes, medium=medium)
)
def _record_all_cleared_event(self): def record_all_cleared(self) -> None:
if self.enable_kv_cache_events: if not self.enabled:
self._enqueue_kv_event(AllBlocksCleared()) return
self.enqueue(AllBlocksCleared())
def take_events(self): def take(self) -> list:
"""Atomically takes all events and clears the queue. """Atomically takes all events and clears the queue.
Returns: Returns:
A list of KV cache events. A list of KV cache events.
""" """
if not self.enable_kv_cache_events: if not self.enabled:
return [] return []
events = self.kv_event_queue events = self._queue
self.kv_event_queue = [] self._queue = []
return events return events
python/sglang/srt/mem_cache/hiradix_cache.py
+10 -10
View File
@@ -908,7 +908,7 @@ class HiRadixCache(RadixCache):
if node.write_through_pending_id == ack_id: if node.write_through_pending_id == ack_id:
node.write_through_pending_id = None node.write_through_pending_id = None
# DMA confirmed -- block is now on host. # DMA confirmed -- block is now on host.
self._record_store_event(node, medium=StorageMedium.CPU) self.kv_events.record_store(node, medium=StorageMedium.CPU)
if self.enable_storage: if self.enable_storage:
self.write_backup_storage(lock_node, backup_len) self.write_backup_storage(lock_node, backup_len)
if release_lock: if release_lock:
@@ -1264,7 +1264,7 @@ class HiRadixCache(RadixCache):
def _detach_backuped(self, node: TreeNode) -> int: def _detach_backuped(self, node: TreeNode) -> int:
# detach nodes from tree while keeping device slots, for write-back eviction # detach nodes from tree while keeping device slots, for write-back eviction
self._record_remove_event(node, medium=StorageMedium.GPU) self.kv_events.record_remove(node, medium=StorageMedium.GPU)
num_evicted = len(node.value) num_evicted = len(node.value)
assert num_evicted > 0 assert num_evicted > 0
self.evictable_size_ -= num_evicted self.evictable_size_ -= num_evicted
@@ -1285,7 +1285,7 @@ class HiRadixCache(RadixCache):
# evict a node not initiated write to host -- emit BlockRemoved # evict a node not initiated write to host -- emit BlockRemoved
assert len(node.children) == 0, f"non-leaf, {node.id=}" assert len(node.children) == 0, f"non-leaf, {node.id=}"
self._record_remove_event(node) self.kv_events.record_remove(node)
self.cache_controller.mem_pool_device_allocator.free(node.value) self.cache_controller.mem_pool_device_allocator.free(node.value)
num_evicted = len(node.value) num_evicted = len(node.value)
self._delete_leaf(node) self._delete_leaf(node)
@@ -1311,11 +1311,11 @@ class HiRadixCache(RadixCache):
freed_device = 0 freed_device = 0
for n in nodes: for n in nodes:
if n.host_value is not None: if n.host_value is not None:
self._record_remove_event(n, medium=StorageMedium.CPU) self.kv_events.record_remove(n, medium=StorageMedium.CPU)
self.cache_controller.evict_host(n.host_value) self.cache_controller.evict_host(n.host_value)
n.host_value = None n.host_value = None
if n.value is not None: if n.value is not None:
self._record_remove_event(n, medium=StorageMedium.GPU) self.kv_events.record_remove(n, medium=StorageMedium.GPU)
self.cache_controller.mem_pool_device_allocator.free(n.value) self.cache_controller.mem_pool_device_allocator.free(n.value)
freed_device += len(n.value) freed_device += len(n.value)
self.evictable_size_ -= len(n.value) self.evictable_size_ -= len(n.value)
@@ -1357,7 +1357,7 @@ class HiRadixCache(RadixCache):
# Block deleted entirely (GPU already evicted, now CPU freed) -- # Block deleted entirely (GPU already evicted, now CPU freed) --
# emit remove(CPU) so the router drops the host-tier entry. # emit remove(CPU) so the router drops the host-tier entry.
self._record_remove_event(x, medium=StorageMedium.CPU) self.kv_events.record_remove(x, medium=StorageMedium.CPU)
num_evicted += self.cache_controller.evict_host(x.host_value) num_evicted += self.cache_controller.evict_host(x.host_value)
key = x.key.child_key(self.page_size) key = x.key.child_key(self.page_size)
@@ -1438,7 +1438,7 @@ class HiRadixCache(RadixCache):
offset += len(node.host_value) offset += len(node.host_value)
# Block promoted from host to GPU -- emit store(GPU) so downstream # Block promoted from host to GPU -- emit store(GPU) so downstream
# indexers see it as device-local again. # indexers see it as device-local again.
self._record_store_event(node, medium=StorageMedium.GPU) self.kv_events.record_store(node, medium=StorageMedium.GPU)
self.evictable_size_ += len(device_indices) self.evictable_size_ += len(device_indices)
self.inc_lock_ref(last_hit_node) self.inc_lock_ref(last_hit_node)
@@ -1849,7 +1849,7 @@ class HiRadixCache(RadixCache):
self._update_host_leaf_status(node) self._update_host_leaf_status(node)
# Publish the newly materialized host suffix immediately so downstream # Publish the newly materialized host suffix immediately so downstream
# cache indexers can resolve descendants that extend this L2-only prefix. # cache indexers can resolve descendants that extend this L2-only prefix.
self._record_store_event(new_node, medium=StorageMedium.CPU) self.kv_events.record_store(new_node, medium=StorageMedium.CPU)
return matched_length return matched_length
@@ -1987,11 +1987,11 @@ class HiRadixCache(RadixCache):
self._update_leaf_status(new_node) self._update_leaf_status(new_node)
# Compute hash_value if storage or kv events are enabled # Compute hash_value if storage or kv events are enabled
if self.enable_storage or self.enable_kv_cache_events: if self.enable_storage or self.kv_events.enabled:
new_node.hash_value = compute_node_hash_values(new_node, self.page_size) new_node.hash_value = compute_node_hash_values(new_node, self.page_size)
# Emit BlockStored so the router indexes this block. # Emit BlockStored so the router indexes this block.
self._record_store_event(new_node) self.kv_events.record_store(new_node)
if self.cache_controller.write_policy != "write_back": if self.cache_controller.write_policy != "write_back":
self._inc_hit_count(new_node, chunked) self._inc_hit_count(new_node, chunked)
python/sglang/srt/mem_cache/mamba_radix_cache.py
+9 -8
View File
@@ -43,7 +43,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchPrefixParams, MatchPrefixParams,
MatchResult, MatchResult,
) )
from sglang.srt.mem_cache.events import KVCacheEventMixin from sglang.srt.mem_cache.events import KVCacheEventRecorder
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
from sglang.srt.mem_cache.multi_ended_allocator import ( from sglang.srt.mem_cache.multi_ended_allocator import (
UnifiedMambaTokenToKVPoolAllocator, UnifiedMambaTokenToKVPoolAllocator,
@@ -441,7 +441,7 @@ class LRUList:
raise Exception(msg) raise Exception(msg)
class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): class MambaRadixCache(BasePrefixCache):
def __init__(self, params: CacheInitParams): def __init__(self, params: CacheInitParams):
assert ( assert (
isinstance(params.token_to_kv_pool_allocator, TokenToKVPoolAllocator) isinstance(params.token_to_kv_pool_allocator, TokenToKVPoolAllocator)
@@ -458,10 +458,11 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
self.page_size = params.page_size self.page_size = params.page_size
self.disable = params.disable self.disable = params.disable
self.enable_kv_cache_events = params.enable_kv_cache_events
self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer
self.enable_mamba_extra_buffer_lazy = params.enable_mamba_extra_buffer_lazy self.enable_mamba_extra_buffer_lazy = params.enable_mamba_extra_buffer_lazy
self.kv_event_queue = [] self.kv_events = KVCacheEventRecorder(
enabled=params.enable_kv_cache_events, page_size=self.page_size
)
if not self.enable_mamba_extra_buffer: if not self.enable_mamba_extra_buffer:
assert ( assert (
@@ -497,7 +498,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
# LRU lists are used to maintain the order of eviction of the nodes in the tree # LRU lists are used to maintain the order of eviction of the nodes in the tree
self.full_lru_list = LRUList(mamba=False) self.full_lru_list = LRUList(mamba=False)
self.mamba_lru_list = LRUList(mamba=True) self.mamba_lru_list = LRUList(mamba=True)
self._record_all_cleared_event() self.kv_events.record_all_cleared()
def match_prefix(self, params: MatchPrefixParams) -> MatchResult: def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
"""Find the matching prefix from the radix tree. """Find the matching prefix from the radix tree.
@@ -819,7 +820,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
assert x.mamba_value is not None, f"leaf node mamba value is not None, {x.id=}" assert x.mamba_value is not None, f"leaf node mamba value is not None, {x.id=}"
# 1. a leaf node, free full tokens and mamba # 1. a leaf node, free full tokens and mamba
self._record_remove_event(x) self.kv_events.record_remove(x)
# Tree values are page-aligned copies of a kv row: page-exact segment. # Tree values are page-aligned copies of a kv row: page-exact segment.
self.token_to_kv_pool_allocator.free_segment(x.value, start_pos=0) self.token_to_kv_pool_allocator.free_segment(x.value, start_pos=0)
full_num_evicted = len(x.value) full_num_evicted = len(x.value)
@@ -1301,7 +1302,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
node.children[child_key] = new_node node.children[child_key] = new_node
self.full_evictable_size_ += len(value) self.full_evictable_size_ += len(value)
self.mamba_evictable_size_ += len(mamba_value) self.mamba_evictable_size_ += len(mamba_value)
self._record_store_event(new_node) self.kv_events.record_store(new_node)
elif node.mamba_value is None: # add for mamba tombstone elif node.mamba_value is None: # add for mamba tombstone
node.mamba_value = mamba_value node.mamba_value = mamba_value
self.full_lru_list.reset_node_mru(node) self.full_lru_list.reset_node_mru(node)
@@ -1330,7 +1331,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
node.parent.mamba_lock_ref == 0 node.parent.mamba_lock_ref == 0
), f"tombstone mamba_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.mamba_lock_ref=}, {node.parent.id=}" ), f"tombstone mamba_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.mamba_lock_ref=}, {node.parent.id=}"
# delete tombstone node evicts full tokens # delete tombstone node evicts full tokens
self._record_remove_event(node.parent) self.kv_events.record_remove(node.parent)
self.token_to_kv_pool_allocator.free_segment(node.parent.value, start_pos=0) self.token_to_kv_pool_allocator.free_segment(node.parent.value, start_pos=0)
full_num_evicted += len(node.parent.value) full_num_evicted += len(node.parent.value)
self.full_lru_list.remove_node(node.parent) self.full_lru_list.remove_node(node.parent)
python/sglang/srt/mem_cache/radix_cache.py
+8 -7
View File
@@ -45,7 +45,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchPrefixParams, MatchPrefixParams,
MatchResult, MatchResult,
) )
from sglang.srt.mem_cache.events import KVCacheEventMixin from sglang.srt.mem_cache.events import KVCacheEventRecorder
from sglang.srt.mem_cache.utils import ( from sglang.srt.mem_cache.utils import (
get_eviction_strategy, get_eviction_strategy,
get_hash_str, get_hash_str,
@@ -300,18 +300,19 @@ class TreeNode:
return self.last_access_time < other.last_access_time return self.last_access_time < other.last_access_time
class RadixCache(KVCacheEventMixin, BasePrefixCache): class RadixCache(BasePrefixCache):
def __init__(self, params: CacheInitParams): def __init__(self, params: CacheInitParams):
self.disable = params.disable self.disable = params.disable
self.req_to_token_pool = params.req_to_token_pool self.req_to_token_pool = params.req_to_token_pool
self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
self.page_size = params.page_size self.page_size = params.page_size
self.enable_kv_cache_events = params.enable_kv_cache_events
self.is_eagle = params.is_eagle self.is_eagle = params.is_eagle
self.disable_finished_insert = params.disable_finished_insert self.disable_finished_insert = params.disable_finished_insert
self.eviction_policy = params.eviction_policy.lower() self.eviction_policy = params.eviction_policy.lower()
self.kv_event_queue = [] self.kv_events = KVCacheEventRecorder(
enabled=params.enable_kv_cache_events, page_size=self.page_size
)
if params.enable_metrics: if params.enable_metrics:
self.init_metrics_collector() self.init_metrics_collector()
@@ -371,7 +372,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
last_host_node=self.root_node, last_host_node=self.root_node,
best_match_node=self.root_node, best_match_node=self.root_node,
) )
self._record_all_cleared_event() self.kv_events.record_all_cleared()
def match_prefix(self, params: MatchPrefixParams) -> MatchResult: def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
"""Find the longest cached prefix of ``key`` in the radix tree. """Find the longest cached prefix of ``key`` in the radix tree.
@@ -614,7 +615,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
new_priority = self.eviction_strategy.get_priority(x.parent) new_priority = self.eviction_strategy.get_priority(x.parent)
heapq.heappush(eviction_heap, (new_priority, x.parent)) heapq.heappush(eviction_heap, (new_priority, x.parent))
self._record_remove_event(x) self.kv_events.record_remove(x)
self.update_eviction_metrics(num_evicted, start_time) self.update_eviction_metrics(num_evicted, start_time)
return EvictResult(num_tokens_evicted=num_evicted) return EvictResult(num_tokens_evicted=num_evicted)
@@ -785,7 +786,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
self._update_leaf_status(node) self._update_leaf_status(node)
self._update_leaf_status(new_node) self._update_leaf_status(new_node)
# Hash will be computed lazily during event emission # Hash will be computed lazily during event emission
self._record_store_event(new_node) self.kv_events.record_store(new_node)
node = new_node node = new_node
return total_prefix_length, node return total_prefix_length, node
python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py
+2 -2
View File
@@ -375,8 +375,8 @@ class FlexKVRadixCache(RadixCache):
self._update_leaf_status(last_node) self._update_leaf_status(last_node)
self._update_leaf_status(new_node) self._update_leaf_status(new_node)
self._record_store_event(new_node.parent) self.kv_events.record_store(new_node.parent)
self._record_store_event(new_node) self.kv_events.record_store(new_node)
return fetched_slots, new_node return fetched_slots, new_node
python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py
+2 -2
View File
@@ -383,8 +383,8 @@ class LMCRadixCache(RadixCache):
self._update_leaf_status(last_node) self._update_leaf_status(last_node)
self._update_leaf_status(new_node) self._update_leaf_status(new_node)
self._record_store_event(new_node.parent) self.kv_events.record_store(new_node.parent)
self._record_store_event(new_node) self.kv_events.record_store(new_node)
return token_slots[:fetched], new_node return token_slots[:fetched], new_node
python/sglang/srt/mem_cache/swa_radix_cache.py
+10 -9
View File
@@ -42,7 +42,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchResult, MatchResult,
) )
from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.events import KVCacheEventMixin from sglang.srt.mem_cache.events import KVCacheEventRecorder
from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.utils import split_node_hash_value from sglang.srt.mem_cache.utils import split_node_hash_value
@@ -342,7 +342,7 @@ class LRUList:
raise Exception(msg) raise Exception(msg)
class SWARadixCache(KVCacheEventMixin, BasePrefixCache): class SWARadixCache(BasePrefixCache):
def __init__(self, params: CacheInitParams): def __init__(self, params: CacheInitParams):
assert isinstance(params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator) assert isinstance(params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator)
self.req_to_token_pool = params.req_to_token_pool self.req_to_token_pool = params.req_to_token_pool
@@ -350,8 +350,9 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
self.page_size = params.page_size self.page_size = params.page_size
self.disable = params.disable self.disable = params.disable
self.is_eagle = params.is_eagle self.is_eagle = params.is_eagle
self.enable_kv_cache_events = params.enable_kv_cache_events self.kv_events = KVCacheEventRecorder(
self.kv_event_queue = [] enabled=params.enable_kv_cache_events, page_size=self.page_size
)
if self.token_to_kv_pool_allocator: if self.token_to_kv_pool_allocator:
self.device = self.token_to_kv_pool_allocator.device self.device = self.token_to_kv_pool_allocator.device
@@ -407,7 +408,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
# LRU lists are used to maintain the order of eviction of the nodes in the tree # LRU lists are used to maintain the order of eviction of the nodes in the tree
self.full_lru_list = LRUList(is_swa_list=False) self.full_lru_list = LRUList(is_swa_list=False)
self.swa_lru_list = LRUList(is_swa_list=True) self.swa_lru_list = LRUList(is_swa_list=True)
self._record_all_cleared_event() self.kv_events.record_all_cleared()
def match_prefix(self, params: MatchPrefixParams) -> MatchResult: def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
"""Find the matching prefix from the radix tree. """Find the matching prefix from the radix tree.
@@ -609,7 +610,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
assert x.full_lock_ref == 0, f"node is in use, {x.id=}" assert x.full_lock_ref == 0, f"node is in use, {x.id=}"
# 1. free node kv indices, evict full and swa tokens # 1. free node kv indices, evict full and swa tokens
self._record_remove_event(x) self.kv_events.record_remove(x)
self.token_to_kv_pool_allocator.free(x.value) self.token_to_kv_pool_allocator.free(x.value)
full_num_evicted += len(x.value) full_num_evicted += len(x.value)
# Tombstoned leaves had their SWA freed earlier in `dec_swa_lock_only` # Tombstoned leaves had their SWA freed earlier in `dec_swa_lock_only`
@@ -674,7 +675,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
x.full_lock_ref == 0 x.full_lock_ref == 0
), f"leaf node with full lock must also have swa lock, {x.id=}" ), f"leaf node with full lock must also have swa lock, {x.id=}"
# 1. a leaf node, free full and swa tokens # 1. a leaf node, free full and swa tokens
self._record_remove_event(x) self.kv_events.record_remove(x)
self.token_to_kv_pool_allocator.free(x.value) self.token_to_kv_pool_allocator.free(x.value)
full_num_evicted += len(x.value) full_num_evicted += len(x.value)
swa_num_evicted += len(x.value) swa_num_evicted += len(x.value)
@@ -1326,7 +1327,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
if not swa_tombstone: if not swa_tombstone:
self.swa_lru_list.insert_mru(new_node) self.swa_lru_list.insert_mru(new_node)
self.swa_evictable_size_ += len(value) self.swa_evictable_size_ += len(value)
self._record_store_event(new_node) self.kv_events.record_store(new_node)
return new_node return new_node
def _iteratively_delete_tombstone_leaf( def _iteratively_delete_tombstone_leaf(
@@ -1344,7 +1345,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
node.parent.swa_lock_ref == 0 node.parent.swa_lock_ref == 0
), f"tombstone swa_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.swa_lock_ref=}, {node.parent.id=}" ), f"tombstone swa_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.swa_lock_ref=}, {node.parent.id=}"
# delete tombstone node evicts full tokens # delete tombstone node evicts full tokens
self._record_remove_event(node.parent) self.kv_events.record_remove(node.parent)
self.token_to_kv_pool_allocator.free(node.parent.value) self.token_to_kv_pool_allocator.free(node.parent.value)
full_num_evicted += len(node.parent.value) full_num_evicted += len(node.parent.value)
self.full_lru_list.remove_node(node.parent) self.full_lru_list.remove_node(node.parent)
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+12 -10
View File
@@ -36,6 +36,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchPrefixParams, MatchPrefixParams,
MatchResult, MatchResult,
) )
from sglang.srt.mem_cache.events import KVCacheEventRecorder
from sglang.srt.mem_cache.hicache_storage import ( from sglang.srt.mem_cache.hicache_storage import (
PoolName, PoolName,
PoolTransfer, PoolTransfer,
@@ -417,8 +418,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self.components_by_type.values() self.components_by_type.values()
) )
self.enable_kv_cache_events = params.enable_kv_cache_events self.kv_events = KVCacheEventRecorder(
self.kv_event_queue = [] enabled=params.enable_kv_cache_events, page_size=self.page_size
)
self.reset() self.reset()
@@ -1119,7 +1121,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self._update_evictable_leaf_sets(new_node) self._update_evictable_leaf_sets(new_node)
self._update_evictable_leaf_sets(parent) self._update_evictable_leaf_sets(parent)
self._record_store_event(new_node) self.kv_events.record_store(new_node)
return new_node return new_node
def _unevict_node_on_insert( def _unevict_node_on_insert(
@@ -1138,7 +1140,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self._update_duplicate_tracking(node) self._update_duplicate_tracking(node)
if node.parent is not None: if node.parent is not None:
self._update_evictable_leaf_sets(node.parent) self._update_evictable_leaf_sets(node.parent)
self._record_store_event(node, medium=StorageMedium.GPU) self.kv_events.record_store(node, medium=StorageMedium.GPU)
def _update_evictable_leaf_sets(self, node: UnifiedTreeNode) -> None: def _update_evictable_leaf_sets(self, node: UnifiedTreeNode) -> None:
"""Update both device and host leaf sets for a node.""" """Update both device and host leaf sets for a node."""
@@ -1303,7 +1305,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
) -> None: ) -> None:
"""Free every component layer on the node and detach it from the LRU """Free every component layer on the node and detach it from the LRU
lists and evictable leaf sets.""" lists and evictable leaf sets."""
self._record_remove_event(node, medium=medium) self.kv_events.record_remove(node, medium=medium)
for comp in self.components: for comp in self.components:
self._evict_component_and_detach_lru( self._evict_component_and_detach_lru(
node, node,
@@ -1422,7 +1424,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
"""Free only the Full host layer; aux host slices stay under their own """Free only the Full host layer; aux host slices stay under their own
pools' LRU (a host-only aux slice may be a sole copy).""" pools' LRU (a host-only aux slice may be a sole copy)."""
assert self._can_reclaim_full_host_duplicate(node) assert self._can_reclaim_full_host_duplicate(node)
self._record_remove_event(node, medium=StorageMedium.CPU) self.kv_events.record_remove(node, medium=StorageMedium.CPU)
self._evict_component_and_detach_lru( self._evict_component_and_detach_lru(
node, node,
self.components_by_type[BASE_COMPONENT_TYPE], self.components_by_type[BASE_COMPONENT_TYPE],
@@ -1447,7 +1449,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
All freed tokens are accumulated into *tracker*.""" All freed tokens are accumulated into *tracker*."""
assert self._is_host_leaf(node), f"node {node.id} is not an H-leaf" assert self._is_host_leaf(node), f"node {node.id} is not an H-leaf"
self._record_remove_event(node, medium=StorageMedium.CPU) self.kv_events.record_remove(node, medium=StorageMedium.CPU)
for comp in self.components: for comp in self.components:
_, hf = self._evict_component_and_detach_lru( _, hf = self._evict_component_and_detach_lru(
node, node,
@@ -1494,7 +1496,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self._cascade_evict( self._cascade_evict(
node, trigger, tracker, device_frees=device_frees, host_frees=host_frees node, trigger, tracker, device_frees=device_frees, host_frees=host_frees
) )
self._record_remove_event(node, medium=StorageMedium.GPU) self.kv_events.record_remove(node, medium=StorageMedium.GPU)
# after device eviction, insert aux components into host LRU. # after device eviction, insert aux components into host LRU.
self._for_each_component_lru( self._for_each_component_lru(
@@ -1980,7 +1982,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
cache_actions=cache_actions, cache_actions=cache_actions,
) )
for nid in kv_xfer.nodes_to_load or (): for nid in kv_xfer.nodes_to_load or ():
self._record_store_event(self.node_by_id(nid), medium=StorageMedium.GPU) self.kv_events.record_store(self.node_by_id(nid), medium=StorageMedium.GPU)
for ct, xfers in comp_xfers.items(): for ct, xfers in comp_xfers.items():
self.components_by_type[ct].commit_hicache_transfer( self.components_by_type[ct].commit_hicache_transfer(
node, node,
@@ -2023,7 +2025,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
node.write_through_pending_id = None node.write_through_pending_id = None
# The backed-up copy becomes a tracked duplicate only now. # The backed-up copy becomes a tracked duplicate only now.
self._update_duplicate_tracking(node) self._update_duplicate_tracking(node)
self._record_store_event(node, medium=StorageMedium.CPU) self.kv_events.record_store(node, medium=StorageMedium.CPU)
def set_component_device_value( def set_component_device_value(
self, node_id: NodeId, component_type: ComponentType, value: torch.Tensor self, node_id: NodeId, component_type: ComponentType, value: torch.Tensor
python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py
+8 -5
View File
@@ -12,8 +12,6 @@ from typing import TYPE_CHECKING, Optional, Sequence
import msgspec import msgspec
from sglang.srt.mem_cache.events import KVCacheEventMixin
# Tree node id -- the node handle used outside the TreeCore. The concrete tree # Tree node id -- the node handle used outside the TreeCore. The concrete tree
# node is a TreeCore-internal type. # node is a TreeCore-internal type.
NodeId = int NodeId = int
@@ -95,6 +93,7 @@ if TYPE_CHECKING:
MatchPrefixParams, MatchPrefixParams,
MatchResult, MatchResult,
) )
from sglang.srt.mem_cache.events import KVCacheEventRecorder
from sglang.srt.mem_cache.hicache_storage import PoolTransfer, PoolTransferResult from sglang.srt.mem_cache.hicache_storage import PoolTransfer, PoolTransferResult
from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.unified_cache.cache_action import ( from sglang.srt.mem_cache.unified_cache.cache_action import (
@@ -112,11 +111,10 @@ if TYPE_CHECKING:
) )
class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC): class UnifiedTreeCoreInterface(ABC):
"""Methods the Controller invokes on the Tree Core. The Controller treats the """Methods the Controller invokes on the Tree Core. The Controller treats the
Tree Core as opaque behind this surface, which grows as tree operations Tree Core as opaque behind this surface, which grows as tree operations
migrate onto the TreeCore. Inherits KVCacheEventMixin for the KV-event API migrate onto the TreeCore."""
(take_events, _record_* recorders)."""
# ==== Tree-owned state the Controller reads (or, via its facade setters, writes) ==== # ==== Tree-owned state the Controller reads (or, via its facade setters, writes) ====
page_size: int page_size: int
@@ -127,9 +125,14 @@ class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC):
write_through_threshold: int write_through_threshold: int
is_write_back: bool is_write_back: bool
has_swa_host_pool: bool has_swa_host_pool: bool
kv_events: KVCacheEventRecorder
# ==== Tree API ==== # ==== Tree API ====
def take_events(self) -> list:
"""Hand the queued KV placement events to the Controller."""
return self.kv_events.take()
@abstractmethod @abstractmethod
def reset(self) -> None: def reset(self) -> None:
"""Drop the entire tree and reinitialize empty state.""" """Drop the entire tree and reinitialize empty state."""
python/sglang/srt/mem_cache/unified_radix_cache.py
+1 -1
View File
@@ -322,7 +322,7 @@ class UnifiedRadixCache(BasePrefixCache):
self.cache_controller.mem_pool_host.clear() self.cache_controller.mem_pool_host.clear()
self.enable_storage = self.cache_controller.enable_storage self.enable_storage = self.cache_controller.enable_storage
self.tree_core._record_all_cleared_event() self.tree_core.kv_events.record_all_cleared()
def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None: def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None:
"""Initialize HiCache infrastructure.""" """Initialize HiCache infrastructure."""
test/registered/unit/mem_cache/test_radix_cache_unit.py
+28 -44
View File
@@ -45,7 +45,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
InsertParams, InsertParams,
MatchPrefixParams, MatchPrefixParams,
) )
from sglang.srt.mem_cache.events import KVCacheEventMixin from sglang.srt.mem_cache.events import KVCacheEventRecorder
from sglang.srt.mem_cache.mamba_radix_cache import TreeNode as MambaTreeNode from sglang.srt.mem_cache.mamba_radix_cache import TreeNode as MambaTreeNode
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
from sglang.srt.utils import get_device from sglang.srt.utils import get_device
@@ -54,12 +54,6 @@ from sglang.srt.utils import get_device
DEFAULT_PAGE_SIZE = 4 DEFAULT_PAGE_SIZE = 4
class _KVCacheEventQueue(KVCacheEventMixin):
def __init__(self):
self.enable_kv_cache_events = True
self.kv_event_queue = []
class TestKVCacheEventQueue(unittest.TestCase): class TestKVCacheEventQueue(unittest.TestCase):
@staticmethod @staticmethod
def _store( def _store(
@@ -87,26 +81,22 @@ class TestKVCacheEventQueue(unittest.TestCase):
) )
def test_enqueue_coalesces_compatible_stores(self): def test_enqueue_coalesces_compatible_stores(self):
queue = _KVCacheEventQueue() queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue._enqueue_kv_event(self._store(1, None)) queue.enqueue(self._store(1, None))
queue._enqueue_kv_event(self._store(2, 1)) queue.enqueue(self._store(2, 1))
events = queue.take_events() events = queue.take()
self.assertEqual(len(events), 1) self.assertEqual(len(events), 1)
self.assertEqual(events[0].block_hashes, [1, 2]) self.assertEqual(events[0].block_hashes, [1, 2])
self.assertEqual(events[0].parent_block_hash, None) self.assertEqual(events[0].parent_block_hash, None)
self.assertEqual(events[0].token_ids, [1, 2, 2, 3]) self.assertEqual(events[0].token_ids, [1, 2, 2, 3])
def test_enqueue_coalesces_compatible_removes(self): def test_enqueue_coalesces_compatible_removes(self):
queue = _KVCacheEventQueue() queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue._enqueue_kv_event( queue.enqueue(BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU))
BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU) queue.enqueue(BlockRemoved(block_hashes=[2, 3], medium=StorageMedium.GPU))
)
queue._enqueue_kv_event(
BlockRemoved(block_hashes=[2, 3], medium=StorageMedium.GPU)
)
events = queue.take_events() events = queue.take()
self.assertEqual(len(events), 1) self.assertEqual(len(events), 1)
self.assertIsInstance(events[0], BlockRemoved) self.assertIsInstance(events[0], BlockRemoved)
self.assertEqual(events[0].block_hashes, [1, 2, 3]) self.assertEqual(events[0].block_hashes, [1, 2, 3])
@@ -119,33 +109,27 @@ class TestKVCacheEventQueue(unittest.TestCase):
self._store(5, None), self._store(5, None),
] ]
for incoming in incompatible_stores: for incoming in incompatible_stores:
queue = _KVCacheEventQueue() queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue._enqueue_kv_event(self._store(1, None)) queue.enqueue(self._store(1, None))
queue._enqueue_kv_event(incoming) queue.enqueue(incoming)
self.assertEqual(len(queue.take_events()), 2) self.assertEqual(len(queue.take()), 2)
queue = _KVCacheEventQueue() queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue._enqueue_kv_event(self._store(1, None)) queue.enqueue(self._store(1, None))
queue._enqueue_kv_event( queue.enqueue(BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU))
BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU) queue.enqueue(AllBlocksCleared())
) queue.enqueue(self._store(2, None))
queue._enqueue_kv_event(AllBlocksCleared()) self.assertEqual(len(queue.take()), 4)
queue._enqueue_kv_event(self._store(2, None))
self.assertEqual(len(queue.take_events()), 4)
queue = _KVCacheEventQueue() queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue._enqueue_kv_event( queue.enqueue(BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU))
BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU) queue.enqueue(BlockRemoved(block_hashes=[2], medium=StorageMedium.CPU))
) self.assertEqual(len(queue.take()), 2)
queue._enqueue_kv_event(
BlockRemoved(block_hashes=[2], medium=StorageMedium.CPU)
)
self.assertEqual(len(queue.take_events()), 2)
queue = _KVCacheEventQueue() queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue._enqueue_kv_event(self._store(1, None, cache_salt="tenant-a")) queue.enqueue(self._store(1, None, cache_salt="tenant-a"))
queue._enqueue_kv_event(self._store(2, 1, cache_salt="tenant-b")) queue.enqueue(self._store(2, 1, cache_salt="tenant-b"))
self.assertEqual(len(queue.take_events()), 2) self.assertEqual(len(queue.take()), 2)
class TestRadixKey(unittest.TestCase): class TestRadixKey(unittest.TestCase):
@@ -410,7 +394,7 @@ class TestRadixCache(unittest.TestCase):
self.assertEqual(cache.page_size, page_size) self.assertEqual(cache.page_size, page_size)
self.assertEqual(cache.disable, disable) self.assertEqual(cache.disable, disable)
self.assertEqual(cache.enable_kv_cache_events, enable_events) self.assertEqual(cache.kv_events.enabled, enable_events)
self.assertEqual(cache.device, torch.device("cpu")) self.assertEqual(cache.device, torch.device("cpu"))
self.assertIsNotNone(cache.root_node) self.assertIsNotNone(cache.root_node)
self.assertEqual(len(cache.root_node.key), 0) self.assertEqual(len(cache.root_node.key), 0)