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

feat(agent sessions): attribute stored KV cache blocks to sessions (#37482)

Browse Source 
Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
This commit is contained in:
ishandhanani
2026-09-14 21:34:49 -07:00
committed by GitHub
parent c9fbe5f655
commit 3f871a246c
22 changed files with 905 additions and 451 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/kv_events.py
+18 -22
View File
@@ -241,11 +241,20 @@ class EventBatch(
class KVCacheEvent(
msgspec.Struct,
array_like=True, # type: ignore[call-arg]
omit_defaults=True, # type: ignore[call-arg]
gc=False, # type: ignore[call-arg]
tag=True,
):
"""Base class for all KV cache-related events"""
"""Base class for all KV cache-related events.
Events are tagged msgpack maps: ``type`` carries the class name and every
other key is a field name. Optional fields left at ``None`` are omitted, so
adding an optional field never changes the shape an older consumer sees.
This is the same encoding vLLM uses for its ``KVCacheEvent``, so a consumer
such as Dynamo decodes both engines with one code path.
``EventBatch`` stays a positional array ``[ts, events, attn_dp_rank]``.
"""
class StorageMedium(str, enum.Enum):
@@ -257,12 +266,6 @@ class StorageMedium(str, enum.Enum):
EXTERNAL = "EXTERNAL" # L4: shared / remote pool (e.g. Mooncake)
class BlockStoredMetadata(msgspec.Struct, omit_defaults=True, gc=False):
"""Typed request metadata attached to a stored KV block."""
cache_salt: str
class OffloadedState(msgspec.Struct):
"""Decode-side offload progress for one request, keyed by Req in the manager."""
@@ -279,16 +282,13 @@ class BlockStored(KVCacheEvent):
block_size: int
lora_id: Optional[int]
medium: Optional[str] = None
class BlockStoredWithMetadata(BlockStored, tag="BlockStored", kw_only=True):
"""BlockStored wire extension used only when typed metadata is present.
A separate struct keeps unsalted events at their legacy array length; an
optional field on BlockStored would still serialize a trailing null.
"""
metadata: BlockStoredMetadata
# Salt of the request that stored these blocks. Block hashes are already
# namespaced by it; consumers index the emitted hashes rather than
# recompute them.
cache_salt: Optional[str] = None
# Session that triggered this store. Attribution only: the blocks may be
# shared with other sessions, and the hash does not depend on it.
session_id: Optional[str] = None
class BlockRemoved(KVCacheEvent):
@@ -301,10 +301,6 @@ class AllBlocksCleared(KVCacheEvent):
class KVEventBatch(EventBatch):
# BlockStoredWithMetadata deliberately stays out of this tagged union.
# Existing typed consumers decode its shared "BlockStored" tag as the base
# type and ignore the trailing metadata; adding both types would give
# msgspec duplicate tags and make the union invalid.
events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]]
python/sglang/srt/mem_cache/base_prefix_cache.py
+1
View File
@@ -90,6 +90,7 @@ class InsertParams:
# General
chunked: bool = False
priority: int = 0
session_id: Optional[str] = None
track_adopted_ranges: bool = False
# Logical-page KV sharding: rotation base of the chain the inserted
python/sglang/srt/mem_cache/events.py
+16 -27
View File
@@ -24,8 +24,6 @@ from sglang.srt.disaggregation.kv_events import (
AllBlocksCleared,
BlockRemoved,
BlockStored,
BlockStoredMetadata,
BlockStoredWithMetadata,
StorageMedium,
)
from sglang.srt.mem_cache.utils import (
@@ -63,19 +61,12 @@ class KVCacheEventRecorder:
return
elif isinstance(tail, BlockStored) and isinstance(event, BlockStored):
tail_metadata = (
tail.metadata if isinstance(tail, BlockStoredWithMetadata) else None
)
event_metadata = (
event.metadata
if isinstance(event, BlockStoredWithMetadata)
else None
)
if (
tail.medium == event.medium
and tail.lora_id == event.lora_id
and tail.block_size == event.block_size
and tail_metadata == event_metadata
and tail.cache_salt == event.cache_salt
and tail.session_id == event.session_id
and tail.block_hashes
and event.parent_block_hash == tail.block_hashes[-1]
):
@@ -112,7 +103,9 @@ class KVCacheEventRecorder:
return None
return hash_str_to_int64(parent_hash_values[-1])
def record_store(self, node: Any, medium=None) -> None:
def record_store(
self, node: Any, medium=None, *, session_id: Optional[str] = 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).
@@ -140,22 +133,18 @@ class KVCacheEventRecorder:
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(
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,
cache_salt=node.key.cache_salt,
session_id=session_id,
)
self.enqueue(event)
)
parent_block_hash = block_hash
page_index += 1
python/sglang/srt/mem_cache/rust_tree_core/adapter.py
+4 -9
View File
@@ -12,8 +12,6 @@ from sglang.srt.disaggregation.kv_events import (
AllBlocksCleared,
BlockRemoved,
BlockStored,
BlockStoredMetadata,
BlockStoredWithMetadata,
StorageMedium,
)
from sglang.srt.mem_cache.base_prefix_cache import (
@@ -84,19 +82,15 @@ def _kv_event_from_tagged(event: tuple):
"""Build the Python KV cache event for one of the binding's tagged tuples."""
tag = event[0]
if tag == "block_stored":
event_args = dict(
return BlockStored(
block_hashes=event[1],
parent_block_hash=event[2],
token_ids=event[3],
block_size=event[4],
lora_id=None,
medium=StorageMedium(event[5]),
)
if event[6] is None:
return BlockStored(**event_args)
return BlockStoredWithMetadata(
**event_args,
metadata=BlockStoredMetadata(cache_salt=event[6]),
cache_salt=event[6],
session_id=event[7],
)
if tag == "block_removed":
return BlockRemoved(block_hashes=event[1], medium=StorageMedium(event[2]))
@@ -588,6 +582,7 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
value=value,
extra_key=key.extra_key,
cache_salt=key.cache_salt,
session_id=params.session_id,
mamba_value=params.mamba_value,
prev_prefix_len=params.prev_prefix_len,
swa_evicted_seqlen=params.swa_evicted_seqlen,
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+17 -4
View File
@@ -1161,7 +1161,11 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
node.priority = max(node.priority, state.priority)
if node.evicted:
self._unevict_node_on_insert(node, state.value[:prefix_len])
self._unevict_node_on_insert(
node,
state.value[:prefix_len],
session_id=state.params.session_id,
)
state.result.record_adopted_range(
BASE_COMPONENT_TYPE,
state.total_prefix_length,
@@ -1235,6 +1239,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
state.key,
state.value,
priority=state.priority,
session_id=state.params.session_id,
rotation_base=state.params.rotation_base,
)
state.is_new_leaf = True
@@ -1360,6 +1365,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
key: RadixKey,
value: torch.Tensor,
priority: int = 0,
session_id: Optional[str] = None,
rotation_base: Optional[int] = None,
) -> UnifiedTreeNode:
new_node = self._new_node(priority=priority)
@@ -1378,11 +1384,14 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self._update_evictable_leaf_sets(new_node)
self._update_evictable_leaf_sets(parent)
self.kv_events.record_store(new_node)
self.kv_events.record_store(new_node, session_id=session_id)
return new_node
def _unevict_node_on_insert(
self, node: UnifiedTreeNode, fresh_value: torch.Tensor
self,
node: UnifiedTreeNode,
fresh_value: torch.Tensor,
session_id: Optional[str] = None,
) -> None:
"""Restore an evicted node's Full device value from fresh KV indices
during insert."""
@@ -1400,7 +1409,11 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self._update_duplicate_tracking(node)
if node.parent is not None:
self._update_evictable_leaf_sets(node.parent)
self.kv_events.record_store(node, medium=StorageMedium.GPU)
self.kv_events.record_store(
node,
medium=StorageMedium.GPU,
session_id=session_id,
)
def _update_evictable_leaf_sets(self, node: UnifiedTreeNode) -> None:
"""Update both device and host leaf sets for a node."""
python/sglang/srt/mem_cache/unified_radix_cache.py
+2
View File
@@ -977,6 +977,7 @@ class UnifiedRadixCache(BasePrefixCache):
insert_params = InsertParams(
prev_prefix_len=req.kv.cache_protected_len,
priority=getattr(req, "priority", 0) or 0,
session_id=req.session_id,
rotation_base=req.kv_rotation_base,
)
@@ -1113,6 +1114,7 @@ class UnifiedRadixCache(BasePrefixCache):
prev_prefix_len=req.kv.cache_protected_len,
chunked=chunked,
priority=getattr(req, "priority", 0) or 0,
session_id=req.session_id,
rotation_base=req.kv_rotation_base,
)
effective_cache_len = len(token_ids)