diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index e40878396..157e117ec 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/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): diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 57e279589..0bb6149b4 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 89fd89496..d543242ff 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/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): """ diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8a7cb5806..23c83ba2f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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() diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index e08e0955b..3faf15cd3 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index b95ff3033..ed094aa56 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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 - if self.enable_storage: + # Compute hash_value if storage or kv events are enabled + 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) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 9292facb0..c9c32f54d 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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() diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index cca592aca..07c24f950 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/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 ) diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 468991795..487482498 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/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: