feat: TTL-based prefix pinning with refresh-on-hit for HiRadixCache (#18941)

Co-authored-by: Claude <noreply@anthropic.com>

python/sglang/srt/entrypoints/http_server.py

@@ -127,6 +127,7 @@ from sglang.srt.managers.io_struct import (
     OpenSessionReqInput,
     ParseFunctionCallReq,
     PauseGenerationReqInput,
+    PinPrefixReqInput,
     ProfileReqInput,
     ReleaseMemoryOccupationReqInput,
     ResumeMemoryOccupationReqInput,
@@ -853,6 +854,25 @@ async def hicache_storage_backend_status():
     }
 
 
+@app.api_route("/hicache/pin_prefix", methods=["POST"])
+@auth_level(AuthLevel.ADMIN_OPTIONAL)
+async def pin_prefix(obj: PinPrefixReqInput):
+    """Pin a prefix by token_ids to resist eviction."""
+    if not _global_state.tokenizer_manager.server_args.admin_api_key:
+        return _admin_api_key_missing_response()
+    ret = await _global_state.tokenizer_manager.pin_prefix(
+        obj.token_ids, obj.ttl_seconds
+    )
+    return ORJSONResponse(
+        content={
+            "status": "ok" if ret.success else "error",
+            "nodes_pinned": ret.nodes_pinned,
+            "message": ret.message,
+        },
+        status_code=200 if ret.success else HTTPStatus.BAD_REQUEST,
+    )
+
+
 @app.api_route("/start_profile", methods=["GET", "POST"])
 @auth_level(AuthLevel.ADMIN_OPTIONAL)
 async def start_profile_async(obj: Optional[ProfileReqInput] = None):

python/sglang/srt/environ.py

@@ -277,6 +277,8 @@ class Envs:
     SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None)
     SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR = EnvStr(None)
     SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR = EnvStr(None)
+    # Max fraction of cache (by token count) that can be pinned; 0 = disable pinning.
+    SGLANG_HICACHE_MAX_PINNED_RATIO = EnvFloat(0.0)
 
     # Mooncake KV Transfer
     SGLANG_MOONCAKE_CUSTOM_MEM_POOL = EnvStr(None)

python/sglang/srt/managers/io_struct.py

@@ -1228,6 +1228,21 @@ class DetachHiCacheStorageReqOutput(BaseReq):
     message: str = ""
 
 
