feat(kv-events): Add medium field to KV event types for storage tier tracking (#18205)
This commit is contained in:
@@ -56,16 +56,23 @@ class KVCacheEvent(
|
|||||||
"""Base class for all KV cache-related events"""
|
"""Base class for all KV cache-related events"""
|
||||||
|
|
||||||
|
|
||||||
|
# Medium values for hicache storage tiers
|
||||||
|
MEDIUM_GPU = "GPU"
|
||||||
|
MEDIUM_CPU = "CPU_PINNED"
|
||||||
|
|
||||||
|
|
||||||
class BlockStored(KVCacheEvent):
|
class BlockStored(KVCacheEvent):
|
||||||
block_hashes: list[int]
|
block_hashes: list[int]
|
||||||
parent_block_hash: Optional[int]
|
parent_block_hash: Optional[int]
|
||||||
token_ids: list[int]
|
token_ids: list[int]
|
||||||
block_size: int
|
block_size: int
|
||||||
lora_id: Optional[int]
|
lora_id: Optional[int]
|
||||||
|
medium: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
class BlockRemoved(KVCacheEvent):
|
class BlockRemoved(KVCacheEvent):
|
||||||
block_hashes: list[int]
|
block_hashes: list[int]
|
||||||
|
medium: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
class AllBlocksCleared(KVCacheEvent):
|
class AllBlocksCleared(KVCacheEvent):
|
||||||
|
|||||||
@@ -35,6 +35,7 @@ import torch
|
|||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
from sglang.srt.disaggregation.kv_events import (
|
from sglang.srt.disaggregation.kv_events import (
|
||||||
|
MEDIUM_GPU,
|
||||||
AllBlocksCleared,
|
AllBlocksCleared,
|
||||||
BlockRemoved,
|
BlockRemoved,
|
||||||
BlockStored,
|
BlockStored,
|
||||||
@@ -822,6 +823,7 @@ class RadixCache(BasePrefixCache):
|
|||||||
token_ids=page_tokens,
|
token_ids=page_tokens,
|
||||||
block_size=len(page_tokens),
|
block_size=len(page_tokens),
|
||||||
lora_id=None,
|
lora_id=None,
|
||||||
|
medium=MEDIUM_GPU,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -843,7 +845,9 @@ class RadixCache(BasePrefixCache):
|
|||||||
|
|
||||||
block_hash = hash_str_to_int64(node.hash_value[page_index])
|
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
|
page_index += 1
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user