[HiCache] L3 storage prefetch lifecycle metrics and cross-tier attribution fixes (#37503)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user