diff --git a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py index ea35536eb..be506660f 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py @@ -1,7 +1,7 @@ from __future__ import annotations import heapq -from typing import TYPE_CHECKING, Callable, Optional +from typing import TYPE_CHECKING, Callable, Optional, Sequence import torch @@ -9,10 +9,15 @@ from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, EvictParams, IncLockRefResult, + InsertResult, MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer +from sglang.srt.mem_cache.hicache_storage import ( + PoolName, + PoolTransfer, + PoolTransferResult, +) from sglang.srt.mem_cache.unified_cache_components.tree_component import ( CacheTransferPhase, ComponentType, @@ -21,6 +26,7 @@ from sglang.srt.mem_cache.unified_cache_components.tree_component import ( ) if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.unified_radix_cache import ( UnifiedTreeNode, ) @@ -246,7 +252,14 @@ class FullComponent(TreeComponent): # ---- HiCache Hooks ---- def build_hicache_transfers( - self, node: UnifiedTreeNode, phase: CacheTransferPhase, **kw + self, + node: UnifiedTreeNode, + phase: CacheTransferPhase, + *, + req: Optional[Req] = None, + token_ids: Optional[Sequence[int]] = None, + prefetch_tokens: int = 0, + last_hash: Optional[str] = None, ) -> Optional[list[PoolTransfer]]: ct = self.component_type @@ -290,7 +303,9 @@ class FullComponent(TreeComponent): node: UnifiedTreeNode, phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), - **kw, + *, + insert_result: Optional[InsertResult] = None, + pool_storage_result: Optional[PoolTransferResult] = None, ) -> None: ct = self.component_type diff --git a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py index b3e8b405b..b731b2d48 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Callable, Optional +from typing import TYPE_CHECKING, Callable, Optional, Sequence import torch @@ -13,7 +13,12 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer +from sglang.srt.mem_cache.hicache_storage import ( + PoolHitPolicy, + PoolName, + PoolTransfer, + PoolTransferResult, +) from sglang.srt.mem_cache.unified_cache_components.tree_component import ( CacheTransferPhase, ComponentType, @@ -369,7 +374,14 @@ class MambaComponent(TreeComponent): # ---- HiCache Hooks ---- def build_hicache_transfers( - self, node: UnifiedTreeNode, phase: CacheTransferPhase, **kw + self, + node: UnifiedTreeNode, + phase: CacheTransferPhase, + *, + req: Optional[Req] = None, + token_ids: Optional[Sequence[int]] = None, + prefetch_tokens: int = 0, + last_hash: Optional[str] = None, ) -> Optional[list[PoolTransfer]]: ct = self.component_type @@ -385,7 +397,6 @@ class MambaComponent(TreeComponent): ] if phase == CacheTransferPhase.LOAD_BACK: - req = kw.get("req") transfers: list[PoolTransfer] = [] cd = node.component_data[ct] @@ -458,7 +469,9 @@ class MambaComponent(TreeComponent): node: UnifiedTreeNode, phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), - **kw, + *, + insert_result: Optional[InsertResult] = None, + pool_storage_result: Optional[PoolTransferResult] = None, ) -> None: ct = self.component_type @@ -488,8 +501,6 @@ class MambaComponent(TreeComponent): return transfer = transfers[0] host_indices = transfer.host_indices - insert_result = kw.get("insert_result") - pool_storage_result = kw.get("pool_storage_result") loaded = ( pool_storage_result is not None and pool_storage_result.extra_pool_hit_pages.get(PoolName.MAMBA, 0) >= 1 diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index 9f70db025..4f9c1519a 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Callable, Optional +from typing import TYPE_CHECKING, Callable, Optional, Sequence import torch @@ -13,7 +13,11 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer +from sglang.srt.mem_cache.hicache_storage import ( + PoolName, + PoolTransfer, + PoolTransferResult, +) from sglang.srt.mem_cache.unified_cache_components.tree_component import ( BASE_COMPONENT_TYPE, CacheTransferPhase, @@ -483,7 +487,14 @@ class SWAComponent(TreeComponent): # ---- HiCache Hooks ---- def build_hicache_transfers( - self, node: UnifiedTreeNode, phase: CacheTransferPhase, **kw + self, + node: UnifiedTreeNode, + phase: CacheTransferPhase, + *, + req: Optional[Req] = None, + token_ids: Optional[Sequence[int]] = None, + prefetch_tokens: int = 0, + last_hash: Optional[str] = None, ) -> Optional[list[PoolTransfer]]: ct = self.component_type @@ -542,7 +553,9 @@ class SWAComponent(TreeComponent): node: UnifiedTreeNode, phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), - **kw, + *, + insert_result: Optional[InsertResult] = None, + pool_storage_result: Optional[PoolTransferResult] = None, ) -> None: ct = self.component_type diff --git a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py index 91be6d4f6..6bf5f8ef3 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py @@ -3,7 +3,7 @@ from __future__ import annotations import dataclasses from abc import ABC, abstractmethod from enum import Enum, IntFlag -from typing import TYPE_CHECKING, Any, Callable, Optional +from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence import torch from numpy import float64 @@ -17,7 +17,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.hicache_storage import PoolTransfer +from sglang.srt.mem_cache.hicache_storage import PoolTransfer, PoolTransferResult if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req @@ -383,7 +383,14 @@ class TreeComponent(ABC): # ---- HiCache Hooks ---- def build_hicache_transfers( - self, node: UnifiedTreeNode, phase: CacheTransferPhase, **kw + self, + node: UnifiedTreeNode, + phase: CacheTransferPhase, + *, + req: Optional[Req] = None, + token_ids: Optional[Sequence[int]] = None, + prefetch_tokens: int = 0, + last_hash: Optional[str] = None, ) -> Optional[list[PoolTransfer]]: """Build transfer descriptors for this component in the given phase. Returns None if the component has nothing to transfer.""" @@ -394,7 +401,9 @@ class TreeComponent(ABC): node: UnifiedTreeNode, phase: CacheTransferPhase, transfers: list[PoolTransfer] = (), - **kw, + *, + insert_result: Optional[InsertResult] = None, + pool_storage_result: Optional[PoolTransferResult] = None, ) -> None: """Post-transfer bookkeeping: store host indices, update LRU, etc.""" pass