[Unified Cache][1/N]: Support cache contract for external linker

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -19,6 +19,7 @@
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.events import KVCacheEventRecorder
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
 from sglang.srt.observability.metrics_collector import (
     STAT_LOGGER_ROLE_RADIX_CACHE,
     RadixCacheMetricsCollector,
@@ -34,9 +35,6 @@
         CacheAction,
         ComponentAction,
     )
-    from sglang.srt.mem_cache.unified_cache.components.tree_component import (
-        ComponentType,
-    )
 
 
 @runtime_checkable
@@ -78,6 +76,7 @@ class InsertParams:
     # General
     chunked: bool = False
     priority: int = 0
+    track_adopted_ranges: bool = False
 
 
 @dataclasses.dataclass
@@ -90,11 +89,24 @@ class InsertResult:
     mamba_exist: bool = False
     inserted_host_node: Any = None
     host_insert_dropped: bool = False
+    adopted_ranges: Optional[dict[ComponentType, list[tuple[int, int]]]] = None
     # Controller-applied actions from the non-stepped channels (e.g. insert_host); the stepped insert emits via InsertStepResult.actions.
     cache_actions: list[CacheAction | ComponentAction] = dataclasses.field(
         default_factory=list
     )
 
+    def record_adopted_range(
+        self, component_type: ComponentType, start: int, end: int
+    ) -> None:
+        if self.adopted_ranges is None or start >= end:
+            return
+        ranges = self.adopted_ranges.setdefault(component_type, [])
+        if ranges and start <= ranges[-1][1]:
+            prev_start, prev_end = ranges[-1]
+            ranges[-1] = (min(prev_start, start), max(prev_end, end))
+        else:
+            ranges.append((start, end))
+
 
 @dataclasses.dataclass
 class EvictParams:

python/sglang/srt/mem_cache/hicache_storage.py

@@ -125,6 +125,13 @@ class PoolTransferResult:
     kv_hit_pages: int
     extra_pool_hit_pages: dict[str, int]
 
+    # Pools with TRAILING_PAGES (SWA, Mamba state) only hold a window that ends on an
+    # offloaded node boundary.
+    # Each rank owns its own shard and may hold a different set, so reducing a
+    # per-rank maximum would pick a length that is illegal on another rank; the
+    # caller intersects these sets instead.
+    restorable_prefix_pages: Optional[List[int]] = None
+
     @classmethod
     def empty(cls) -> PoolTransferResult:
         return cls(0, {})

python/sglang/srt/mem_cache/unified_cache/components/__init__.py

@@ -8,6 +8,8 @@
     ComponentData,
     ComponentType,
     EvictLayer,
+    ExternalLinkerLoadPhase,
+    LinkerTransferPhase,
     LRURefreshPhase,
     PrepareLoadBackResult,
     PreparePrefetchResult,
@@ -20,6 +22,8 @@
     "BASE_COMPONENT_TYPE",
     "ComponentData",
     "ComponentType",
+    "ExternalLinkerLoadPhase",
+    "LinkerTransferPhase",
     "EvictLayer",
     "FullComponent",
     "CacheTransferPhase",

python/sglang/srt/mem_cache/unified_cache/components/full_component.py

@@ -7,6 +7,7 @@
 
 from sglang.srt.mem_cache.base_prefix_cache import (
     DecLockRefParams,
+    EvictParams,
     IncLockRefResult,
     InsertResult,
     MatchPrefixParams,
@@ -22,10 +23,13 @@
     CacheTransferPhase,
     ComponentType,
     EvictLayer,
+    ExternalLinkerLoadPhase,
+    LinkerTransferPhase,
     TreeComponent,
 )
 
 if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
     from sglang.srt.mem_cache.unified_cache.cache_action import (
         CacheAction,
         ComponentAction,
@@ -424,6 +428,68 @@ def commit_hicache_transfer(
 
             self.tree_core._update_evictable_leaf_sets(node)
 
+    def _full_allocator(self):
+        """The allocator that owns the full-attention pool alone."""
+        allocator = self.cache.token_to_kv_pool_allocator
+        return allocator.full_attn_allocator if self.cache.is_swa_enabled else allocator
+
+    def build_external_linker_transfer(
+        self,
+        phase: LinkerTransferPhase,
+        node: Optional[UnifiedTreeNode],
+        keys: Optional[Sequence[str]],
+    ) -> Optional[PoolTransfer]:
+        if phase == LinkerTransferPhase.OFFLOAD:
+            if node is None or not node.hash_value:
+                return None
+            value = node.component_data[self.component_type].value
+            if value is None:
+                return None
+            return PoolTransfer(
+                name=PoolName.KV,
+                device_indices=value.to(torch.int64),
+                keys=list(node.hash_value),
+            )
+
+        if not keys:
+            return None
+
+        if phase == LinkerTransferPhase.LOOKUP:
+            return PoolTransfer(name=PoolName.KV, keys=list(keys))
+
+        if phase == LinkerTransferPhase.LOAD:
+            allocator = self._full_allocator()
+            num_tokens = len(keys) * self.cache.page_size
+            shortfall = max(0, num_tokens - allocator.available_size())
+            if shortfall:
+                self.cache.evict(EvictParams(num_tokens=shortfall))
+            slots = allocator.alloc(num_tokens)
+            if slots is None:
+                return None
+
+            return PoolTransfer(
+                name=PoolName.KV,
+                device_indices=slots.to(torch.int64),
+                keys=list(keys),
+            )
+
+    def update_external_linker_load(
+        self,
+        phase: ExternalLinkerLoadPhase,
+        req: Req,
+        full_transfer: PoolTransfer,
+        transfer: PoolTransfer,
+        prefix_len: int,
+        *,
+        insert_result: Optional[InsertResult] = None,
+        canonical_full: Optional[torch.Tensor] = None,
+    ) -> Optional[PoolTransfer]:
+        if phase == ExternalLinkerLoadPhase.ABORT:
+            self._full_allocator().free(transfer.device_indices)
+            return None
+
+        return transfer
+
     def free_host_values(self, host_values: list[torch.Tensor]) -> None:
         if self._full_kv_pool_host is None:
             return

python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py

@@ -29,6 +29,7 @@
     CacheTransferPhase,
     ComponentType,
     EvictLayer,
+    LinkerTransferPhase,
     LRURefreshPhase,
     PrepareLoadBackResult,
     PreparePrefetchResult,
@@ -648,6 +649,16 @@ def cleanup_after_caching_req(
                 self._free_mamba_value(insert_params.mamba_value)
             req.mamba_last_track_seqlen = None
 
+    def build_external_linker_transfer(
+        self,
+        phase: LinkerTransferPhase,
+        node: Optional[UnifiedTreeNode],
+        keys: Optional[Sequence[str]],
+    ) -> Optional[PoolTransfer]:
+        raise AssertionError(
+            "MambaComponent does not support external linker mode, will support soon"
+        )
+
     # ---- HiCache Hooks ----
 
     def prepare_load_back(