chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
ishandhanani
2026-03-01 23:27:22 -08:00
committed by GitHub
co-authored by Claude
parent e42fa009d4
commit f7da379b61
9 changed files with 296 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/entrypoints/http_server.py
+20
View File
@@ -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
+2
View File
@@ -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
+15
View File
@@ -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
+34
View File
@@ -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
+22
View File
@@ -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
+152 -7
View File
@@ -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)
python/sglang/srt/mem_cache/radix_cache.py
+4
View File
@@ -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
+30
View File
@@ -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
+17
View File
@@ -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: