diff --git a/python/sglang/srt/disaggregation/kv_events.py b/python/sglang/srt/disaggregation/kv_events.py index 22c7aeeb3..b8048ee6b 100644 --- a/python/sglang/srt/disaggregation/kv_events.py +++ b/python/sglang/srt/disaggregation/kv_events.py @@ -56,16 +56,23 @@ class KVCacheEvent( """Base class for all KV cache-related events""" +# Medium values for hicache storage tiers +MEDIUM_GPU = "GPU" +MEDIUM_CPU = "CPU_PINNED" + + class BlockStored(KVCacheEvent): block_hashes: list[int] parent_block_hash: Optional[int] token_ids: list[int] block_size: int lora_id: Optional[int] + medium: Optional[str] = None class BlockRemoved(KVCacheEvent): block_hashes: list[int] + medium: Optional[str] = None class AllBlocksCleared(KVCacheEvent): diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index fcec2c9ac..9dba0b349 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -35,6 +35,7 @@ import torch logger = logging.getLogger(__name__) from sglang.srt.disaggregation.kv_events import ( + MEDIUM_GPU, AllBlocksCleared, BlockRemoved, BlockStored, @@ -822,6 +823,7 @@ class RadixCache(BasePrefixCache): token_ids=page_tokens, block_size=len(page_tokens), lora_id=None, + medium=MEDIUM_GPU, ) ) @@ -843,7 +845,9 @@ class RadixCache(BasePrefixCache): block_hash = hash_str_to_int64(node.hash_value[page_index]) - self.kv_event_queue.append(BlockRemoved(block_hashes=[block_hash])) + self.kv_event_queue.append( + BlockRemoved(block_hashes=[block_hash], medium=MEDIUM_GPU) + ) page_index += 1