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
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.observability.metrics_collector import (
STAT_LOGGER_ROLE_RADIX_CACHE,
@@ -238,6 +239,8 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
None # metrics collector for the cache
)
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):
from sglang.srt.runtime_context import get_server_args
@@ -377,7 +380,7 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
raise NotImplementedError()
def take_events(self):
return []
return [] if self.kv_events is None else self.kv_events.take()
def supports_swa(self) -> bool:
return False
python/sglang/srt/mem_cache/events.py
+113 -96
View File
@@ -11,13 +11,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""KV cache placement event emission mixin.
"""KV cache placement event recording.
The mixin produces the ``BlockStored`` / ``BlockRemoved`` / ``AllBlocksCleared``
events consumed by KV-aware routers (e.g. dynamo).
Produces the ``BlockStored`` / ``BlockRemoved`` / ``AllBlocksCleared`` events
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 (
AllBlocksCleared,
@@ -34,16 +35,27 @@ from sglang.srt.mem_cache.utils import (
)
class KVCacheEventMixin:
def _enqueue_kv_event(self, event):
class KVCacheEventRecorder:
"""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.
KV event batches already support multiple block hashes. Combining them
here avoids emitting one event per page while preserving ordering and
the parent-linked store chains consumers use to rebuild the cache tree.
"""
if self.kv_event_queue:
tail = self.kv_event_queue[-1]
if self._queue:
tail = self._queue[-1]
if isinstance(tail, BlockRemoved) and isinstance(event, BlockRemoved):
if tail.medium == event.medium:
@@ -71,117 +83,122 @@ class KVCacheEventMixin:
tail.token_ids.extend(event.token_ids)
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.
# ``medium`` defaults to StorageMedium.GPU but callers may override
# for lower-tier insertions (e.g. StorageMedium.CPU for host/L2 cache).
if self.enable_kv_cache_events:
if medium is None:
medium = StorageMedium.GPU
if not self.enabled:
return
if medium is None:
medium = StorageMedium.GPU
# 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)
event_hash_values = (
compute_node_event_hash_values(node, self.page_size)
if node.key.cache_salt is not None
else node.hash_value
)
event_hash_values = self._node_event_hash_values(node)
parent_block_hash = self._parent_block_hash(node)
# 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
logical_len = len(node.key)
is_bigram = node.key.is_bigram
raw = node.key.token_ids
for start in range(0, logical_len, self.page_size):
end = min(start + self.page_size, logical_len)
if end <= start:
continue
# Preserve historical event payload: bigram pages expose tuples.
if is_bigram:
page_tokens = [(raw[j], raw[j + 1]) for j in range(start, end)]
else:
page_tokens = list(raw[start:end])
page_index = 0
logical_len = len(node.key)
is_bigram = node.key.is_bigram
raw = node.key.token_ids
for start in range(0, logical_len, self.page_size):
end = min(start + self.page_size, logical_len)
if end <= start:
continue
# Preserve historical event payload: bigram pages expose tuples.
if is_bigram:
page_tokens = [(raw[j], raw[j + 1]) for j in range(start, end)]
else:
page_tokens = list(raw[start:end])
block_hash = hash_str_to_int64(event_hash_values[page_index])
block_hash = hash_str_to_int64(event_hash_values[page_index])
event_args = {
"block_hashes": [block_hash],
"parent_block_hash": parent_block_hash,
"token_ids": page_tokens,
"block_size": len(page_tokens),
"lora_id": None,
"medium": medium,
}
if node.key.cache_salt is None:
event = BlockStored(**event_args)
else:
event = BlockStoredWithMetadata(
**event_args,
metadata=BlockStoredMetadata(cache_salt=node.key.cache_salt),
)
self.enqueue(event)
event_args = {
"block_hashes": [block_hash],
"parent_block_hash": parent_block_hash,
"token_ids": page_tokens,
"block_size": len(page_tokens),
"lora_id": None,
"medium": medium,
}
if node.key.cache_salt is None:
event = BlockStored(**event_args)
else:
event = BlockStoredWithMetadata(
**event_args,
metadata=BlockStoredMetadata(cache_salt=node.key.cache_salt),
)
self._enqueue_kv_event(event)
parent_block_hash = block_hash
page_index += 1
parent_block_hash = block_hash
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.
# ``medium`` defaults to StorageMedium.GPU but callers may override for
# lower-tier removals (e.g. StorageMedium.CPU when evicting from host).
if self.enable_kv_cache_events:
if medium is None:
medium = StorageMedium.GPU
if not self.enabled:
return
if medium is None:
medium = StorageMedium.GPU
# 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)
event_hash_values = (
compute_node_event_hash_values(node, self.page_size)
if node.key.cache_salt is not None
else node.hash_value
)
# Hash values must match what was stored.
event_hash_values = self._node_event_hash_values(node)
block_hashes = []
logical_len = len(node.key)
page_index = 0
for start in range(0, logical_len, self.page_size):
end = min(start + self.page_size, logical_len)
if end <= start:
continue
block_hashes = []
logical_len = len(node.key)
page_index = 0
for start in range(0, logical_len, self.page_size):
end = min(start + self.page_size, logical_len)
if end <= start:
continue
block_hashes.append(hash_str_to_int64(event_hash_values[page_index]))
page_index += 1
block_hashes.append(hash_str_to_int64(event_hash_values[page_index]))
page_index += 1
if block_hashes:
self._enqueue_kv_event(
BlockRemoved(block_hashes=block_hashes, medium=medium)
)
if block_hashes:
self.enqueue(BlockRemoved(block_hashes=block_hashes, medium=medium))
def _record_all_cleared_event(self):
if self.enable_kv_cache_events:
self._enqueue_kv_event(AllBlocksCleared())
def record_all_cleared(self) -> None:
if not self.enabled:
return
self.enqueue(AllBlocksCleared())
def take_events(self):
def take(self) -> list:
"""Atomically takes all events and clears the queue.
Returns:
A list of KV cache events.
"""
if not self.enable_kv_cache_events:
if not self.enabled:
return []
events = self.kv_event_queue
self.kv_event_queue = []
events = self._queue
self._queue = []
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:
node.write_through_pending_id = None
# 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:
self.write_backup_storage(lock_node, backup_len)
if release_lock:
@@ -1264,7 +1264,7 @@ class HiRadixCache(RadixCache):
def _detach_backuped(self, node: TreeNode) -> int:
# 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)
assert num_evicted > 0
self.evictable_size_ -= num_evicted
@@ -1285,7 +1285,7 @@ class HiRadixCache(RadixCache):
# evict a node not initiated write to host -- emit BlockRemoved
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)
num_evicted = len(node.value)
self._delete_leaf(node)
@@ -1311,11 +1311,11 @@ class HiRadixCache(RadixCache):
freed_device = 0
for n in nodes:
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)
n.host_value = 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)
freed_device += 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) --
# 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)
key = x.key.child_key(self.page_size)
@@ -1438,7 +1438,7 @@ class HiRadixCache(RadixCache):
offset += len(node.host_value)
# Block promoted from host to GPU -- emit store(GPU) so downstream
# 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.inc_lock_ref(last_hit_node)
@@ -1849,7 +1849,7 @@ class HiRadixCache(RadixCache):
self._update_host_leaf_status(node)
# Publish the newly materialized host suffix immediately so downstream
# 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
@@ -1987,11 +1987,11 @@ class HiRadixCache(RadixCache):
self._update_leaf_status(new_node)
# 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)
# 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":
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,
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.multi_ended_allocator import (
UnifiedMambaTokenToKVPoolAllocator,
@@ -441,7 +441,7 @@ class LRUList:
raise Exception(msg)
class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
class MambaRadixCache(BasePrefixCache):
def __init__(self, params: CacheInitParams):
assert (
isinstance(params.token_to_kv_pool_allocator, TokenToKVPoolAllocator)
@@ -458,10 +458,11 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
self.page_size = params.page_size
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_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:
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
self.full_lru_list = LRUList(mamba=False)
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:
"""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=}"
# 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.
self.token_to_kv_pool_allocator.free_segment(x.value, start_pos=0)
full_num_evicted = len(x.value)
@@ -1301,7 +1302,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
node.children[child_key] = new_node
self.full_evictable_size_ += len(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
node.mamba_value = mamba_value
self.full_lru_list.reset_node_mru(node)
@@ -1330,7 +1331,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
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=}"
# 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)
full_num_evicted += len(node.parent.value)
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,
MatchResult,
)
from sglang.srt.mem_cache.events import KVCacheEventMixin
from sglang.srt.mem_cache.events import KVCacheEventRecorder
from sglang.srt.mem_cache.utils import (
get_eviction_strategy,
get_hash_str,
@@ -300,18 +300,19 @@ class TreeNode:
return self.last_access_time < other.last_access_time
class RadixCache(KVCacheEventMixin, BasePrefixCache):
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 = []
self.kv_events = KVCacheEventRecorder(
enabled=params.enable_kv_cache_events, page_size=self.page_size
)
if params.enable_metrics:
self.init_metrics_collector()
@@ -371,7 +372,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
last_host_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:
"""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)
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)
return EvictResult(num_tokens_evicted=num_evicted)
@@ -785,7 +786,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
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)
self.kv_events.record_store(new_node)
node = new_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(new_node)
self._record_store_event(new_node.parent)
self._record_store_event(new_node)
self.kv_events.record_store(new_node.parent)
self.kv_events.record_store(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(new_node)
self._record_store_event(new_node.parent)
self._record_store_event(new_node)
self.kv_events.record_store(new_node.parent)
self.kv_events.record_store(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,
)
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.utils import split_node_hash_value
@@ -342,7 +342,7 @@ class LRUList:
raise Exception(msg)
class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
class SWARadixCache(BasePrefixCache):
def __init__(self, params: CacheInitParams):
assert isinstance(params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator)
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.disable = params.disable
self.is_eagle = params.is_eagle
self.enable_kv_cache_events = params.enable_kv_cache_events
self.kv_event_queue = []
self.kv_events = KVCacheEventRecorder(
enabled=params.enable_kv_cache_events, page_size=self.page_size
)
if self.token_to_kv_pool_allocator:
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
self.full_lru_list = LRUList(is_swa_list=False)
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:
"""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=}"
# 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)
full_num_evicted += len(x.value)
# 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
), f"leaf node with full lock must also have swa lock, {x.id=}"
# 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)
full_num_evicted += len(x.value)
swa_num_evicted += len(x.value)
@@ -1326,7 +1327,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
if not swa_tombstone:
self.swa_lru_list.insert_mru(new_node)
self.swa_evictable_size_ += len(value)
self._record_store_event(new_node)
self.kv_events.record_store(new_node)
return new_node
def _iteratively_delete_tombstone_leaf(
@@ -1344,7 +1345,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
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=}"
# 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)
full_num_evicted += len(node.parent.value)
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,
MatchResult,
)
from sglang.srt.mem_cache.events import KVCacheEventRecorder
from sglang.srt.mem_cache.hicache_storage import (
PoolName,
PoolTransfer,
@@ -417,8 +418,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self.components_by_type.values()
)
self.enable_kv_cache_events = params.enable_kv_cache_events
self.kv_event_queue = []
self.kv_events = KVCacheEventRecorder(
enabled=params.enable_kv_cache_events, page_size=self.page_size
)
self.reset()
@@ -1119,7 +1121,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self._update_evictable_leaf_sets(new_node)
self._update_evictable_leaf_sets(parent)
self._record_store_event(new_node)
self.kv_events.record_store(new_node)
return new_node
def _unevict_node_on_insert(
@@ -1138,7 +1140,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self._update_duplicate_tracking(node)
if node.parent is not None:
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:
"""Update both device and host leaf sets for a node."""
@@ -1303,7 +1305,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
) -> None:
"""Free every component layer on the node and detach it from the LRU
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:
self._evict_component_and_detach_lru(
node,
@@ -1422,7 +1424,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
"""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)."""
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(
node,
self.components_by_type[BASE_COMPONENT_TYPE],
@@ -1447,7 +1449,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
All freed tokens are accumulated into *tracker*."""
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:
_, hf = self._evict_component_and_detach_lru(
node,
@@ -1494,7 +1496,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self._cascade_evict(
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.
self._for_each_component_lru(
@@ -1980,7 +1982,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
cache_actions=cache_actions,
)
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():
self.components_by_type[ct].commit_hicache_transfer(
node,
@@ -2023,7 +2025,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
node.write_through_pending_id = None
# The backed-up copy becomes a tracked duplicate only now.
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(
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
from sglang.srt.mem_cache.events import KVCacheEventMixin
# Tree node id -- the node handle used outside the TreeCore. The concrete tree
# node is a TreeCore-internal type.
NodeId = int
@@ -95,6 +93,7 @@ if TYPE_CHECKING:
MatchPrefixParams,
MatchResult,
)
from sglang.srt.mem_cache.events import KVCacheEventRecorder
from sglang.srt.mem_cache.hicache_storage import PoolTransfer, PoolTransferResult
from sglang.srt.mem_cache.radix_cache import RadixKey
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
Tree Core as opaque behind this surface, which grows as tree operations
migrate onto the TreeCore. Inherits KVCacheEventMixin for the KV-event API
(take_events, _record_* recorders)."""
migrate onto the TreeCore."""
# ==== Tree-owned state the Controller reads (or, via its facade setters, writes) ====
page_size: int
@@ -127,9 +125,14 @@ class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC):
write_through_threshold: int
is_write_back: bool
has_swa_host_pool: bool
kv_events: KVCacheEventRecorder
# ==== Tree API ====
def take_events(self) -> list:
"""Hand the queued KV placement events to the Controller."""
return self.kv_events.take()
@abstractmethod
def reset(self) -> None:
"""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.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:
"""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,
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.radix_cache import RadixCache, RadixKey, TreeNode
from sglang.srt.utils import get_device
@@ -54,12 +54,6 @@ from sglang.srt.utils import get_device
DEFAULT_PAGE_SIZE = 4
class _KVCacheEventQueue(KVCacheEventMixin):
def __init__(self):
self.enable_kv_cache_events = True
self.kv_event_queue = []
class TestKVCacheEventQueue(unittest.TestCase):
@staticmethod
def _store(
@@ -87,26 +81,22 @@ class TestKVCacheEventQueue(unittest.TestCase):
)
def test_enqueue_coalesces_compatible_stores(self):
queue = _KVCacheEventQueue()
queue._enqueue_kv_event(self._store(1, None))
queue._enqueue_kv_event(self._store(2, 1))
queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue.enqueue(self._store(1, None))
queue.enqueue(self._store(2, 1))
events = queue.take_events()
events = queue.take()
self.assertEqual(len(events), 1)
self.assertEqual(events[0].block_hashes, [1, 2])
self.assertEqual(events[0].parent_block_hash, None)
self.assertEqual(events[0].token_ids, [1, 2, 2, 3])
def test_enqueue_coalesces_compatible_removes(self):
queue = _KVCacheEventQueue()
queue._enqueue_kv_event(
BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU)
)
queue._enqueue_kv_event(
BlockRemoved(block_hashes=[2, 3], medium=StorageMedium.GPU)
)
queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue.enqueue(BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU))
queue.enqueue(BlockRemoved(block_hashes=[2, 3], medium=StorageMedium.GPU))
events = queue.take_events()
events = queue.take()
self.assertEqual(len(events), 1)
self.assertIsInstance(events[0], BlockRemoved)
self.assertEqual(events[0].block_hashes, [1, 2, 3])
@@ -119,33 +109,27 @@ class TestKVCacheEventQueue(unittest.TestCase):
self._store(5, None),
]
for incoming in incompatible_stores:
queue = _KVCacheEventQueue()
queue._enqueue_kv_event(self._store(1, None))
queue._enqueue_kv_event(incoming)
self.assertEqual(len(queue.take_events()), 2)
queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue.enqueue(self._store(1, None))
queue.enqueue(incoming)
self.assertEqual(len(queue.take()), 2)
queue = _KVCacheEventQueue()
queue._enqueue_kv_event(self._store(1, None))
queue._enqueue_kv_event(
BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU)
)
queue._enqueue_kv_event(AllBlocksCleared())
queue._enqueue_kv_event(self._store(2, None))
self.assertEqual(len(queue.take_events()), 4)
queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue.enqueue(self._store(1, None))
queue.enqueue(BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU))
queue.enqueue(AllBlocksCleared())
queue.enqueue(self._store(2, None))
self.assertEqual(len(queue.take()), 4)
queue = _KVCacheEventQueue()
queue._enqueue_kv_event(
BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU)
)
queue._enqueue_kv_event(
BlockRemoved(block_hashes=[2], medium=StorageMedium.CPU)
)
self.assertEqual(len(queue.take_events()), 2)
queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue.enqueue(BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU))
queue.enqueue(BlockRemoved(block_hashes=[2], medium=StorageMedium.CPU))
self.assertEqual(len(queue.take()), 2)
queue = _KVCacheEventQueue()
queue._enqueue_kv_event(self._store(1, None, cache_salt="tenant-a"))
queue._enqueue_kv_event(self._store(2, 1, cache_salt="tenant-b"))
self.assertEqual(len(queue.take_events()), 2)
queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE)
queue.enqueue(self._store(1, None, cache_salt="tenant-a"))
queue.enqueue(self._store(2, 1, cache_salt="tenant-b"))
self.assertEqual(len(queue.take()), 2)
class TestRadixKey(unittest.TestCase):
@@ -410,7 +394,7 @@ class TestRadixCache(unittest.TestCase):
self.assertEqual(cache.page_size, page_size)
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.assertIsNotNone(cache.root_node)
self.assertEqual(len(cache.root_node.key), 0)