diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index cceac56e3..d42980ca9 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index d8263b63b..418a01f20 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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 diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 5aa27c38e..d08e86093 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -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() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 128ba41ee..4c0189a21 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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), diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index f57ee42a1..9a7017256 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -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 diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 2998b18e4..37891f6ea 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py index 11bc49860..6de4d0cd2 100644 --- a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py +++ b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py index 8a0d2e8bf..15b6db174 100644 --- a/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py +++ b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py @@ -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() diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 7dc939be4..9ef9bbd33 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index 57bc54b8c..ec91abc9c 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -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): diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 1c603cac3..9f10f0f31 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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: diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 156454338..0a9a6f08d 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -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"], ) diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index da18238ea..d9b271c35 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -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() diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 6525c3ce9..d74db97d6 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -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), diff --git a/test/registered/unit/managers/test_schedule_batch_out_of_place.py b/test/registered/unit/managers/test_schedule_batch_out_of_place.py index c91500ff4..67f551a03 100644 --- a/test/registered/unit/managers/test_schedule_batch_out_of_place.py +++ b/test/registered/unit/managers/test_schedule_batch_out_of_place.py @@ -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. diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index ff6f65593..b58344a39 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py @@ -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( diff --git a/test/registered/unit/mem_cache/test_mamba_path_state_cap.py b/test/registered/unit/mem_cache/test_mamba_path_state_cap.py index ec4f425af..940a2c350 100644 --- a/test/registered/unit/mem_cache/test_mamba_path_state_cap.py +++ b/test/registered/unit/mem_cache/test_mamba_path_state_cap.py @@ -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) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 0ec3d265c..5c55c7684 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -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) diff --git a/test/registered/unit/observability/test_stat_loggers_di.py b/test/registered/unit/observability/test_stat_loggers_di.py index 99a087470..2888f26de 100644 --- a/test/registered/unit/observability/test_stat_loggers_di.py +++ b/test/registered/unit/observability/test_stat_loggers_di.py @@ -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()