+@dataclass
+class PinPrefixReqInput(BaseReq):
+    """Pin a prefix by token_ids to resist eviction."""
+
+    token_ids: List[int] = field(default_factory=list)
+    ttl_seconds: int = 300  # TTL in seconds, default 5 minutes
+
+
+@dataclass
+class PinPrefixReqOutput(BaseReq):
+    success: bool
+    nodes_pinned: int = 0
+    message: str = ""
+
+
 @dataclass
 class PauseGenerationReqInput(BaseReq):
     """

python/sglang/srt/managers/scheduler.py

@@ -113,6 +113,8 @@ from sglang.srt.managers.io_struct import (
     LoadLoRAAdapterReqOutput,
     OpenSessionReqInput,
     PauseGenerationReqInput,
+    PinPrefixReqInput,
+    PinPrefixReqOutput,
     ProfileReq,
     ReleaseMemoryOccupationReqInput,
     ResumeMemoryOccupationReqInput,
@@ -1045,6 +1047,7 @@ class Scheduler(
                 (ClearHiCacheReqInput, self.clear_hicache_storage_wrapped),
                 (AttachHiCacheStorageReqInput, self.attach_hicache_storage_wrapped),
                 (DetachHiCacheStorageReqInput, self.detach_hicache_storage_wrapped),
+                (PinPrefixReqInput, self.pin_prefix_wrapped),
                 (AbortReq, self.abort_request),
                 (OpenSessionReqInput, self.open_session),
                 (CloseSessionReqInput, self.close_session),
@@ -2599,6 +2602,37 @@ class Scheduler(
 
         return DetachHiCacheStorageReqOutput(success=False, message=msg)
 
+    def pin_prefix_wrapped(self, recv_req: PinPrefixReqInput):
+        if not hasattr(self.tree_cache, "pin_prefix"):
+            return PinPrefixReqOutput(
+                success=False,
+                nodes_pinned=0,
+                message="PIN requires --enable-hierarchical-cache",
+            )
+        if getattr(self.tree_cache, "_max_pinned_tokens", 0) <= 0:
+            return PinPrefixReqOutput(
+                success=False,
+                nodes_pinned=0,
+                message="Pinning is disabled (SGLANG_HICACHE_MAX_PINNED_RATIO is 0)",
+            )
+        nodes_pinned, reject_reason = self.tree_cache.pin_prefix(
+            recv_req.token_ids, recv_req.ttl_seconds
+        )
+        if nodes_pinned == 0:
+            return PinPrefixReqOutput(
+                success=False,
+                nodes_pinned=0,
+                message=reject_reason or "No matching prefix found in cache to pin",
+            )
+        msg = f"Pinned {nodes_pinned} nodes (ttl={recv_req.ttl_seconds}s)"
+        if reject_reason:
+            msg += f"; {reject_reason}"
+        return PinPrefixReqOutput(
+            success=True,
+            nodes_pinned=nodes_pinned,
+            message=msg,
+        )
+
     def _is_no_request(self):
         no_request = (
             self.running_batch.is_empty()

python/sglang/srt/managers/tokenizer_communicator_mixin.py

@@ -59,6 +59,8 @@ from sglang.srt.managers.io_struct import (
     LoadLoRAAdapterReqOutput,
     LoRAUpdateOutput,
     OpenSessionReqInput,
+    PinPrefixReqInput,
+    PinPrefixReqOutput,
     ProfileReq,
     ProfileReqOutput,
     ProfileReqType,
@@ -214,6 +216,9 @@ class TokenizerCommunicatorMixin:
         self.detach_hicache_storage_communicator = _Communicator(
             self.send_to_scheduler, server_args.dp_size
         )
+        self.pin_prefix_communicator = _Communicator(
+            self.send_to_scheduler, server_args.dp_size
+        )
         self.profile_communicator = _Communicator(
             self.send_to_scheduler, server_args.dp_size
         )
@@ -304,6 +309,10 @@ class TokenizerCommunicatorMixin:
                     DetachHiCacheStorageReqOutput,
                     self.detach_hicache_storage_communicator.handle_recv,
                 ),
+                (
+                    PinPrefixReqOutput,
+                    self.pin_prefix_communicator.handle_recv,
+                ),
                 (
                     FlushCacheReqOutput,
                     self.flush_cache_communicator.handle_recv,
@@ -408,6 +417,19 @@ class TokenizerCommunicatorMixin:
             self.server_args.hicache_storage_backend_extra_config = None
         return out
 
+    async def pin_prefix(
+        self: TokenizerManager, token_ids: List[int], ttl_seconds: int = 300
+    ) -> PinPrefixReqOutput:
+        """Pin a prefix by token_ids to resist eviction."""
+        results = await self.pin_prefix_communicator(
+            PinPrefixReqInput(token_ids=token_ids, ttl_seconds=ttl_seconds)
+        )
+        all_success, all_message = _Communicator.merge_results(results)
+        total = sum(r.nodes_pinned for r in results)
+        return PinPrefixReqOutput(
+            success=all_success, nodes_pinned=total, message=all_message
+        )
+
     async def start_profile(
         self: TokenizerManager,
         output_dir: Optional[str] = None,

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -8,10 +8,11 @@ import os
 import threading
 import time
 from queue import Empty
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
 import torch
 
+from sglang.srt.environ import envs
 from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation
 from sglang.srt.mem_cache.base_prefix_cache import (
     EvictParams,
@@ -165,6 +166,18 @@ class HiRadixCache(RadixCache):
 
         self.evictable_host_leaves = set()
 
+        # Pin budget: max tokens that can be pinned = ratio * host pool capacity.
+        pin_ratio = envs.SGLANG_HICACHE_MAX_PINNED_RATIO.get()
+        if pin_ratio < 0 or pin_ratio >= 1:
+            raise ValueError(
+                f"SGLANG_HICACHE_MAX_PINNED_RATIO must be in [0, 1), got {pin_ratio}"
+            )
+        self._max_pinned_tokens = int(self.token_to_kv_pool_host.size * pin_ratio)
+        self.pinned_size_ = 0
+        logger.info(
+            "Pin budget: %d tokens (ratio=%.3f)", self._max_pinned_tokens, pin_ratio
+        )
+
         super().__init__(params=params)
 
     def shutdown(self):
@@ -581,6 +594,7 @@ class HiRadixCache(RadixCache):
         # Clear per-request tracking dicts
         self.prefetch_loaded_tokens_by_reqid.clear()
         self.evictable_host_leaves.clear()
+        self.pinned_size_ = 0
         super().reset()
 
     def get_height(self, node: TreeNode):
@@ -720,6 +734,87 @@ class HiRadixCache(RadixCache):
     def evictable_size(self):
         return self.evictable_size_
 
+    def _is_pinned(self, node: TreeNode) -> bool:
+        """Check if a node has an active (non-expired) pin."""
+        return node.pin_expiry > 0 and time.monotonic() <= node.pin_expiry
+
+    def _clear_pin(self, node: TreeNode):
+        """Clear expired pin state and release host_ref_counter hold."""
+        if node.pin_expiry > 0:
+            self.pinned_size_ = max(0, self.pinned_size_ - len(node.key))
+            node.host_ref_counter = max(0, node.host_ref_counter - 1)
+        node.pin_expiry = 0.0
+        node.pin_ttl = 0
+
+    def pin_prefix(
+        self, token_ids: List[int], ttl_seconds: int = 300
+    ) -> Tuple[int, Optional[str]]:
+        """Pin nodes along a prefix path. Returns (nodes_pinned, reject_reason)."""
+        if self.disable or not token_ids:
+            return (0, None)
+
+        key, _ = self.maybe_bigram_convert(self._to_radix_key(token_ids))
+        if self.page_size != 1:
+            page_aligned_len = len(key) // self.page_size * self.page_size
+            key = key[:page_aligned_len]
+        if len(key) == 0:
+            return (0, None)
+
+        expiry = time.monotonic() + ttl_seconds
+        nodes_pinned = 0
+        budget_exceeded = False
+        node = self.root_node
+        child_key = self.get_child_key_fn(key)
+
+        while len(key) > 0 and child_key in node.children:
+            child = node.children[child_key]
+            prefix_len = self.key_match_fn(child.key, key)
+
+            # First pin on this node: check budget, then acquire hold
+            if child.pin_expiry == 0:
+                if self.pinned_size_ + len(child.key) > self._max_pinned_tokens:
+                    budget_exceeded = True
+                    break
+                child.host_ref_counter += 1
+                self.pinned_size_ += len(child.key)
+
+                # Eagerly back up to host so eviction finds pinned nodes
+                # already backuped and never enters the write_back drain
+                # path, which would leak lock_ref on in-flight
+                # write-through entries. No-op under write_back policy.
+                self._inc_hit_count(child)
+
+            # Extend expiry and store TTL for refresh-on-hit
+            child.pin_expiry = max(child.pin_expiry, expiry)
+            child.pin_ttl = max(child.pin_ttl, ttl_seconds)
+            nodes_pinned += 1
+
+            if prefix_len < len(child.key):
+                break
+
+            node = child
+            key = key[prefix_len:]
+            if len(key):
+                child_key = self.get_child_key_fn(key)
+
+        logger.info(
+            "[PIN] pin_prefix: nodes_pinned=%d, ttl=%ds", nodes_pinned, ttl_seconds
+        )
+        if budget_exceeded:
+            msg = f"Pin budget exhausted ({self.pinned_size_}/{self._max_pinned_tokens} tokens pinned)"
+            if nodes_pinned == 0:
+                return (0, msg)
+            return (nodes_pinned, f"prefix partially pinned; {msg}")
+        return (nodes_pinned, None)
+
+    def _to_radix_key(self, token_ids: List[int]) -> RadixKey:
+        """Convert raw token_ids to a RadixKey for tree walking.
+
+        Must use list (not tuple) to match scheduler's RadixKey format,
+        since _key_match_paged compares slices directly and list != tuple.
+        """
+        return RadixKey(token_ids=list(token_ids))
+
     def inc_lock_ref(self, node: TreeNode):
         if self.disable:
             return 0
@@ -788,6 +883,26 @@ class HiRadixCache(RadixCache):
             if x.lock_ref > 0:
                 continue
 
+            if self._is_pinned(x):
+                # Still active: demote to host if possible
+                if x.backuped:
+                    num_evicted += self._evict_backuped(x)
+                    continue
+                written = self.write_backup(x, write_back=True)
+                if written > 0:
+                    num_evicted += written
+                    write_back_nodes.append(x)
+                    continue  # backup succeeded, pin holds on host
+                # Host full -- drop pin so GPU can be freed
+                self._clear_pin(x)
+                logger.warning(
+                    "[PIN] evict: can't backup node %d to host, releasing pin",
+                    x.id,
+                )
+            elif x.pin_expiry > 0:
+                # Expired pin: clear and fall through to normal eviction
+                self._clear_pin(x)
+
             if not x.backuped:
                 if self.cache_controller.write_policy == "write_back":
                     # write to host if the node is not backuped
@@ -818,7 +933,7 @@ class HiRadixCache(RadixCache):
         return EvictResult(num_tokens_evicted=num_evicted)
 
     def _evict_backuped(self, node: TreeNode):
-        # evict a node already written to host
+        # GPU -> CPU demotion: no BlockRemoved since block is still reachable via load_back
         num_evicted = self.cache_controller.evict_device(node.value)
         assert num_evicted > 0
         self.evictable_size_ -= num_evicted
@@ -830,7 +945,8 @@ class HiRadixCache(RadixCache):
         return num_evicted
 
     def _evict_regular(self, node: TreeNode):
-        # evict a node not initiated write to host
+        # evict a node not initiated write to host -- emit BlockRemoved
+        self._record_remove_event(node)
         self.cache_controller.mem_pool_device_allocator.free(node.value)
         num_evicted = len(node.value)
         self._delete_leaf(node)
@@ -852,10 +968,17 @@ class HiRadixCache(RadixCache):
             if not x.evicted:
                 continue
 
-            # node is protected from eviction as it has ongoing prefetch or backup to storage
+            # Expire stale pins before checking host_ref_counter
+            if x.pin_expiry > 0 and time.monotonic() > x.pin_expiry:
+                self._clear_pin(x)
+
+            # node is protected from eviction as it has ongoing prefetch, backup, or pin
             if x.host_ref_counter > 0:
                 continue
 
+            # Block deleted entirely (GPU already evicted, now CPU freed) --
+            # emit BlockRemoved so the router removes this block from its index.
+            self._record_remove_event(x)
             num_evicted += self.cache_controller.evict_host(x.host_value)
 
             key = self.get_child_key_fn(x.key)
@@ -872,7 +995,6 @@ class HiRadixCache(RadixCache):
     def load_back(
         self, node: TreeNode, mem_quota: Optional[int] = None
     ) -> Optional[torch.Tensor]:
-        # todo: more loading policies
 
         start_time = time.perf_counter()
         last_hit_node = node
@@ -909,6 +1031,13 @@ class HiRadixCache(RadixCache):
         self.dec_lock_ref(ancester_node)
         if device_indices is None:
             # no sufficient GPU memory to load back KV caches
+            logger.warning(
+                "load_back: FAILED to load %d tokens for node %d "
+                "even after eviction (evictable_size=%d)",
+                len(host_indices),
+                last_hit_node.id,
+                self.evictable_size_,
+            )
             return None
 
         self.ongoing_load_back[last_hit_node.id] = last_hit_node
@@ -1135,6 +1264,7 @@ class HiRadixCache(RadixCache):
                 host_hit_length=0,
             )
 
+        page_aligned_len = len(key)
         if self.page_size != 1:
             page_aligned_len = len(key) // self.page_size * self.page_size
             key = key[:page_aligned_len]
@@ -1213,6 +1343,9 @@ class HiRadixCache(RadixCache):
         while len(key) > 0 and child_key in node.children.keys():
             node = node.children[child_key]
             node.last_access_time = time.monotonic()
+            # Refresh pin TTL on host insert hit
+            if self._is_pinned(node):
+                node.pin_expiry = time.monotonic() + node.pin_ttl
             prefix_len = self.key_match_fn(node.key, key)
             key = key[prefix_len:]
             host_value = host_value[prefix_len:]
@@ -1248,6 +1381,9 @@ class HiRadixCache(RadixCache):
         while len(key) > 0 and child_key in node.children.keys():
             child = node.children[child_key]
             child.last_access_time = time.monotonic()
+            # Refresh pin TTL on cache hit
+            if self._is_pinned(child):
+                child.pin_expiry = time.monotonic() + child.pin_ttl
             prefix_len = self.key_match_fn(child.key, key)
             if prefix_len < len(child.key):
                 new_node = self._split_node(child.key, child, prefix_len)
@@ -1272,6 +1408,11 @@ class HiRadixCache(RadixCache):
         new_node.children = {self.get_child_key_fn(key[split_len:]): child}
         new_node.parent = child.parent
         new_node.lock_ref = child.lock_ref
+        new_node.pin_expiry = child.pin_expiry
+        new_node.pin_ttl = child.pin_ttl
+        # If child is pinned, new parent inherits a host_ref_counter hold
+        if child.pin_expiry > 0:
+            new_node.host_ref_counter += 1
         new_node.key = child.key[:split_len]
         new_node.hit_count = child.hit_count
 
@@ -1291,6 +1432,7 @@ class HiRadixCache(RadixCache):
         child.parent = new_node
         child.key = child.key[split_len:]
         new_node.parent.children[self.get_child_key_fn(key)] = new_node
+
         return new_node
 
     def insert(self, params: InsertParams) -> InsertResult:
@@ -1366,10 +1508,13 @@ class HiRadixCache(RadixCache):
             self._update_leaf_status(node)
             self._update_leaf_status(new_node)
 
-            # Compute hash_value if storage is enabled
+            # Compute hash_value if storage or kv events are enabled
-            if self.enable_storage:
+            if self.enable_storage or self.enable_kv_cache_events:
                 new_node.hash_value = compute_node_hash_values(new_node, self.page_size)
 
+            # Emit BlockStored so the router indexes this block.
+            self._record_store_event(new_node)
+
             if self.cache_controller.write_policy != "write_back":
                 self._inc_hit_count(new_node, chunked)
         return InsertResult(prefix_len=total_prefix_length)

python/sglang/srt/mem_cache/radix_cache.py

@@ -124,6 +124,10 @@ class TreeNode:
         self.key: RadixKey = None
         self.value: Optional[torch.Tensor] = None
         self.lock_ref = 0
+        self.pin_expiry: float = (
+            0.0  # absolute expiry time (time.monotonic()), 0 = not pinned
+        )
+        self.pin_ttl: int = 0  # original TTL in seconds, for refresh-on-hit
         self.last_access_time = time.monotonic()
         self.creation_time = time.monotonic()
 

python/sglang/srt/observability/metrics_collector.py

@@ -101,6 +101,10 @@ class SchedulerStats:
     lora_pool_slots_total: int = 0
     lora_pool_utilization: float = 0.0
 
+    # HiCache metrics
+    hicache_host_used_tokens: int = 0
+    hicache_host_total_tokens: int = 0
+
     # Routing key metrics
     num_unique_running_routing_keys: int = 0
     routing_key_running_req_counts: List[int] = field(default_factory=list)
@@ -144,6 +148,7 @@ class SchedulerMetricsCollector:
         self,
         labels: Dict[str, str],
         enable_lora: bool = False,
+        enable_hierarchical_cache: bool = False,
         server_args: Optional["ServerArgs"] = None,
     ) -> None:
         # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
@@ -151,6 +156,7 @@ class SchedulerMetricsCollector:
 
         self.labels = labels
         self.enable_lora = enable_lora
+        self.enable_hierarchical_cache = enable_hierarchical_cache
         self.last_log_time = time.perf_counter()
 
         self.num_running_reqs = Gauge(
@@ -599,6 +605,21 @@ class SchedulerMetricsCollector:
                 multiprocess_mode="mostrecent",
             )
 
+        # HiCache host-tier metrics (only created when hierarchical cache is enabled)
+        if self.enable_hierarchical_cache:
+            self.hicache_host_used_tokens = Gauge(
+                name="sglang:hicache_host_used_tokens",
+                documentation="Number of tokens currently used in the host KV cache.",
+                labelnames=labels.keys(),
+                multiprocess_mode="mostrecent",
+            )
+            self.hicache_host_total_tokens = Gauge(
+                name="sglang:hicache_host_total_tokens",
+                documentation="Total capacity of the host KV cache in tokens.",
+                labelnames=labels.keys(),
+                multiprocess_mode="mostrecent",
+            )
+
         self.num_unique_running_routing_keys = Gauge(
             name="sglang:num_unique_running_routing_keys",
             documentation="Number of unique routing keys in running batch.",
@@ -894,6 +915,15 @@ class SchedulerMetricsCollector:
             self._log_gauge(self.lora_pool_slots_total, stats.lora_pool_slots_total)
             self._log_gauge(self.lora_pool_utilization, stats.lora_pool_utilization)
 
+        # HiCache host-tier metrics (only logged if hierarchical cache is enabled)
+        if self.enable_hierarchical_cache:
+            self._log_gauge(
+                self.hicache_host_used_tokens, stats.hicache_host_used_tokens
+            )
+            self._log_gauge(
+                self.hicache_host_total_tokens, stats.hicache_host_total_tokens
+            )
+
         self._log_gauge(
             self.num_unique_running_routing_keys, stats.num_unique_running_routing_keys
         )

python/sglang/srt/observability/scheduler_metrics_mixin.py

@@ -125,6 +125,7 @@ class SchedulerMetricsMixin:
             self.metrics_collector = SchedulerMetricsCollector(
                 labels=labels,
                 enable_lora=self.enable_lora,
+                enable_hierarchical_cache=self.enable_hierarchical_cache,
                 server_args=self.server_args,
             )
 
@@ -297,6 +298,7 @@ class SchedulerMetricsMixin:
             # Others
             self.calculate_utilization()
             self.update_lora_metrics()
+            self._log_hicache_stats()
             self.metrics_collector.log_stats(self.stats)
             self._emit_kv_metrics()
         self._publish_kv_events()
@@ -469,6 +471,7 @@ class SchedulerMetricsMixin:
             # Others
             self.calculate_utilization()
             self.update_lora_metrics()
+            self._log_hicache_stats()
             self.metrics_collector.log_stats(self.stats)
             self._emit_kv_metrics()
         self._publish_kv_events()
@@ -530,6 +533,20 @@ class SchedulerMetricsMixin:
             batch = KVEventBatch(ts=time.time(), events=events)
             self.kv_event_publisher.publish(batch)
 
+    def _log_hicache_stats(self: Scheduler):
+        """Populate HiCache host-tier stats on self.stats.
+
+        These are pushed to Prometheus by SchedulerMetricsCollector.log_stats().
+        """
+        if not self.enable_hierarchical_cache:
+            return
+
+        host_pool = self.tree_cache.token_to_kv_pool_host
+        self.stats.hicache_host_used_tokens = (
+            host_pool.size - host_pool.available_size()
+        )
+        self.stats.hicache_host_total_tokens = host_pool.size
+
     def update_lora_metrics(self: Scheduler):
         """Update LoRA pool metrics for monitoring and autoscaling."""
         if not self.enable_lora: