[UnifiedRadixTree]: Add HiCache hook interface for TreeComponent (#22924)
This commit is contained in:
@@ -17,6 +17,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
|
|||||||
MatchPrefixParams,
|
MatchPrefixParams,
|
||||||
MatchResult,
|
MatchResult,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.mem_cache.hicache_storage import PoolTransfer
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.managers.schedule_batch import Req
|
from sglang.srt.managers.schedule_batch import Req
|
||||||
@@ -62,6 +63,16 @@ class ComponentData:
|
|||||||
value: Optional[torch.Tensor] = None
|
value: Optional[torch.Tensor] = None
|
||||||
lock_ref: int = 0
|
lock_ref: int = 0
|
||||||
metadata: dict[str, Any] = dataclasses.field(default_factory=dict)
|
metadata: dict[str, Any] = dataclasses.field(default_factory=dict)
|
||||||
|
host_value: Optional[torch.Tensor] = None
|
||||||
|
host_lock_ref: int = 0
|
||||||
|
|
||||||
|
|
||||||
|
class CacheTransferPhase(str, Enum):
|
||||||
|
|
||||||
|
BACKUP_HOST = "backup_host" # D→H
|
||||||
|
LOAD_BACK = "load_back" # H→D
|
||||||
|
BACKUP_STORAGE = "backup_storage" # H→Storage
|
||||||
|
PREFETCH = "prefetch" # Storage→H
|
||||||
|
|
||||||
|
|
||||||
def get_and_increase_time_counter() -> float64:
|
def get_and_increase_time_counter() -> float64:
|
||||||
@@ -290,3 +301,29 @@ class TreeComponent(ABC):
|
|||||||
``insert_params`` is None only on the disabled path; on early-return
|
``insert_params`` is None only on the disabled path; on early-return
|
||||||
paths it is still provided so components can free their resources."""
|
paths it is still provided so components can free their resources."""
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
# ---- HiCache Hooks ----
|
||||||
|
|
||||||
|
def build_hicache_transfers(
|
||||||
|
self, node: UnifiedTreeNode, phase: CacheTransferPhase, **kw
|
||||||
|
) -> Optional[list[PoolTransfer]]:
|
||||||
|
"""Build transfer descriptors for this component in the given phase.
|
||||||
|
Returns None if the component has nothing to transfer."""
|
||||||
|
return None
|
||||||
|
|
||||||
|
def commit_hicache_transfer(
|
||||||
|
self,
|
||||||
|
node: UnifiedTreeNode,
|
||||||
|
phase: CacheTransferPhase,
|
||||||
|
transfers: list[PoolTransfer] = (),
|
||||||
|
) -> None:
|
||||||
|
"""Post-transfer bookkeeping: store host indices, update LRU, etc."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
def drive_host_eviction(
|
||||||
|
self, num_tokens: int, tracker: dict[ComponentType, int]
|
||||||
|
) -> None:
|
||||||
|
"""Evict from this component's host-side resources.
|
||||||
|
Called by HostPoolGroup when the host pool is full.
|
||||||
|
Default no-op for components without host storage."""
|
||||||
|
pass
|
||||||
|
|||||||
Reference in New Issue
Block a user