[HiCache] L3 storage prefetch lifecycle metrics and cross-tier attribution fixes (#37503)

This commit is contained in:
Zhiqiang Xie
2026-09-03 16:00:04 -07:00
committed by GitHub
parent 0610a6539d
commit a480f388b2
19 changed files with 828 additions and 143 deletions
@@ -235,7 +235,8 @@ class StorageOperation:
self.all_hash_values: Optional[List[str]] = None
# Prefetch-outcome accounting, set at enqueue by the tree cache.
self.stats_requested_tokens = 0
self.stats_total_tokens = 0
# Absolute token offset at which this storage-prefetched span starts.
self.storage_start = 0
self.id = StorageOperation.counter
StorageOperation.counter += 1
+57 -10
View File
@@ -214,17 +214,39 @@ def sanity_check_mm_pad_shift_value(vocab_size: int) -> None:
def split_cached_prefix_by_tier(
prefix_len: int, host_hit_len: int, storage_hit_len: int
prefix_len: int,
host_hit_len: int,
storage_hit_len: int,
storage_hit_start: Optional[int] = None,
host_hit_is_storage: bool = False,
) -> tuple[int, int, int]:
"""Split a request's cached prefix into (device, host, storage) tokens.
prefix_len is len(prefix_indices) AFTER host load-back, so it contains the
host-loaded portion; host_hit_len in turn contains the storage-prefetched
portion (storage is clamped to it to handle edge cases).
``host_hit_len`` is the materialized host hit. An exact L3-loaded span is
preserved after H2D promotion, while an evicted tail is clipped by
``prefix_len``. In buffer mode host memory is only L3 staging.
"""
storage = min(host_hit_len, storage_hit_len)
host = host_hit_len - storage
device = max(0, prefix_len - host_hit_len)
host_hit_len = min(prefix_len, host_hit_len)
host_start = prefix_len - host_hit_len
if storage_hit_start is None:
if host_hit_is_storage:
storage = host_hit_len
host = 0
else:
storage = min(host_hit_len, storage_hit_len)
host = host_hit_len - storage
else:
storage_end = storage_hit_start + storage_hit_len
storage = max(
0,
min(prefix_len, storage_end) - max(0, storage_hit_start),
)
storage_in_host = max(
0,
min(prefix_len, storage_end) - max(host_start, storage_hit_start),
)
host = 0 if host_hit_is_storage else host_hit_len - storage_in_host
device = prefix_len - host - storage
return device, host, storage
@@ -1088,6 +1110,13 @@ class Req(ReqDllmMixin):
self.num_matched_prefix_tokens = 0
# Tokens loaded from storage backend (L3) during prefetch for this request
self.storage_hit_length = 0
self.storage_hit_start: Optional[int] = None
# FULL host-hit tokens actually spliced to device by init_load_back at
# admission; less than host_hit_length when the load-back was declined
# or the staged splice dropped (that shortfall is re-prefilled).
self.host_loaded_length = 0
# Buffer-mode host memory is transport staging, not an L2 cache tier.
self.host_hit_is_storage = False
# Storage prefetch retry state while queued
# (see Scheduler._retry_missed_storage_prefetches).
self.storage_prefetch_retry_pending = False
@@ -1322,6 +1351,22 @@ class Req(ReqDllmMixin):
or self.mamba_host_hit_length > 0
)
def materialized_host_hit_len(self) -> int:
"""Host-hit tokens that actually reached the device: the metrics tier
split must not credit host/storage for a declined or dropped
load-back, whose span was re-prefilled and counted as input."""
return min(self.host_hit_length, self.host_loaded_length)
def fulfilled_storage_hit_len(self, prefix_len: int) -> int:
"""L3-fetched tokens covered by the admitted cached prefix."""
if self.storage_hit_start is None:
return min(prefix_len, self.storage_hit_length)
return max(
0,
min(prefix_len, self.storage_hit_start + self.storage_hit_length)
- self.storage_hit_start,
)
def detach_kv(self) -> ReqKvInfo:
# Hand the KV record to a new holder; the req keeps a fresh empty one.
kv, self.kv = self.kv, ReqKvInfo()
@@ -2608,16 +2653,18 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# Only compute once on FIRST chunk - subsequent chunks in chunked prefill
# would incorrectly count previously computed tokens as cache hits.
if not req._cache_breakdown_computed:
# storage_hit_length is set by scheduler.pop_prefetch_loaded_tokens()
# after prefetch completes.
# storage_hit_length is set after pop_prefetch_loaded_span()
# returns a completed prefetch.
(
req.cached_tokens_device,
req.cached_tokens_host,
req.cached_tokens_storage,
) = split_cached_prefix_by_tier(
prefix_len=len(req.prefix_indices),
host_hit_len=req.host_hit_length,
host_hit_len=req.materialized_host_hit_len(),
storage_hit_len=req.storage_hit_length,
storage_hit_start=req.storage_hit_start,
host_hit_is_storage=req.host_hit_is_storage,
)
req._cache_breakdown_computed = True
+32 -11
View File
@@ -835,8 +835,6 @@ class PrefillAdder:
max_new_tokens: int,
retracted_stain: bool,
mamba_gap_reserve: int = 0,
host_hit_len: int = 0,
storage_hit_len: int = 0,
):
# TODO(lsyin): check this workaround logic, which only ensures the prefill will not out of memory, and may be too conservative
extend_input_len = self.ceil_paged_tokens(extend_input_len)
@@ -876,16 +874,41 @@ class PrefillAdder:
if retracted_stain:
self.reprocessed_log_hit_tokens += prefix_len
self.reprocessed_log_input_tokens += extend_input_len
elif prefix_len > 0:
def _account_prefill_cache_admission(self, req: Req, prefix_len: int) -> None:
if req.retracted_stain:
# Retraction attribution is intentionally omitted for now; discard
# its lifecycle state so a later abort cannot report it as a drop.
self.tree_cache.discard_storage_prefetch_accounting(req.rid)
return
if prefix_len > 0:
device_hit, host_hit, storage_hit = split_cached_prefix_by_tier(
prefix_len=prefix_len,
host_hit_len=host_hit_len,
storage_hit_len=storage_hit_len,
host_hit_len=req.materialized_host_hit_len(),
storage_hit_len=req.storage_hit_length,
storage_hit_start=req.storage_hit_start,
host_hit_is_storage=req.host_hit_is_storage,
)
self.log_device_hit_tokens += device_hit
self.log_host_hit_tokens += host_hit
self.log_storage_hit_tokens += storage_hit
fulfilled_storage_hit = req.fulfilled_storage_hit_len(prefix_len)
reason = None
if fulfilled_storage_hit < req.storage_hit_length:
reason = (
"device_capacity"
if req.needs_host_load_back()
and req.host_loaded_length < req.host_hit_length
else "shrunk"
)
self.tree_cache.finish_storage_prefetch_admission(
req.rid,
fulfilled_tokens=fulfilled_storage_hit,
reason=reason,
)
def _get_dllm_remain_tokens(self) -> int:
_rem_tokens = min(
self.rem_dllm_tokens,
@@ -917,9 +940,8 @@ class PrefillAdder:
0,
req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req),
host_hit_len=req.host_hit_length,
storage_hit_len=req.storage_hit_length,
)
self._account_prefill_cache_admission(req, prefix_len)
def _req_inc_lock_ref(self, req: Req):
result = self.tree_cache.inc_lock_ref(req.last_node)
@@ -1292,6 +1314,7 @@ class PrefillAdder:
req=req,
)
)
req.host_loaded_length = len(new_indices)
req.prefix_indices = torch.cat([req.prefix_indices, new_indices])
prefix_len = len(req.prefix_indices)
req.kv.cache_protected_len = prefix_len
@@ -1347,9 +1370,8 @@ class PrefillAdder:
),
req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req),
host_hit_len=req.host_hit_length,
storage_hit_len=req.storage_hit_length,
)
self._account_prefill_cache_admission(req, prefix_len)
else:
# Make sure at least one page is available
trunc_len = chunk_tokens_limit // self.page_size * self.page_size
@@ -1395,9 +1417,8 @@ class PrefillAdder:
0,
req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req),
host_hit_len=req.host_hit_length,
storage_hit_len=req.storage_hit_length,
)
self._account_prefill_cache_admission(req, prefix_len)
return self.budget_state()
+16 -2
View File
@@ -3756,10 +3756,17 @@ class Scheduler(
if not prefetch_done:
# skip staging requests that are ongoing prefetch
continue
# Pop the number of tokens loaded from storage (L3 hits)
loaded_tokens = self.tree_cache.pop_prefetch_loaded_tokens(req.rid)
# Pop the L3-loaded span. Unified cache exposes its absolute
# start so cache-mode L2/L3 attribution survives L3-tail eviction.
loaded_tokens, loaded_start = self.tree_cache.pop_prefetch_loaded_span(
req.rid
)
if loaded_tokens > 0:
req.storage_hit_length = loaded_tokens
req.storage_hit_start = loaded_start
# Cache-mode host memory is a resident L2 tier. Buffer mode
# marks the staged span below once it is surfaced.
req.host_hit_is_storage = False
req.init_next_round_input(self.tree_cache)
if (
@@ -3781,6 +3788,13 @@ class Scheduler(
if held_tokens > 0:
req.host_hit_length = held_tokens
req.swa_host_hit_length = held_swa_tokens
req.storage_hit_length = held_tokens
req.storage_hit_start = len(req.prefix_indices)
req.host_hit_is_storage = True
elif not (req.host_hit_is_storage and req.host_loaded_length > 0):
req.storage_hit_length = 0
req.storage_hit_start = None
req.host_hit_is_storage = False
res = adder.add_one_req(
req,
has_chunked_req=(self.chunked_req is not None),
@@ -48,6 +48,11 @@ RECORD_STEP_TIME = envs.SGLANG_RECORD_STEP_TIME.get()
LOG_FORWARD_ITERS = envs.SGLANG_LOG_FORWARD_ITERS.get()
ENABLE_METRICS_DEVICE_TIMER = envs.SGLANG_ENABLE_METRICS_DEVICE_TIMER.get()
CACHE_HIT_RATE_WINDOW_SECONDS = envs.SGLANG_CACHE_HIT_RATE_WINDOW_SECONDS.get()
# gen_throughput is computed only on decode-stats ticks; when decode is starved
# (e.g. long chunked-prefill stretches) the last window's value would otherwise
# be re-exported indefinitely. 30s is far above any healthy decode-stats gap,
# so past it the true recent decode throughput is ~0.
GEN_THROUGHPUT_STALENESS_SECONDS = 30.0
class _CacheHitRateWindow:
@@ -153,6 +158,16 @@ class SchedulerMetricsReporter:
# cache_hit_rate stats keep their per-report semantics.
self.recent_cache_hit_rate = 0.0
def _current_gen_throughput(self, now: float) -> float:
"""last_gen_throughput, decayed to 0 once decode-stats stop arriving.
Mirrors the pause-path zeroing in Scheduler.pause_generation: a stale
decode window must not keep exporting its throughput forever.
"""
if now - self.last_decode_stats_tic > GEN_THROUGHPUT_STALENESS_SECONDS:
self.last_gen_throughput = 0.0
return self.last_gen_throughput
def _init_metrics(
self,
tp_rank: int,
@@ -717,6 +732,9 @@ class SchedulerMetricsReporter:
)
self.stats.num_grammar_queue_reqs = len(self.scheduler.grammar_manager)
self.stats.cache_hit_rate = cache_hit_rate
# Refresh here too: prefill-heavy stretches can run long between
# decode-stats ticks, and the gauge must decay rather than hold.
self.stats.gen_throughput = self._current_gen_throughput(now)
# Memory pool usage ratios / Absolute token counts
pool_stats.update_scheduler_stats(self.stats)
@@ -875,8 +893,6 @@ class SchedulerMetricsReporter:
spec_num_steps = spec_snapshot["num_steps"]
spec_num_draft_tokens = spec_snapshot["num_draft_tokens"]
cache_hit_rate = 0.0
if self.scheduler.disaggregation_mode == DisaggregationMode.DECODE:
msg += f"pre-allocated usage: {self.scheduler.disagg_decode_prealloc_queue.num_tokens_pre_allocated / self.scheduler.max_total_num_tokens:.2f}, "
msg += f"#prealloc-req: {len(self.scheduler.disagg_decode_prealloc_queue.queue)}, "
@@ -934,7 +950,9 @@ class SchedulerMetricsReporter:
)
self.stats.num_grammar_queue_reqs = len(self.scheduler.grammar_manager)
self.stats.gen_throughput = self.last_gen_throughput
self.stats.cache_hit_rate = cache_hit_rate
# cache_hit_rate is prefill-owned (per-report semantics); decode
# ticks must not reset it, or the exported gauge reads 0 whenever
# a decode report lands between prefill reports.
self.stats.decode_sum_seq_lens = _decode_total_seq_lens(batch)
# Memory pool usage ratios / Absolute token counts
@@ -439,6 +439,21 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
"""
raise NotImplementedError()
def finish_storage_prefetch_admission(
self, req_id: str, fulfilled_tokens: int, reason: Optional[str]
) -> None:
"""Resolve storage-hit accounting once a request is admitted.
Non-storage caches have no lifecycle state to resolve.
"""
def discard_storage_prefetch_accounting(self, req_id: str) -> None:
"""Forget storage-hit lifecycle state without emitting a result."""
def pop_prefetch_loaded_span(self, req_id: str) -> tuple[int, Optional[int]]:
"""Pop L3-loaded tokens and their absolute prefix start, if known."""
return self.pop_prefetch_loaded_tokens(req_id), None
def ready_to_load_host_cache(self) -> Any:
"""
Notify the cache controller to start the KV cache loading
@@ -309,7 +309,7 @@ class BufferModePipeline:
def enqueue_backup_intent(self, node_id: NodeId) -> None:
"""Snapshot a backup intent and commit it to the write queue.
Admission gates: belief skip, parent-cover, backlog cap, oversize.
Drops are silent; the node re-triggers on a later hit."""
Rejected intents are counted; the node re-triggers on a later hit."""
if not self._cache.enable_storage:
return
if node_id in self.inflight_backup_node_ids:
@@ -438,7 +438,7 @@ class BufferModePipeline:
) -> Optional[BufferBackupState]:
# Arena-lookup failure = deleted, key-length mismatch vs the snapshot
# = split, a None FULL device value = evicted. Stale
# intents drop silently; the node re-triggers on a later hit.
# intents are counted as dropped; the node re-triggers on a later hit.
snapshot = intent.snapshot
return self._cache.tree_core.validate_buffer_backup(
snapshot.node_id, len(snapshot.key)
@@ -453,16 +453,20 @@ class BufferModePipeline:
page_size = self._cache.page_size
survivors: deque[_UnifiedBackupIntent] = deque()
states: dict[NodeId, BufferBackupState] = {}
swept_tokens = 0
for intent in self.pending_write_queue:
snapshot = intent.snapshot
state = self._validate_backup_intent(intent)
if state is None:
self.inflight_backup_node_ids.discard(snapshot.node_id)
self.write_backlog_tokens_ -= len(snapshot.hash_values) * page_size
intent_tokens = len(snapshot.hash_values) * page_size
self.write_backlog_tokens_ -= intent_tokens
swept_tokens += intent_tokens
continue
survivors.append(intent)
states[snapshot.node_id] = state
self.pending_write_queue = survivors
self._log_backup_dropped(swept_tokens)
return states
def flush_pending_writes(self) -> None:
@@ -829,12 +833,14 @@ class BufferModePipeline:
if num_tokens == 0 or prefix_tokens is None:
# Nothing usable fetched: recompute.
cache.discard_storage_prefetch_accounting(req_id)
self.release_anchor_lock(req_id)
cc.append_host_mem_release(
host_indices[:num_tokens], extra_pools=aux_xfers or None
)
cc.prefetch_tokens_occupied -= self._occupied_span(host_indices)
cache.prefetch_loaded_tokens_by_reqid[req_id] = 0
cache.prefetch_loaded_storage_start_by_reqid.pop(req_id, None)
return True
staged_pages = num_tokens // cache.page_size
@@ -860,6 +866,7 @@ class BufferModePipeline:
operation_id=operation.id,
)
cache.prefetch_loaded_tokens_by_reqid[req_id] = num_tokens
cache.prefetch_loaded_storage_start_by_reqid[req_id] = operation.storage_start
return True
def plan_staged_splice(
@@ -874,6 +881,7 @@ class BufferModePipeline:
return 0, 0
splice_tokens = staged_splice_tokens(f, device_prefix_len)
if splice_tokens == 0:
covered_tokens = self._resolve_staged_device_coverage(f, device_prefix_len)
logger.info(
"HiCache staged prefetch released req=%s matched=%d "
"device_prefix=%d tokens=%d",
@@ -882,10 +890,18 @@ class BufferModePipeline:
device_prefix_len,
f.num_tokens,
)
self.release_staged_hold(req_id)
reason = None if covered_tokens == f.num_tokens else "shrunk"
self.release_staged_hold(req_id, reason=reason)
return 0, 0
return splice_tokens, self.staged_prefetch_swa_tokens(req_id)
def _resolve_staged_device_coverage(
self, f: _StagedPrefetch, device_prefix_len: int
) -> int:
covered_tokens = min(max(device_prefix_len - f.matched_len, 0), f.num_tokens)
self._cache._resolve_storage_prefetch_tokens(f.req_id, covered_tokens)
return covered_tokens
def staged_prefetch_swa_tokens(self, req_id: str) -> int:
"""SWA device tokens consuming this staged prefetch will allocate (the
staged trailing window); surfaced as the request's swa_host_hit_length
@@ -921,13 +937,17 @@ class BufferModePipeline:
return unchanged
cc = cache.cache_controller
def _drop() -> tuple[torch.Tensor, NodeId]:
def _drop(reason: Optional[str]) -> tuple[torch.Tensor, NodeId]:
cache._finish_storage_prefetch(req.rid, fulfilled_tokens=0, reason=reason)
self.release_anchor_lock(req.rid)
self._free_staging_now(f.host_indices, f.aux_xfers)
cc.prefetch_tokens_occupied -= f.occupied_tokens
# Nothing spliced: keep the surfaced host-hit fields truthful.
req.host_hit_length = 0
req.swa_host_hit_length = 0
req.storage_hit_length = 0
req.storage_hit_start = None
req.host_hit_is_storage = False
return unchanged
# A hold staged under a different namespace than the consuming request
@@ -941,11 +961,12 @@ class BufferModePipeline:
(f.extra_key, f.cache_salt),
(req.extra_key, req.cache_salt),
)
return _drop()
return _drop("dropped")
splice_base = len(req.prefix_indices)
splice_tokens = staged_splice_tokens(f, splice_base)
if splice_tokens == 0:
covered_tokens = self._resolve_staged_device_coverage(f, splice_base)
logger.warning(
"HiCache staged prefetch dropped req=%s matched=%d now=%d "
"tokens_wasted=%d locked=%s",
@@ -955,12 +976,14 @@ class BufferModePipeline:
f.num_tokens,
req.rid in self.anchor_locks,
)
return _drop()
reason = None if covered_tokens == f.num_tokens else "shrunk"
return _drop(reason)
trim_tokens = splice_base - f.matched_len
assert trim_tokens % cache.page_size == 0, (
f"staged splice trim not page-aligned req={req.rid}: "
f"matched={f.matched_len} splice_base={splice_base}"
)
cache._resolve_storage_prefetch_tokens(req.rid, trim_tokens)
key = RadixKey(
array("q", f.key_tokens),
@@ -990,7 +1013,14 @@ class BufferModePipeline:
f.num_tokens,
req.rid in self.anchor_locks,
)
return _drop()
available_end = min(
span_end,
len(live.device_indices),
live.full_kv_hit_length,
)
available_overlap = max(0, available_end - splice_base)
cache._resolve_storage_prefetch_tokens(req.rid, available_overlap)
return _drop(None if available_overlap == splice_tokens else "shrunk")
# Evict-before-alloc (mirrors _load_back_transfers): the budget gate
# counts evictable pages, but cc.load draws from free slots only.
@@ -1007,7 +1037,7 @@ class BufferModePipeline:
avail = cache.token_to_kv_pool_allocator.available_size()
if avail < splice_tokens:
# Genuinely no room (locked pages): recompute.
return _drop()
return _drop("device_capacity")
load_back_id = -(f.operation_id) - 1
device_indices = cc.load(
@@ -1018,7 +1048,7 @@ class BufferModePipeline:
if device_indices is None:
# Transient allocator shortfall despite the evict: recompute
# (init_load_back's degrade contract).
return _drop()
return _drop("device_capacity")
swa_dev = next(
(
@@ -1102,11 +1132,12 @@ class BufferModePipeline:
cc.prefetch_tokens_occupied,
self.anchor_locked_tokens_,
)
if cache.enable_storage_metrics and cache.storage_metrics_collector is not None:
cache.storage_metrics_collector.log_prefetched_tokens(f.num_tokens)
cache._finish_storage_prefetch(
f.req_id, fulfilled_tokens=f.num_tokens, reason=None
)
return True
def release_staged_hold(self, rid: str) -> bool:
def release_staged_hold(self, rid: str, reason: Optional[str] = None) -> bool:
"""Free a staged hold outright — anchor pin, host bounce (KV + aux),
occupancy grant; nothing device-side exists yet. Called for aborts
and for holds that can no longer splice. Returns True when a hold
@@ -1115,6 +1146,7 @@ class BufferModePipeline:
staged = self.staged_prefetches.pop(rid, None)
if staged is None:
return False
self._cache._finish_storage_prefetch(rid, fulfilled_tokens=0, reason=reason)
self._free_staging_now(staged.host_indices, staged.aux_xfers)
self._cache.cache_controller.prefetch_tokens_occupied -= staged.occupied_tokens
return True
@@ -367,6 +367,7 @@ class StorageAttachment:
for req_id in list(cache.ongoing_prefetch):
info = cache.ongoing_prefetch[req_id]
try:
cache.discard_storage_prefetch_accounting(req_id)
if info.host_indices is None:
# Host pages were never allocated for this operation.
cache.revoke_pending_prefetch(req_id)
@@ -394,4 +395,7 @@ class StorageAttachment:
except Exception:
logger.exception("Failed to release host lock for backup op %s", ack_id)
for req_id in list(cache._storage_prefetch_hit_remaining_by_reqid):
cache.discard_storage_prefetch_accounting(req_id)
cache.prefetch_loaded_tokens_by_reqid.clear()
cache.prefetch_loaded_storage_start_by_reqid.clear()
@@ -445,6 +445,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
# The single in-flight resumable insert, if suspended at a barrier.
self._ongoing_insert_walk_state: Optional[_InsertWalkState] = None
self._tracked_unbacked_tokens = 0
self._is_tracking_unbacked_tokens = False
self.root_node = self._new_node()
self.root_node.priority = -sys.maxsize
@@ -1317,6 +1319,18 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
"""Begin a component's device-eviction walk for up to request_cnt tokens."""
self.components_by_type[component_type].evict_device_start(request_cnt)
def _begin_tracking_unbacked_tokens(self) -> None:
assert not self._is_tracking_unbacked_tokens
assert self._tracked_unbacked_tokens == 0
self._is_tracking_unbacked_tokens = True
def _finish_tracking_unbacked_tokens(self) -> int:
assert self._is_tracking_unbacked_tokens
self._is_tracking_unbacked_tokens = False
result = self._tracked_unbacked_tokens
self._tracked_unbacked_tokens = 0
return result
def evict_device_next_node(
self, component_type: ComponentType, tracker: dict[ComponentType, int]
) -> EvictDeviceNextNodeResult:
@@ -1325,9 +1339,15 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
# The walk reads running totals for its doneness check; the result
# carries only this step's delta.
updated_tracker = defaultdict(int, tracker)
result.node_id = self.components_by_type[component_type].evict_device_next_node(
updated_tracker, result.device_frees, result.host_frees
)
self._begin_tracking_unbacked_tokens()
try:
result.node_id = self.components_by_type[
component_type
].evict_device_next_node(
updated_tracker, result.device_frees, result.host_frees
)
finally:
result.unbacked_tokens = self._finish_tracking_unbacked_tokens()
for ct, n in updated_tracker.items():
delta = n - tracker.get(ct, 0)
if delta:
@@ -1348,25 +1368,31 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
result = EvictDeviceLeafResult()
node = self.node_by_id(node_id)
assert self._is_device_leaf(node), f"node {node.id} is not a D-leaf"
if not node.backuped:
if is_write_back:
result.backup_kv = self._build_backup_kv_action(node, write_back=True)
self._begin_tracking_unbacked_tokens()
try:
if not node.backuped:
if is_write_back:
result.backup_kv = self._build_backup_kv_action(
node, write_back=True
)
return result
# Write-through: node has no backup, delete entirely.
self._delete_unbacked_device_leaf(
node,
result.tracker,
device_frees=result.device_frees,
host_frees=result.host_frees,
)
return result
# Write-through: node has no backup, delete entirely.
self._delete_unbacked_device_leaf(
self._demote(
node,
result.tracker,
device_frees=result.device_frees,
host_frees=result.host_frees,
)
return result
self._demote(
node,
result.tracker,
device_frees=result.device_frees,
host_frees=result.host_frees,
)
return result
finally:
result.unbacked_tokens = self._finish_tracking_unbacked_tokens()
def drop_subtree_no_host(self, node_id: NodeId) -> DropSubtreeNoHostResult:
"""Write-back fallback when a D-leaf's D->H backup fails under host
@@ -1727,6 +1753,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
target: EvictLayer = EvictLayer.DEVICE,
tracker: Optional[dict[ComponentType, int]] = None,
) -> tuple[int, int]:
had_host_copy = node.component_data[comp.component_type].host_value is not None
device_freed, host_freed = comp.evict_component(
node, target=target, device_frees=device_frees, host_frees=host_frees
)
@@ -1735,6 +1762,13 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
tracker[comp.component_type] += device_freed
elif EvictLayer.HOST in target:
tracker[comp.component_type] += host_freed
if (
self._is_tracking_unbacked_tokens
and comp.component_type == BASE_COMPONENT_TYPE
and device_freed > 0
and not had_host_copy
):
self._tracked_unbacked_tokens += device_freed
# Detach from the appropriate LRU list(s)
ct = comp.component_type
@@ -50,10 +50,12 @@ class EvictDeviceNextNodeResult(BaseEvictionResult):
node_id: Optional[NodeId] = None
made_progress: bool = False
unbacked_tokens: int = 0
class EvictDeviceLeafResult(BaseEvictionResult):
backup_kv: Optional[BackupKV] = None
unbacked_tokens: int = 0
class DemoteResult(BaseEvictionResult):
@@ -108,15 +108,6 @@ from sglang.srt.utils.rank_consensus_checker import rank_consensus
T = TypeVar("T")
# Metric label per component, matching the host pool names used by
# hicache_backup_tokens_total and the host occupancy gauges.
_COMPONENT_POOL_LABEL = {
ComponentType.FULL: PoolName.KV.value,
ComponentType.SWA: PoolName.SWA.value,
ComponentType.MAMBA: PoolName.MAMBA.value,
}
COMPONENT_REGISTRY: dict[ComponentType, type[TreeComponent]] = {
ComponentType.FULL: FullComponent,
ComponentType.MAMBA: MambaComponent,
@@ -271,9 +262,6 @@ class UnifiedRadixCache(BasePrefixCache):
"l3_demand_requests": 0,
"l3_miss_tokens": 0,
"l1l2_miss_tokens": 0,
"l3_demand_total_tokens": 0,
"l3_sum_rate_all": 0.0,
"l3_sum_rate_main_weighted": 0.0,
}
self.reset()
@@ -368,7 +356,12 @@ class UnifiedRadixCache(BasePrefixCache):
self.ongoing_load_back: dict[int, _OngoingLoadBack] = {}
self.enable_storage = False
self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {}
self.prefetch_loaded_storage_start_by_reqid: dict[str, int] = {}
self.ongoing_prefetch: dict[str, _OngoingPrefetch] = {}
# Rank-agreed L3-hit tokens not yet resolved as usable or unfulfilled.
# Cache-mode entries survive L3->L2 until H2D succeeds or admission
# fails; buffer-mode entries survive staging until the H2D ack.
self._storage_prefetch_hit_remaining_by_reqid: dict[str, int] = {}
# Rids whose storage prefetch resolved without a usable result;
# popped by the scheduler to pace availability-check retries.
self._storage_prefetch_missed_rids: set[str] = set()
@@ -474,13 +467,16 @@ class UnifiedRadixCache(BasePrefixCache):
self.cache_controller is not None
and self.cache_controller.write_policy == "write_back"
)
# Pre-seed the dropped-tokens series at 0 per pool
# Pre-seed the logical dropped-tokens series.
if self.metrics_collector is not None and self.cache_controller is not None:
for ct in self.tree_components:
reasons = ["host_pressure"]
if self._tracks_write_through_unbacked_evictions():
reasons.append("write_through_unbacked_eviction")
for reason in reasons:
self.metrics_collector.increment_dropped_tokens(
num_tokens=0,
reason="host_pressure",
pool=_COMPONENT_POOL_LABEL[ct],
reason=reason,
pool=PoolName.KV.value,
)
self.load_back_threshold = 10
self.prefetch_stop_policy = get_memory().hicache_storage_prefetch_policy
@@ -668,6 +664,11 @@ class UnifiedRadixCache(BasePrefixCache):
"""Advance the eviction walk one node, consuming its step result."""
result = self.tree_core.evict_device_next_node(component_type, tracker)
self._free_values(result.device_frees, result.host_frees)
if self._tracks_write_through_unbacked_evictions():
self._record_dropped_tokens(
result.unbacked_tokens,
reason="write_through_unbacked_eviction",
)
self._accumulate_tracker(tracker, result.tracker)
return result.node_id, result.made_progress
@@ -678,6 +679,11 @@ class UnifiedRadixCache(BasePrefixCache):
deferred write-back BackupKV when one must run before the demote."""
result = self.tree_core.evict_device_leaf(node_id, self.is_write_back)
self._free_values(result.device_frees, result.host_frees)
if self._tracks_write_through_unbacked_evictions():
self._record_dropped_tokens(
result.unbacked_tokens,
reason="write_through_unbacked_eviction",
)
self._accumulate_tracker(tracker, result.tracker)
return result.backup_kv
@@ -693,6 +699,10 @@ class UnifiedRadixCache(BasePrefixCache):
"""Run the write-back drop fallback, consuming its step result."""
result = self.tree_core.drop_subtree_no_host(node_id)
self._free_values(result.device_frees, result.host_frees)
if result.is_dropped:
self._record_dropped_tokens(
result.tracker.get(BASE_COMPONENT_TYPE, 0), reason="host_pressure"
)
self._accumulate_tracker(tracker, result.tracker)
return result.is_dropped
@@ -739,12 +749,10 @@ class UnifiedRadixCache(BasePrefixCache):
written = self._execute_and_commit_kv_backup(
backup_kv, write_back=True
)
freed_before_drop = dict(tracker)
if written > 0:
self.writing_check(write_back=True)
self._demote(node_id, tracker)
elif self._drop_subtree_no_host(node_id, tracker):
self._record_dropped_tokens(tracker, freed_before_drop)
logger.warning(
"write_back: KV subtree dropped without backup "
"due to host memory pressure, root node %d",
@@ -761,22 +769,23 @@ class UnifiedRadixCache(BasePrefixCache):
finally:
self.tree_core.evict_device_end(ct)
def _record_dropped_tokens(
self,
tracker: dict[ComponentType, int],
freed_before_drop: dict[ComponentType, int],
) -> None:
"""Record per-pool tokens dropped without backup under host pressure."""
if self.metrics_collector is None:
def _tracks_write_through_unbacked_evictions(self) -> bool:
return (
isinstance(self.tree_core, UnifiedTreeCore)
and self.host_memory_mode == "cache"
and self.cache_controller is not None
and self.cache_controller.write_policy == "write_through"
)
def _record_dropped_tokens(self, dropped_tokens: int, reason: str) -> None:
"""Record logical KV tokens irreversibly dropped without a host backup."""
if self.metrics_collector is None or dropped_tokens <= 0:
return
for ct, freed in tracker.items():
dropped = freed - freed_before_drop[ct]
if dropped > 0:
self.metrics_collector.increment_dropped_tokens(
num_tokens=dropped,
reason="host_pressure",
pool=_COMPONENT_POOL_LABEL[ct],
)
self.metrics_collector.increment_dropped_tokens(
num_tokens=dropped_tokens,
reason=reason,
pool=PoolName.KV.value,
)
def inc_lock_ref(
self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = ()
@@ -1803,12 +1812,10 @@ class UnifiedRadixCache(BasePrefixCache):
extra_pools=aux_xfers or None,
)
stats["issued"] += 1
# Snapshots for the L3 miss accounting at the query outcome (the
# hit/revoke drains): requested span and total prompt length.
# Snapshot the requested span for L3 miss-token accounting at the
# rank-synchronized query outcome.
operation.stats_requested_tokens = prefetch_length
operation.stats_total_tokens = prefetch_length + len(
matched_prefix_tokens or []
)
operation.storage_start = len(matched_prefix_tokens or [])
self.ongoing_prefetch[req_id] = _OngoingPrefetch(
last_host_node_id,
prefetch_key,
@@ -1923,6 +1930,20 @@ class UnifiedRadixCache(BasePrefixCache):
# Hybrid all-or-nothing check failed; result already discarded.
return
allocated_tokens = len(host_indices)
if completed_tokens < allocated_tokens:
self._resolve_storage_prefetch_tokens(
req_id,
allocated_tokens - completed_tokens,
reason="storage_transfer",
)
if (
completed_tokens > 0
and self.enable_storage_metrics
and self.storage_metrics_collector is not None
):
self.storage_metrics_collector.log_prefetched_tokens(completed_tokens)
if self.buffer_pipeline is not None:
# No graft: release the rank-local tail beyond the synced usable
# length, then park the bounce for admission-time consumption.
@@ -1942,6 +1963,8 @@ class UnifiedRadixCache(BasePrefixCache):
self._apply_cache_actions(insert_result.cache_actions)
if insert_result.host_insert_dropped:
self._resolve_storage_prefetch_tokens(req_id, insert_result.prefix_len)
self._finish_storage_prefetch(req_id, fulfilled_tokens=0, reason="dropped")
self.cache_controller.append_host_mem_release(
host_indices=host_indices[:completed_tokens],
extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
@@ -1965,12 +1988,22 @@ class UnifiedRadixCache(BasePrefixCache):
host_indices[: insert_result.prefix_len]
)
loaded_from_storage = completed_tokens - insert_result.prefix_len
# Cache mode has only completed L3 -> L2 here. Keep the usable
# storage span unresolved until admission proves that L2 -> L1
# load-back actually materialized it for this request.
self._resolve_storage_prefetch_tokens(req_id, insert_result.prefix_len)
self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
del self.ongoing_prefetch[req_id]
self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key)
self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage
if loaded_from_storage > 0:
self.prefetch_loaded_storage_start_by_reqid[req_id] = (
operation.storage_start + insert_result.prefix_len
)
else:
self.prefetch_loaded_storage_start_by_reqid.pop(req_id, None)
logger.info(
"HiCache prefetch %s req=%s completed=%d matched=%d loaded=%d occupied=%d",
"dropped" if insert_result.host_insert_dropped else "success",
@@ -1980,8 +2013,6 @@ class UnifiedRadixCache(BasePrefixCache):
loaded_from_storage,
self.cache_controller.prefetch_tokens_occupied,
)
if self.enable_storage_metrics and self.storage_metrics_collector is not None:
self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage)
return
def _check_hybrid_prefetch_result(
@@ -2055,6 +2086,9 @@ class UnifiedRadixCache(BasePrefixCache):
host_indices=host_indices[:completed_tokens],
extra_pools=pool_transfers if operation.pool_transfers_done else None,
)
self._finish_storage_prefetch(
req_id, fulfilled_tokens=0, reason="storage_transfer"
)
if anchor_lock_params is not None:
self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
if self.buffer_pipeline is not None:
@@ -2065,6 +2099,7 @@ class UnifiedRadixCache(BasePrefixCache):
self._prefetch_occupied_span(prefetch_key, host_indices)
)
self.prefetch_loaded_tokens_by_reqid[req_id] = 0
self.prefetch_loaded_storage_start_by_reqid.pop(req_id, None)
logger.warning(
"HiCache hybrid prefetch discarded req=%s completed=%d requested=%d "
"kv_beliefs_kept_pages=%d",
@@ -2076,11 +2111,108 @@ class UnifiedRadixCache(BasePrefixCache):
return False
return True
def _record_storage_prefetch_hit(self, req_id: str, num_tokens: int) -> None:
"""Start accounting one rank-agreed positive L3 query result."""
if (
num_tokens <= 0
or not self.enable_storage_metrics
or self.storage_metrics_collector is None
):
return
if req_id in self._storage_prefetch_hit_remaining_by_reqid:
logger.warning(
"Replacing unresolved storage-hit accounting req=%s old=%d new=%d",
req_id,
self._storage_prefetch_hit_remaining_by_reqid[req_id],
num_tokens,
)
self.discard_storage_prefetch_accounting(req_id)
self._storage_prefetch_hit_remaining_by_reqid[req_id] = num_tokens
self.storage_metrics_collector.log_storage_prefetch_hit_tokens(num_tokens)
def _resolve_storage_prefetch_tokens(
self, req_id: str, num_tokens: int, reason: Optional[str] = None
) -> None:
if num_tokens <= 0:
return
remaining = self._storage_prefetch_hit_remaining_by_reqid.get(req_id)
if remaining is None:
return
dropped = min(num_tokens, remaining)
if num_tokens > remaining:
logger.warning(
"Storage-prefetch accounting exceeded remaining "
"tokens req=%s requested=%d remaining=%d reason=%s",
req_id,
num_tokens,
remaining,
reason,
)
if reason is not None:
self.storage_metrics_collector.log_storage_prefetch_unfulfilled_tokens(
dropped, reason
)
remaining -= dropped
if remaining:
self._storage_prefetch_hit_remaining_by_reqid[req_id] = remaining
else:
self._storage_prefetch_hit_remaining_by_reqid.pop(req_id, None)
def _finish_storage_prefetch(
self, req_id: str, fulfilled_tokens: int, reason: Optional[str]
) -> None:
remaining = self._storage_prefetch_hit_remaining_by_reqid.pop(req_id, None)
if remaining is None:
return
if fulfilled_tokens > remaining:
logger.warning(
"Storage-prefetch fulfilled accounting exceeded remaining "
"tokens req=%s fulfilled=%d remaining=%d",
req_id,
fulfilled_tokens,
remaining,
)
unfulfilled = max(0, remaining - fulfilled_tokens)
if reason is not None:
self.storage_metrics_collector.log_storage_prefetch_unfulfilled_tokens(
unfulfilled, reason
)
def finish_storage_prefetch_admission(
self, req_id: str, fulfilled_tokens: int, reason: Optional[str]
) -> None:
"""Resolve a cache-mode L3 hit after request admission.
Buffer mode resolves at the H2D completion ack instead, because its
host allocation is transport staging rather than a resident L2 hit.
"""
if self.host_memory_mode == "cache":
self._finish_storage_prefetch(req_id, fulfilled_tokens, reason)
def discard_storage_prefetch_accounting(self, req_id: str) -> None:
"""Drop lifecycle state for cases intentionally excluded from metrics."""
self._storage_prefetch_hit_remaining_by_reqid.pop(req_id, None)
def _handle_storage_prefetch_anchor_loss(self, req_id: str) -> None:
self._finish_storage_prefetch(req_id, fulfilled_tokens=0, reason="shrunk")
# The span is still L3-resident; retry from the shorter live match.
self._storage_prefetch_missed_rids.add(req_id)
self.revoke_pending_prefetch(req_id)
def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
# The request is being scheduled; a still-unserved miss marker is moot.
self._storage_prefetch_missed_rids.discard(req_id)
self.prefetch_loaded_storage_start_by_reqid.pop(req_id, None)
return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)
def pop_prefetch_loaded_span(self, req_id: str) -> tuple[int, Optional[int]]:
"""Pop the loaded L3 token count and its absolute prefix start."""
self._storage_prefetch_missed_rids.discard(req_id)
return (
self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0),
self.prefetch_loaded_storage_start_by_reqid.pop(req_id, None),
)
def pop_storage_prefetch_miss(self, req_id: str) -> bool:
"""True once per resolved storage-prefetch miss for a live request;
the scheduler uses it to arm the paced availability-check retry."""
@@ -2110,12 +2242,14 @@ class UnifiedRadixCache(BasePrefixCache):
if self.linker is not None:
self.linker.release_request(rid)
self.prefetch_loaded_tokens_by_reqid.pop(rid, None)
self.prefetch_loaded_storage_start_by_reqid.pop(rid, None)
self._storage_prefetch_missed_rids.discard(rid)
if (
self.buffer_pipeline is not None
and self.buffer_pipeline.release_staged_hold(rid)
):
return
self.discard_storage_prefetch_accounting(rid)
if rid not in self.ongoing_prefetch:
return
@@ -2177,23 +2311,12 @@ class UnifiedRadixCache(BasePrefixCache):
else:
stats["revoked_full_miss"] += 1
miss = requested - hit
total = max(operation.stats_total_tokens, requested, 1)
stats["l3_demand_requests"] += 1
stats["l1l2_miss_tokens"] += requested
stats["l3_miss_tokens"] += miss
stats["l3_demand_total_tokens"] += total
stats["l3_sum_rate_all"] += miss / total
stats["l3_sum_rate_main_weighted"] += (miss / requested) * total
def prefetch_outcome_stats_snapshot(self) -> dict:
"""Cumulative counters + instantaneous occupancy, in the schema
log_prefetch_stats consumers expect."""
cc = self.cache_controller
cap = max(cc.prefetch_capacity_limit, 1)
return {
**self._prefetch_outcome_stats,
"occupancy_ratio": cc.prefetch_tokens_occupied / cap,
}
return self._prefetch_outcome_stats.copy()
def _prefetch_occupied_span(self, prefetch_key, host_indices) -> int:
"""Occupancy units held by a prefetch: cache mode reserves the
@@ -2205,6 +2328,7 @@ class UnifiedRadixCache(BasePrefixCache):
def revoke_pending_prefetch(self, req_id: str) -> None:
info = self.ongoing_prefetch.pop(req_id, None)
self._finish_storage_prefetch(req_id, fulfilled_tokens=0, reason="dropped")
if info is None:
return
(
@@ -2274,6 +2398,7 @@ class UnifiedRadixCache(BasePrefixCache):
(buffer mode parks and retries; cache mode revokes)."""
req_id = operation.request_id
info = self.ongoing_prefetch.get(req_id)
hit_tokens = operation.storage_hit_count
if info is None:
return True # aborted/cleaned; nothing to retry
if operation.is_terminated():
@@ -2291,10 +2416,7 @@ class UnifiedRadixCache(BasePrefixCache):
# splice base is gone is not worth its storage read.
if self.buffer_pipeline.try_lock_anchor(req_id) == "anchor_lost":
self._prefetch_outcome_stats["declined_anchor_lost"] += 1
# Span still L3-resident: arm the paced retry to re-fetch
# from the shorter post-loss match.
self._storage_prefetch_missed_rids.add(req_id)
self.revoke_pending_prefetch(req_id)
self._handle_storage_prefetch_anchor_loss(req_id)
return True
if self.buffer_pipeline.staged_span_covered(
req_id, operation.storage_hit_count
@@ -2302,9 +2424,12 @@ class UnifiedRadixCache(BasePrefixCache):
# Live tree already covers the span: nothing left to
# splice, so skip the storage read.
self._prefetch_outcome_stats["declined_device_covered"] += 1
self._finish_storage_prefetch(
req_id, fulfilled_tokens=0, reason=None
)
self.revoke_pending_prefetch(req_id)
return True
alloc_len = operation.storage_hit_count
alloc_len = hit_tokens
host_indices = cc.mem_pool_host.alloc(alloc_len)
if host_indices is None:
self.evict_host(alloc_len)
@@ -2314,7 +2439,7 @@ class UnifiedRadixCache(BasePrefixCache):
# (Cache mode only — buffer mode parks for the full hit.)
available_size = cc.mem_pool_host.available_size()
alloc_len = min(
operation.storage_hit_count,
hit_tokens,
available_size - (available_size % self.page_size),
)
if alloc_len >= self.prefetch_threshold:
@@ -2322,9 +2447,15 @@ class UnifiedRadixCache(BasePrefixCache):
if host_indices is None:
if buffer_mode:
return False
self._finish_storage_prefetch(
req_id, fulfilled_tokens=0, reason="host_capacity"
)
self.revoke_pending_prefetch(req_id)
return True
self._resolve_storage_prefetch_tokens(
req_id, hit_tokens - alloc_len, reason="host_capacity"
)
operation.storage_hit_count = alloc_len
operation.hash_value = operation.hash_value[: alloc_len // self.page_size]
operation.host_indices = host_indices
@@ -2344,22 +2475,39 @@ class UnifiedRadixCache(BasePrefixCache):
parked.popleft()
for operation in _drain_queue(cc.prefetch_hit_queue, n_storage_hit):
req_id = operation.request_id
hit_tokens = operation.storage_hit_count
info = self.ongoing_prefetch.get(req_id)
if info is None:
# Request already aborted/cleaned up; still flush the
# query's absent-hash feedback.
self._invalidate_absent_from_hit_query(operation)
if hit_tokens > 0:
self.discard_storage_prefetch_accounting(req_id)
continue
if hit_tokens > 0:
self._record_storage_prefetch_hit(req_id, hit_tokens)
if operation.is_terminated():
# Controller-side miss termination (retryable) or an abort
# race (abort cleanup discards the marker).
if hit_tokens > 0:
self._finish_storage_prefetch(
req_id,
fulfilled_tokens=0,
reason=(
"below_threshold"
if hit_tokens < self.prefetch_threshold
else None
),
)
self._storage_prefetch_missed_rids.add(req_id)
self.revoke_pending_prefetch(req_id)
continue
if operation.storage_hit_count < self.prefetch_threshold:
# Below-threshold hit: classify + feed the L3 miss
# accounting, then revoke (not enough benefit).
if hit_tokens < self.prefetch_threshold:
# Below-threshold hits are not worth the transfer.
self._account_prefetch_outcome(operation, revoked=True)
self._finish_storage_prefetch(
req_id, fulfilled_tokens=0, reason="below_threshold"
)
self._storage_prefetch_missed_rids.add(req_id)
self.revoke_pending_prefetch(req_id)
continue
@@ -2858,10 +3006,9 @@ class UnifiedRadixCache(BasePrefixCache):
if self.enable_storage_metrics and self.storage_metrics_collector is not None:
storage_metrics = self.cache_controller.storage_backend.get_stats()
if storage_metrics is None:
# Backends without native stats (e.g. file) still carry the
# controller-side prefetch outcome counters.
storage_metrics = StorageMetrics()
storage_metrics.prefetch_stats = self.prefetch_outcome_stats_snapshot()
if not hasattr(storage_metrics, "prefetch_stats"):
storage_metrics.prefetch_stats = self.prefetch_outcome_stats_snapshot()
self.storage_metrics_collector.log_storage_metrics(storage_metrics)
def ready_to_load_host_cache(self) -> int:
@@ -1780,10 +1780,7 @@ class TokenizerMetricsCollector(_StatLoggerDIMixin):
if "storage" in cached_tokens_details:
storage_tokens = cached_tokens_details.get("storage", 0)
if storage_tokens > 0:
backend = (
cached_tokens_details.get("storage_backend") or "unknown"
)
report_cache_source(f"storage_{backend}", storage_tokens)
report_cache_source("storage", storage_tokens)
else:
# Fallback for backward compatibility
labels_total = {**labels, "cache_source": "total"}
@@ -1869,7 +1866,8 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
self.prefetched_tokens_total = Counter(
name="sglang:prefetched_tokens_total",
documentation="Number of prefetched prompt tokens.",
documentation="Prompt tokens successfully transferred from L3 "
"storage into host memory, before admission-time reuse or discard.",
labelnames=labels.keys(),
)
@@ -1879,10 +1877,38 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
labelnames=labels.keys(),
)
self.storage_prefetch_hit_tokens_total = Counter(
name="sglang:storage_prefetch_hit_tokens_total",
documentation="Storage-hit tokens returned by an L3 query before "
"host allocation, transfer, or cache insertion.",
labelnames=labels.keys(),
)
self.storage_prefetch_unfulfilled_tokens_total = Counter(
name="sglang:storage_prefetch_unfulfilled_tokens_total",
documentation="Storage-hit tokens that did not become a usable "
"prefetch result, by terminal reason.",
labelnames=list(labels.keys()) + ["reason"],
)
for reason in (
"below_threshold",
"host_capacity",
"device_capacity",
"storage_transfer",
"shrunk",
"dropped",
):
self.storage_prefetch_unfulfilled_tokens_total.labels(
**self.labels, reason=reason
)
self.backup_dropped_tokens_total = Counter(
name="sglang:hicache_backup_dropped_tokens_total",
documentation="Buffer-mode backup tokens dropped by write-path rate "
"limiting (backlog cap or dropped-parent cascade).",
documentation="Buffer-mode backup tokens that never reached L3 "
"because admission rejected them or the source became stale "
"before the D2H staging launched; the write path lagging device "
"churn is the common cause, and the span rewrites only on a "
"later recompute.",
labelnames=labels.keys(),
)
@@ -1928,14 +1954,19 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
self.histogram_prefetch_bandwidth = Histogram(
name="sglang:prefetch_bandwidth",
documentation="Histogram of prefetch bandwidth in GB/s.",
documentation="Histogram of per-op L3 prefetch wire bandwidth in "
"GB/s: transferred bytes (compressed size when KV compression is "
"on) over the op's IO wall time.",
labelnames=labels.keys(),
buckets=bucket_bandwidth,
)
self.histogram_backup_bandwidth = Histogram(
name="sglang:backup_bandwidth",
documentation="Histogram of backup bandwidth in GB/s.",
documentation="Histogram of per-op L3 backup wire bandwidth in "
"GB/s: transferred bytes over the op's IO wall time, excluding "
"pages the backend already held (the log line's "
"processing_throughput includes them).",
labelnames=labels.keys(),
buckets=bucket_bandwidth,
)
@@ -1948,6 +1979,18 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
if backuped_tokens > 0:
self.backuped_tokens_total.labels(**self.labels).inc(backuped_tokens)
def log_storage_prefetch_hit_tokens(self, num_tokens: int) -> None:
if num_tokens > 0:
self.storage_prefetch_hit_tokens_total.labels(**self.labels).inc(num_tokens)
def log_storage_prefetch_unfulfilled_tokens(
self, num_tokens: int, reason: str
) -> None:
if num_tokens > 0:
self.storage_prefetch_unfulfilled_tokens_total.labels(
**self.labels, reason=reason
).inc(num_tokens)
def log_backup_dropped_tokens(self, dropped_tokens: int):
if dropped_tokens > 0:
self.backup_dropped_tokens_total.labels(**self.labels).inc(dropped_tokens)
@@ -2010,6 +2053,11 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin):
"SGLANG_BUCKET_EVICTION_DURATION"
)
if bucket_eviction_duration is None:
# Keep within a 1.0s ceiling: downstream metrics gateways with a
# fixed bucket preset silently drop the series when a boundary
# falls outside it. The >1s regime is covered by the
# backup-duration histogram, and the env override widens this
# where no such preset applies.
bucket_eviction_duration = [
0.001,
0.002,
@@ -2029,11 +2077,6 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin):
0.2,
0.5,
1.0,
2.0,
5.0,
10.0,
30.0,
60.0,
]
bucket_load_back_duration = get_histogram_conf_from_env(
"SGLANG_BUCKET_LOAD_BACK_DURATION"
@@ -2062,12 +2105,13 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin):
# D->H backups include blocking merged ops issued during eviction under
# --hicache-write-policy write_back, which can run for seconds -- hence
# the wider default range than load-back.
# Keep the boundaries to a coarse, widely supported set: downstream
# metrics gateways with a fixed bucket preset drop the series when a
# boundary is not in it.
bucket_backup_duration = [
0.001,
0.002,
0.005,
0.01,
0.02,
0.05,
0.1,
0.2,
@@ -2101,7 +2145,9 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin):
self.load_back_duration_seconds = Histogram(
name="sglang:load_back_duration_seconds",
documentation="Time taken to load KV cache from CPU back to GPU in seconds.",
documentation="GPU-stream span of a merged host-to-device load-back "
"copy loop, including gaps between per-layer I/O submissions but "
"excluding pre-transfer queue and fence waits.",
labelnames=labels.keys(),
buckets=bucket_load_back_duration,
)
@@ -2138,8 +2184,8 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin):
documentation="Bytes loaded back from local host DRAM (L2) to "
"GPU, all pools combined, including draft/sidecar transfers that "
"the token counter excludes. Divided by the rate of "
"load_back_duration_seconds_sum, gives the achieved H->D "
"bandwidth while transferring.",
"load_back_duration_seconds_sum, gives the achieved bandwidth "
"over each merged H2D load operation.",
labelnames=labels.keys(),
)
@@ -2154,9 +2200,11 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin):
self.hicache_dropped_tokens = Counter(
name="sglang:hicache_dropped_tokens_total",
documentation="The number of device KV tokens destroyed without a "
"host backup, by pool (kv, swa, ...) and reason (e.g. write-back "
"backup failure under host memory pressure).",
documentation="The number of logical device KV tokens destroyed "
"without a host backup. The pool label is always kv; reason is "
"host_pressure for write-back failure, or "
"write_through_unbacked_eviction when "
"eager write-through did not complete before eviction.",
labelnames=list(labels.keys()) + ["reason", "pool"],
)
@@ -565,6 +565,17 @@ class StreamingSession(BasePrefixCache):
def init_load_back(self, params: InitLoadBackParams):
return self.inner.init_load_back(params)
def pop_prefetch_loaded_span(self, req_id: str) -> tuple[int, Optional[int]]:
return self.inner.pop_prefetch_loaded_span(req_id)
def finish_storage_prefetch_admission(
self, req_id: str, fulfilled_tokens: int, reason: Optional[str]
) -> None:
self.inner.finish_storage_prefetch_admission(req_id, fulfilled_tokens, reason)
def discard_storage_prefetch_accounting(self, req_id: str) -> None:
self.inner.discard_storage_prefetch_accounting(req_id)
def ready_to_load_host_cache(self):
return self.inner.ready_to_load_host_cache()
@@ -101,6 +101,11 @@ class TestPrefillAdder(CustomTestCase):
req.retracted_stain = False
req.host_hit_length = 0
req.storage_hit_length = 0
req.storage_hit_start = None
req.host_hit_is_storage = False
req.host_loaded_length = 0
req.materialized_host_hit_len.return_value = 0
req.fulfilled_storage_hit_len.return_value = 0
req.finished.return_value = False
req.needs_host_load_back.return_value = False
return req
@@ -120,6 +125,51 @@ class TestPrefillAdder(CustomTestCase):
defaults.update(kwargs)
return PrefillAdder(**defaults)
def test_storage_prefetch_fulfillment_resolves_at_admission(self):
adder = self.create_adder(self.create_running_batch())
req = self.create_mock_req("storage-hit", priority=0, max_new_tokens=1)
req.host_hit_length = 12
req.host_loaded_length = 4
req.storage_hit_length = 8
req.storage_hit_start = 4
req.materialized_host_hit_len.return_value = 4
req.fulfilled_storage_hit_len.return_value = 8
req.needs_host_load_back.return_value = True
adder._account_prefill_cache_admission(req, prefix_len=12)
self.mock_tree_cache.finish_storage_prefetch_admission.assert_called_once_with(
"storage-hit",
fulfilled_tokens=8,
reason=None,
)
self.assertEqual(adder.log_device_hit_tokens, 4)
self.assertEqual(adder.log_host_hit_tokens, 0)
self.assertEqual(adder.log_storage_hit_tokens, 8)
self.mock_tree_cache.finish_storage_prefetch_admission.reset_mock()
req.host_loaded_length = 0
req.materialized_host_hit_len.return_value = 0
req.fulfilled_storage_hit_len.return_value = 0
adder._account_prefill_cache_admission(req, prefix_len=0)
self.mock_tree_cache.finish_storage_prefetch_admission.assert_called_once_with(
"storage-hit", fulfilled_tokens=0, reason="device_capacity"
)
def test_retracted_storage_prefetch_accounting_is_omitted(self):
adder = self.create_adder(self.create_running_batch())
req = self.create_mock_req(
"retracted-storage-hit", priority=0, max_new_tokens=1
)
req.retracted_stain = True
adder._account_prefill_cache_admission(req, prefix_len=8)
self.mock_tree_cache.discard_storage_prefetch_accounting.assert_called_once_with(
"retracted-storage-hit"
)
self.mock_tree_cache.finish_storage_prefetch_admission.assert_not_called()
def test_preempt_success_high_priority_values_first(self):
params = [
("run1", 0, 50),
@@ -11,7 +11,10 @@ from sglang.test.test_utils import maybe_stub_sgl_kernel
maybe_stub_sgl_kernel()
from sglang.srt.managers.schedule_batch import ScheduleBatch # noqa: E402
from sglang.srt.managers.schedule_batch import ( # noqa: E402
ScheduleBatch,
split_cached_prefix_by_tier,
)
from sglang.srt.model_executor.forward_batch_info import ForwardMode # noqa: E402
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm # noqa: E402
from sglang.srt.utils.common import Range # noqa: E402
@@ -21,6 +24,65 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
AUTO_FILL_EXCLUDED_FIELDS = ["reqs"]
class TestCachedPrefixTierAttribution(unittest.TestCase):
def test_split_cached_prefix_by_tier(self):
cases = [
(
{"prefix_len": 100, "host_hit_len": 70, "storage_hit_len": 20},
(30, 50, 20),
),
(
{
"prefix_len": 80,
"host_hit_len": 50,
"storage_hit_len": 20,
"storage_hit_start": 80,
},
(30, 50, 0),
),
(
{
"prefix_len": 90,
"host_hit_len": 60,
"storage_hit_len": 20,
"storage_hit_start": 80,
},
(30, 50, 10),
),
(
{
"prefix_len": 100,
"host_hit_len": 70,
"storage_hit_len": 70,
"host_hit_is_storage": True,
},
(30, 0, 70),
),
(
{
"prefix_len": 100,
"host_hit_len": 0,
"storage_hit_len": 20,
"storage_hit_start": 80,
},
(80, 0, 20),
),
(
{
"prefix_len": 100,
"host_hit_len": 0,
"storage_hit_len": 20,
"storage_hit_start": 80,
"host_hit_is_storage": True,
},
(80, 0, 20),
),
]
for kwargs, expected in cases:
with self.subTest(**kwargs):
self.assertEqual(split_cached_prefix_by_tier(**kwargs), expected)
def make_schedule_batch(bs: int, **overrides) -> ScheduleBatch:
batch = ScheduleBatch(reqs=overrides.pop("reqs"))
# init_new always sets a SpeculativeAlgorithm enum, never None.
@@ -9,6 +9,7 @@ import torch
from sglang.srt.managers.cache_controller import CacheOperation, HiCacheController
from sglang.srt.mem_cache import l2_transfer as transfer_module
from sglang.srt.mem_cache.buffer_mode.pipeline import BufferModePipeline
from sglang.srt.mem_cache.hicache_storage import (
PoolHitPolicy,
PoolName,
@@ -262,6 +263,35 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
controller._num_tokens_by_pool.assert_called_once_with(merged_op)
self.assertEqual(controller.ack_load_queue[0].node_ids, [7, 7])
def test_short_staged_swa_tail_resolves_device_covered_head(self):
pipeline = BufferModePipeline.__new__(BufferModePipeline)
pipeline._cache = mock.Mock()
pipeline.release_staged_hold = mock.Mock(return_value=True)
pipeline.staged_prefetches = {
"r": SimpleNamespace(
req_id="r",
key_tokens=list(range(8)),
extra_key=None,
cache_salt=None,
matched_len=2,
num_tokens=6,
occupied_tokens=6,
host_indices=_indices(0, 6),
aux_xfers=[
PoolTransfer(
name=PoolName.SWA,
host_indices=_indices(0, 4),
)
],
hash_values=[],
operation_id=1,
)
}
self.assertEqual(pipeline.plan_staged_splice("r", device_prefix_len=6), (0, 0))
pipeline._cache._resolve_storage_prefetch_tokens.assert_called_once_with("r", 4)
pipeline.release_staged_hold.assert_called_once_with("r", reason="shrunk")
def test_l2_transfer_maps_global_layers(self):
host_pool = mock.Mock()
transfer = L2Transfer(
@@ -47,6 +47,10 @@ class _FakeTreeCore:
}
self.evicted = []
self.cascaded = []
# The real eviction helper reports unbacked FULL evictions to the
# write-through drop counter; this fake never tracks a walk.
self._is_tracking_unbacked_tokens = False
self._tracked_unbacked_tokens = 0
def _evict_component_and_detach_lru(self, node, component, *args, **kwargs):
self.evicted.append(node)
@@ -3676,6 +3676,8 @@ class UnifiedRadixCacheSuite:
self.cfg, enable_kv_cache_events=True
)
self._init_buffer_hicache(cons, storage_dir)
cons.enable_storage_metrics = True
cons.storage_metrics_collector = mock.Mock()
cons.take_events()
avail0 = self._host_avail_sizes(cons)
dev_avail0 = cons.token_to_kv_pool_allocator.available_size()
@@ -3783,8 +3785,14 @@ class UnifiedRadixCacheSuite:
and e.medium == StorageMedium.CPU
]
self.assertEqual(cpu_events, [])
cons.storage_metrics_collector.log_storage_prefetch_hit_tokens.assert_called_once_with(
len(seq)
)
cons.storage_metrics_collector.log_prefetched_tokens.assert_called_once_with(
len(seq)
)
cons.storage_metrics_collector.log_storage_prefetch_unfulfilled_tokens.assert_not_called()
self.assertIn("occupancy_ratio", cons.prefetch_outcome_stats_snapshot())
cons.sanity_check()
def test_buffer_only_cache_salt_uses_the_request_namespace(self):
@@ -4148,6 +4156,8 @@ class UnifiedRadixCacheSuite:
cons, cons_alloc, cons_rtp = build_fixture(self.cfg)
self._init_buffer_hicache(cons, storage_dir)
cons.enable_storage_metrics = True
cons.storage_metrics_collector = mock.Mock()
avail0 = self._host_avail_sizes(cons)
req_id = "sibling-publish"
@@ -4188,6 +4198,7 @@ class UnifiedRadixCacheSuite:
k, v = self._snapshot_full_kv(cons_alloc, m.device_indices)
self.assertTrue(torch.equal(k, sib_kv[0]))
self.assertTrue(torch.equal(v, sib_kv[1]))
cons.storage_metrics_collector.log_storage_prefetch_unfulfilled_tokens.assert_not_called()
cons.sanity_check()
def test_buffer_only_load_back_drops_on_full_overlap_masked_by_swa_tombstone(
@@ -8448,6 +8459,7 @@ class TestPrefetchCommitOrdering(CustomTestCase):
cache.tree_core.insert_host.return_value = insert_result
operation = mock.MagicMock()
operation.request_id = "req"
operation.completed_tokens = 8
cache.ongoing_prefetch = {
operation.request_id: (
7,
@@ -8588,6 +8600,8 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase):
"""The caller owns every completed buffer when host insertion drops."""
cache, allocator, _ = build_fixture(self.cfg)
self._init_hicache(cache)
cache.enable_storage_metrics = True
cache.storage_metrics_collector = mock.Mock()
parent_id = self._insert_device(
cache, allocator, list(range(1, 1 + 3 * self.ps))
@@ -8633,6 +8647,7 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase):
anchor_lock_params,
comp_xfers,
)
cache._record_storage_prefetch_hit(req_id, completed_tokens)
cache.cache_controller.prefetch_tokens_occupied = completed_tokens
hashes = [f"h{i}" for i in range(completed_tokens // self.ps)]
operation.hash_value = hashes
@@ -8681,9 +8696,38 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase):
)
self.assertIs(drop_releases[0].kwargs["extra_pools"][0], swa_transfer)
self.assertIs(drop_releases[0].kwargs["extra_pools"][1], mamba_transfer)
cache.storage_metrics_collector.log_storage_prefetch_hit_tokens.assert_called_once_with(
completed_tokens
)
cache.storage_metrics_collector.log_storage_prefetch_unfulfilled_tokens.assert_called_once_with(
completed_tokens, "dropped"
)
cache.sanity_check()
def test_write_through_eviction_counts_unbacked_tokens(self):
if _selected_tree_core_test_backend() == "rust":
# The unbacked-eviction tracker is a Python tree-core feature;
# UnifiedRadixCache only enables it for that backend.
self.skipTest(
"write-through unbacked-eviction tracking is Python-core only"
)
cache, allocator, _ = build_fixture(self.cfg)
self._init_hicache(cache)
cache.metrics_collector = mock.Mock()
seq = list(range(1, 1 + 2 * self.ps))
self._insert_device(cache, allocator, seq)
result = cache.evict(EvictParams(num_tokens=len(seq)))
self.assertGreaterEqual(result.num_tokens_evicted, len(seq))
cache.metrics_collector.increment_dropped_tokens.assert_called_once_with(
num_tokens=len(seq),
reason="write_through_unbacked_eviction",
pool=PoolName.KV.value,
)
cache.sanity_check()
def test_prefetch_refill_leaves_eviction_path_uncorrupted(self):
"""Write-through: eviction after such a refill must not corrupt the tree."""
cache, allocator, _ = build_fixture(self.cfg)
@@ -8946,6 +8990,20 @@ class TestAnchorLockOutcomePolicy(CustomTestCase):
self.assertEqual(pipeline.anchor_locks, {})
self.assertEqual(pipeline.anchor_locked_tokens_, 0)
def test_positive_hit_with_lost_anchor_is_reported_as_shrunk(self):
cache = UnifiedRadixCache.__new__(UnifiedRadixCache)
cache._storage_prefetch_missed_rids = set()
cache._finish_storage_prefetch = mock.Mock()
cache.revoke_pending_prefetch = mock.Mock()
cache._handle_storage_prefetch_anchor_loss(self._REQ)
cache._finish_storage_prefetch.assert_called_once_with(
self._REQ, fulfilled_tokens=0, reason="shrunk"
)
self.assertIn(self._REQ, cache._storage_prefetch_missed_rids)
cache.revoke_pending_prefetch.assert_called_once_with(self._REQ)
def test_over_cap_reports_cap_skip_before_matching(self):
cache = self._make_cache(live_match_len=len(self._PREFIX))
pipeline = self._make_pipeline(cache, cap_tokens=len(self._PREFIX) - 1)
@@ -42,6 +42,48 @@ from sglang.srt.observability.metrics_collector import (
from sglang.srt.runtime_context import get_context, reset_context
class _BoundRecordingMetric:
def __init__(self, metric, labels):
self.metric = metric
self.labels = labels
def inc(self, value=1):
self.metric.increments.append((self.labels, value))
def observe(self, value):
self.metric.observations.append((self.labels, value))
def set(self, value):
self.metric.sets.append((self.labels, value))
class _RecordingMetric:
"""Small prometheus_client-compatible metric that preserves labels."""
def __init__(self, *args, name=None, labelnames=(), **kwargs):
self.name = name if name is not None else args[0]
self.labelnames = tuple(labelnames)
self.increments = []
self.observations = []
self.sets = []
def labels(self, *values, **labels):
if values:
labels = dict(zip(self.labelnames, values, strict=True))
return _BoundRecordingMetric(self, labels)
class _RecordingTokenizerMetricsCollector(TokenizerMetricsCollector):
_counter_cls = _RecordingMetric
_gauge_cls = _RecordingMetric
_histogram_cls = _RecordingMetric
class _RecordingStorageMetricsCollector(StorageMetricsCollector):
_counter_cls = _RecordingMetric
_histogram_cls = _RecordingMetric
class TestCollectorClassAttrs(unittest.TestCase):
"""All five collectors expose four DI hook class attrs, all defaulting to
None so the existing prometheus_client backend is used unchanged."""
@@ -137,5 +179,50 @@ class TestDefaultBackend(unittest.TestCase):
)
class TestHiCacheMetrics(unittest.TestCase):
def test_cached_tokens_uses_literal_storage_source(self):
labels = {"model_name": "test"}
with get_context().override_server_args(
prompt_tokens_buckets=None, generation_tokens_buckets=None
):
collector = _RecordingTokenizerMetricsCollector(labels=labels)
collector.observe_one_finished_request(
labels=labels,
prompt_tokens=20,
generation_tokens=2,
cached_tokens=12,
e2e_latency=0.1,
has_grammar=False,
cached_tokens_details={
"device": 3,
"host": 4,
"storage": 5,
"storage_backend": "BackendShim",
},
)
by_source = {
metric_labels["cache_source"]: value
for metric_labels, value in collector.cached_tokens_total.increments
}
self.assertEqual(by_source, {"device": 3, "host": 4, "storage": 5})
def test_storage_prefetch_lifecycle_metrics(self):
labels = {"model_name": "test"}
collector = _RecordingStorageMetricsCollector(labels=labels)
collector.log_storage_prefetch_hit_tokens(21)
collector.log_storage_prefetch_unfulfilled_tokens(4, "storage_transfer")
self.assertEqual(
collector.storage_prefetch_hit_tokens_total.increments, [(labels, 21)]
)
self.assertEqual(
collector.storage_prefetch_unfulfilled_tokens_total.increments,
[({**labels, "reason": "storage_transfer"}, 4)],
)
if __name__ == "__main__":
unittest.main()