From 106bcc12935569d32a3f2a241f4f17cf76fe8a30 Mon Sep 17 00:00:00 2001 From: Zhiqiang Xie Date: Wed, 5 Aug 2026 14:13:06 -0700 Subject: [PATCH] Observability enhancement for HiCache (#32388) --- .../sglang/srt/managers/cache_controller.py | 32 +++- python/sglang/srt/managers/schedule_batch.py | 42 +++--- python/sglang/srt/managers/schedule_policy.py | 26 +++- .../scheduler_components/metrics_reporter.py | 24 +-- .../srt/mem_cache/hi_mamba_radix_cache.py | 8 +- python/sglang/srt/mem_cache/hiradix_cache.py | 31 +++- .../hybrid_cache/hybrid_cache_controller.py | 55 ++++++- .../srt/mem_cache/unified_radix_cache.py | 64 +++++++- .../srt/observability/metrics_collector.py | 139 +++++++++++++++++- .../unit/managers/test_prefill_adder.py | 2 + .../test_hicache_load_back_timing.py | 6 +- ...test_hicache_staged_write_back_dispatch.py | 40 ++++- 12 files changed, 420 insertions(+), 49 deletions(-) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 0f3234b3b..beb4b3980 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -163,6 +163,11 @@ class HiCacheAck(NamedTuple): node_ids: List[int] num_tokens: int = 0 timing_enabled: bool = False + # Tokens transferred per host pool (PoolName value -> count). + num_tokens_by_pool: Optional[dict[str, int]] = None + # Total bytes moved by the op across all pools, including draft piggyback + # and sidecar transfers that the per-pool token counts exclude. + num_bytes: int = 0 class StorageOperation: @@ -714,11 +719,12 @@ class HiCacheController: self.write_queue.clear() start_event = device_module.Event() - finish_event = device_module.Event() + ack_start_event, ack_finish_event, timing_enabled = make_timing_event_pair() start_event.record() with device_module.stream(self.write_stream): start_event.wait(self.write_stream) + ack_start_event.record() self.mem_pool_host.backup_from_device_all_layer( self.mem_pool_device, host_indices, device_indices, self.io_backend ) @@ -729,7 +735,7 @@ class HiCacheController: device_indices, self.io_backend, ) - finish_event.record() + ack_finish_event.record() # NOTE: We must save the host indices and device indices here, # this is because we need to guarantee that these tensors are # still alive when the write stream is executing. @@ -738,7 +744,25 @@ class HiCacheController: if device_indices.is_cuda: device_indices.record_stream(self.write_stream) - self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) + self.ack_write_queue.append( + HiCacheAck( + start_event=ack_start_event, + finish_event=ack_finish_event, + node_ids=op.node_ids, + num_tokens=len(op.device_indices), + timing_enabled=timing_enabled, + num_tokens_by_pool={PoolName.KV.value: len(op.device_indices)}, + num_bytes=self._transfer_num_bytes(op), + ) + ) + + def _transfer_num_bytes(self, op: CacheOperation) -> int: + """Total bytes moved by a merged transfer op (draft piggyback included).""" + num_tokens = len(op.device_indices) + num_bytes = num_tokens * self.mem_pool_host.size_per_token + if self.has_draft: + num_bytes += num_tokens * self.mem_pool_host_draft.size_per_token + return num_bytes def load( self, @@ -830,6 +854,8 @@ class HiCacheController: node_ids=op.node_ids, num_tokens=len(op.device_indices), timing_enabled=timing_enabled, + num_tokens_by_pool={PoolName.KV.value: len(op.device_indices)}, + num_bytes=self._transfer_num_bytes(op), ) ) return producer_id diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index e20854b33..7ceeabf77 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -191,6 +191,21 @@ 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 +) -> 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). + """ + storage = min(host_hit_len, storage_hit_len) + host = host_hit_len - storage + device = max(0, prefix_len - host_hit_len) + return device, host, storage + + def _compute_pad_value(hash: int) -> int: """Compute pad value from hash.""" return MM_PAD_SHIFT_VALUE + (hash % (1 << 30)) @@ -2394,24 +2409,17 @@ 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: - # At this point, prefix_indices has been extended with host data - # via init_load_back in schedule_policy, so: - # - len(prefix_indices) = device_original + host_loaded - # - host_hit_length = total tokens from host cache (including storage-prefetched) - # - storage_hit_length = tokens loaded from storage backend (L3 hits) - # - device_portion = len(prefix_indices) - host_hit_length - # - # Storage hits are now tracked via scheduler after prefetch completes. # storage_hit_length is set by scheduler.pop_prefetch_loaded_tokens() - host_total = req.host_hit_length - # Clamp storage to host_total to handle edge cases - storage_portion = min(host_total, req.storage_hit_length) - host_portion = host_total - storage_portion - device_portion = max(0, len(req.prefix_indices) - host_total) - - req.cached_tokens_device = device_portion - req.cached_tokens_host = host_portion - req.cached_tokens_storage = storage_portion + # after prefetch completes. + ( + 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, + storage_hit_len=req.storage_hit_length, + ) req._cache_breakdown_computed = True req.already_computed = seq_len diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 39bc02b78..65526e595 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -38,7 +38,11 @@ import torch from sglang.srt.dllm.config import DllmConfig from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled -from sglang.srt.managers.schedule_batch import Req, ScheduleBatch +from sglang.srt.managers.schedule_batch import ( + Req, + ScheduleBatch, + split_cached_prefix_by_tier, +) from sglang.srt.mem_cache.allocator.hisparse import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, ) @@ -483,6 +487,9 @@ class PrefillAdder: self.new_chunked_req = None self.log_hit_tokens = 0 self.reprocessed_log_hit_tokens = 0 + self.log_device_hit_tokens = 0 + self.log_host_hit_tokens = 0 + self.log_storage_hit_tokens = 0 # TODO(lsyin): report the real input tokens excluding page alignment self.log_input_tokens = 0 self.reprocessed_log_input_tokens = 0 @@ -745,6 +752,8 @@ 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) @@ -784,6 +793,15 @@ class PrefillAdder: if retracted_stain: self.reprocessed_log_hit_tokens += prefix_len self.reprocessed_log_input_tokens += extend_input_len + elif 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, + ) + self.log_device_hit_tokens += device_hit + self.log_host_hit_tokens += host_hit + self.log_storage_hit_tokens += storage_hit def _get_dllm_remain_tokens(self) -> int: _rem_tokens = min( @@ -816,6 +834,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, ) def _req_inc_lock_ref(self, req: Req): @@ -1209,6 +1229,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, ) else: # Make sure at least one page is available @@ -1250,6 +1272,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, ) return self.budget_state() diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 13ba75da5..9e1f28cef 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -7,13 +7,7 @@ import tempfile import time from collections import defaultdict from dataclasses import dataclass -from typing import ( - TYPE_CHECKING, - List, - Optional, - Tuple, - Union, -) +from typing import TYPE_CHECKING, List, Optional, Tuple, Union from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.environ import envs @@ -64,6 +58,9 @@ class PrefillStats: num_new_seqs: int # len(can_run_list) reprocessed_log_input_tokens: int = 0 reprocessed_log_hit_tokens: int = 0 + log_device_hit_tokens: int = 0 + log_host_hit_tokens: int = 0 + log_storage_hit_tokens: int = 0 num_pending_tokens: int = 0 @classmethod @@ -79,6 +76,9 @@ class PrefillStats: log_hit_tokens=adder.log_hit_tokens, reprocessed_log_input_tokens=adder.reprocessed_log_input_tokens, reprocessed_log_hit_tokens=adder.reprocessed_log_hit_tokens, + log_device_hit_tokens=adder.log_device_hit_tokens, + log_host_hit_tokens=adder.log_host_hit_tokens, + log_storage_hit_tokens=adder.log_storage_hit_tokens, new_token_ratio=adder.new_token_ratio, num_running_reqs=QueueCount.from_reqs( running_reqs, enable_priority_scheduling @@ -637,6 +637,12 @@ class SchedulerMetricsReporter: cache_hit_rate = ( effective_hit_tokens / total_tokens if total_tokens > 0 else 0.0 ) + self.metrics_collector.increment_effective_prefill_tokens( + input_tokens=effective_input_tokens, + device_hit_tokens=prefill_stats.log_device_hit_tokens, + host_hit_tokens=prefill_stats.log_host_hit_tokens, + storage_hit_tokens=prefill_stats.log_storage_hit_tokens, + ) # Basics if ( @@ -970,9 +976,7 @@ class SchedulerMetricsReporter: if not self.scheduler.enable_fpm: return - from sglang.srt.observability.forward_pass_metrics import ( - ForwardPassMetrics, - ) + from sglang.srt.observability.forward_pass_metrics import ForwardPassMetrics if self.scheduler._fpm_uses_device_timer: self.forward_pass_device_timer._report() diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index dbcec3547..8671a7a75 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -455,7 +455,13 @@ class HiMambaRadixCache(MambaRadixCache): self.dec_lock_ref(end_node) if self.metrics_collector is not None: - self.metrics_collector.increment_load_back_num_tokens(ack.num_tokens) + for pool, num_tokens in (ack.num_tokens_by_pool or {}).items(): + if num_tokens > 0: + self.metrics_collector.increment_load_back_num_tokens( + num_tokens=num_tokens, pool=pool + ) + if ack.num_bytes > 0: + self.metrics_collector.increment_load_back_num_bytes(ack.num_bytes) if ack.timing_enabled: duration_ms = ack.start_event.elapsed_time(ack.finish_event) self.metrics_collector.observe_load_back_duration( diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 5ebac56f3..74dec738f 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1028,6 +1028,7 @@ class HiRadixCache(RadixCache): ack.finish_event.synchronize() for ack_id in ack.node_ids: self._finish_write_through_ack(ack_id, release_lock=False) + self._log_write_ack_metrics(ack) self.cache_controller.ack_write_queue.clear() assert len(self.ongoing_write_through) == 0 return @@ -1055,8 +1056,24 @@ class HiRadixCache(RadixCache): ack.finish_event.synchronize() for ack_id in ack.node_ids: self._finish_write_through_ack(ack_id, release_lock=True) + self._log_write_ack_metrics(ack) finish_count -= 1 + def _log_write_ack_metrics(self, ack) -> None: + """Record D->H backup volume and duration for a completed write ack.""" + if self.metrics_collector is None: + return + for pool, num_tokens in (ack.num_tokens_by_pool or {}).items(): + if num_tokens > 0: + self.metrics_collector.increment_backup_num_tokens( + num_tokens=num_tokens, pool=pool + ) + if ack.num_bytes > 0: + self.metrics_collector.increment_backup_num_bytes(ack.num_bytes) + if ack.timing_enabled: + duration_ms = ack.start_event.elapsed_time(ack.finish_event) + self.metrics_collector.observe_backup_duration(duration_ms / 1000.0) + def loading_check(self, finish_count: Optional[int] = None): if finish_count is None: finish_count = 0 @@ -1080,7 +1097,13 @@ class HiRadixCache(RadixCache): self.dec_lock_ref(end_node) if self.metrics_collector is not None: - self.metrics_collector.increment_load_back_num_tokens(ack.num_tokens) + for pool, num_tokens in (ack.num_tokens_by_pool or {}).items(): + if num_tokens > 0: + self.metrics_collector.increment_load_back_num_tokens( + num_tokens=num_tokens, pool=pool + ) + if ack.num_bytes > 0: + self.metrics_collector.increment_load_back_num_bytes(ack.num_bytes) if ack.timing_enabled: duration_ms = ack.start_event.elapsed_time(ack.finish_event) self.metrics_collector.observe_load_back_duration( @@ -1300,6 +1323,12 @@ class HiRadixCache(RadixCache): root.parent.children.pop(key, None) self._update_leaf_status(root.parent) self._update_host_leaf_status(root.parent) + if freed_device > 0 and self.metrics_collector is not None: + self.metrics_collector.increment_dropped_tokens( + num_tokens=freed_device, + reason="host_pressure", + pool=PoolName.KV.value, + ) return freed_device def evict_host(self, num_tokens: int): diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 5764a9f60..191cd8a90 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -419,10 +419,11 @@ class HybridCacheController(BaseHiCacheController): ) self.write_queue.clear() start_event = device_module.Event() - finish_event = device_module.Event() + ack_start_event, ack_finish_event, timing_enabled = make_timing_event_pair() start_event.record() with device_module.stream(self.write_stream): start_event.wait(self.write_stream) + ack_start_event.record() self.mem_pool_host.backup_from_device_all_layer( self.mem_pool_device, host_indices, @@ -437,14 +438,60 @@ class HybridCacheController(BaseHiCacheController): device_indices, self.io_backend, ) - finish_event.record() + ack_finish_event.record() self._record_transfer_indices_on_stream( self.write_stream, host_indices, device_indices, resolved_pool_transfers, ) - self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) + self.ack_write_queue.append( + HiCacheAck( + start_event=ack_start_event, + finish_event=ack_finish_event, + node_ids=op.node_ids, + num_tokens=len(op.device_indices), + timing_enabled=timing_enabled, + num_tokens_by_pool=self._num_tokens_by_pool(op), + num_bytes=self._transfer_num_bytes(op), + ) + ) + + def _num_tokens_by_pool(self, op: CacheOperation) -> dict[str, int]: + """Per-pool token counts for a merged transfer op (anchor + extra + pools), shared by D->H write and H->D load acks; sidecar transfers + reusing another pool's indices are excluded.""" + counts = {self.mem_pool_host.anchor_entry.name.value: len(op.device_indices)} + for transfer in op.pool_transfers or []: + if transfer.indices_from_pool is not None or transfer.host_indices is None: + continue + name = transfer.name.value + counts[name] = counts.get(name, 0) + len(transfer.host_indices) + return counts + + def _transfer_num_bytes(self, op: CacheOperation) -> int: + """Total bytes moved by a merged transfer op across all pools, + including draft piggyback and sidecar transfers riding another + pool's indices (both excluded from the per-pool token counts).""" + kv_tokens = len(op.device_indices) + num_bytes = kv_tokens * self.mem_pool_host.anchor_entry.host_pool.size_per_token + if self.has_draft: + num_bytes += kv_tokens * self.mem_pool_host_draft.size_per_token + # Slot counts of the pools sidecars can ride on. + source_len = {self.mem_pool_host.anchor_entry.name: kv_tokens} + for t in op.pool_transfers or []: + if t.indices_from_pool is None and t.host_indices is not None: + source_len[t.name] = len(t.host_indices) + for t in op.pool_transfers or []: + entry = self.mem_pool_host.entry_map.get(t.name) + if entry is None: + continue + if t.indices_from_pool is not None: + num_slots = source_len.get(t.indices_from_pool, 0) + else: + num_slots = len(t.host_indices) if t.host_indices is not None else 0 + num_bytes += num_slots * entry.host_pool.size_per_token + return num_bytes def load( self, @@ -542,6 +589,8 @@ class HybridCacheController(BaseHiCacheController): op.node_ids, num_tokens=len(op.device_indices), timing_enabled=timing_enabled, + num_tokens_by_pool=self._num_tokens_by_pool(op), + num_bytes=self._transfer_num_bytes(op), ) ) return producer_id diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index d43623826..a85a7157e 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -72,6 +72,7 @@ from sglang.srt.observability.metrics_collector import ( from sglang.srt.session.streaming_session import StreamingSession if TYPE_CHECKING: + from sglang.srt.managers.cache_controller import HiCacheAck from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( @@ -82,6 +83,14 @@ if TYPE_CHECKING: 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, @@ -360,6 +369,14 @@ 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 + if self.metrics_collector is not None and self.cache_controller is not None: + for ct in self.tree_components: + self.metrics_collector.increment_dropped_tokens( + num_tokens=0, + reason="host_pressure", + pool=_COMPONENT_POOL_LABEL[ct], + ) self.load_back_threshold = 10 self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy @@ -436,7 +453,8 @@ class UnifiedRadixCache(BasePrefixCache): ): self.writing_check(write_back=True) - self.update_eviction_metrics(sum(tracker.values()), start_time) + # Report full-layer tokens only + self.update_eviction_metrics(tracker[BASE_COMPONENT_TYPE], start_time) return EvictResult( num_tokens_evicted=tracker[BASE_COMPONENT_TYPE], swa_num_tokens_evicted=tracker.get(ComponentType.SWA, 0), @@ -519,10 +537,12 @@ 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", @@ -539,6 +559,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: + 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], + ) + def inc_lock_ref( self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = () ) -> IncLockRefResult: @@ -1832,6 +1869,7 @@ class UnifiedRadixCache(BasePrefixCache): for ack_id in ack.node_ids: if ack_id in self.ongoing_write_through: self._finish_write_through_ack(ack_id) + self._log_write_ack_metrics(ack) cc.ack_write_queue.clear() assert len(self.ongoing_write_through) == 0 return @@ -1854,8 +1892,24 @@ class UnifiedRadixCache(BasePrefixCache): ack.finish_event.synchronize() for ack_id in ack.node_ids: self._finish_write_through_ack(ack_id) + self._log_write_ack_metrics(ack) finish_count -= 1 + def _log_write_ack_metrics(self, ack: HiCacheAck) -> None: + """Record D->H backup volume and duration for a completed write ack.""" + if self.metrics_collector is None: + return + for pool, num_tokens in (ack.num_tokens_by_pool or {}).items(): + if num_tokens > 0: + self.metrics_collector.increment_backup_num_tokens( + num_tokens=num_tokens, pool=pool + ) + if ack.num_bytes > 0: + self.metrics_collector.increment_backup_num_bytes(ack.num_bytes) + if ack.timing_enabled: + duration_ms = ack.start_event.elapsed_time(ack.finish_event) + self.metrics_collector.observe_backup_duration(duration_ms / 1000.0) + def loading_check(self, finish_count: Optional[int] = None) -> None: """Poll load-back completions.""" cc = self.cache_controller @@ -1882,7 +1936,13 @@ class UnifiedRadixCache(BasePrefixCache): self.dec_host_lock_ref(node, host_lock_params) if self.metrics_collector is not None: - self.metrics_collector.increment_load_back_num_tokens(ack.num_tokens) + for pool, num_tokens in (ack.num_tokens_by_pool or {}).items(): + if num_tokens > 0: + self.metrics_collector.increment_load_back_num_tokens( + num_tokens=num_tokens, pool=pool + ) + if ack.num_bytes > 0: + self.metrics_collector.increment_load_back_num_bytes(ack.num_bytes) if ack.timing_enabled: duration_ms = ack.start_event.elapsed_time(ack.finish_event) self.metrics_collector.observe_load_back_duration( diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index d08ce2aff..a5976c120 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -889,6 +889,20 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): ), labelnames=list(labels.keys()) + ["mode"], ) + self.prefill_effective_tokens_total = Counter( + name="sglang:prefill_effective_tokens_total", + documentation=( + "Effective prefill tokens with retracted-request re-counts " + "excluded, updated on each log interval. mode: device_hit, " + "host_hit, storage_hit, input. Windowed prefix cache hit " + "rate = rate(sum of *_hit) / rate(sum of all modes); " + "per-tier rate uses a single *_hit mode in the numerator." + ), + labelnames=list(labels.keys()) + ["mode"], + ) + # Pre-seed every mode at 0 so per-tier ratio charts get a complete operand set + for mode in ("input", "device_hit", "host_hit", "storage_hit"): + self.prefill_effective_tokens_total.labels(**labels, mode=mode) self.forward_execution_seconds_total = Counter( name="sglang:forward_execution_seconds_total", documentation=( @@ -1248,6 +1262,24 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): **dp_cooperation_info.to_labels(), ).inc(delta) + def increment_effective_prefill_tokens( + self, + input_tokens: int, + device_hit_tokens: int, + host_hit_tokens: int, + storage_hit_tokens: int, + ) -> None: + for mode, delta in [ + ("input", input_tokens), + ("device_hit", device_hit_tokens), + ("host_hit", host_hit_tokens), + ("storage_hit", storage_hit_tokens), + ]: + if delta > 0: + self.prefill_effective_tokens_total.labels( + **self.labels, mode=mode + ).inc(delta) + def increment_forward_execution_seconds( self, category: str, @@ -1964,6 +1996,11 @@ 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" @@ -1989,16 +2026,43 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin): 0.5, 1.0, ] + # 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. + bucket_backup_duration = [ + 0.001, + 0.002, + 0.005, + 0.01, + 0.02, + 0.05, + 0.1, + 0.2, + 0.5, + 1.0, + 2.0, + 5.0, + 10.0, + 30.0, + 60.0, + ] + self.eviction_duration_seconds = Histogram( name="sglang:eviction_duration_seconds", - documentation="Time taken to evict memory from GPU to CPU in seconds.", + documentation="End-to-end time of a device eviction pass in " + "seconds; under --hicache-write-policy write_back this includes " + "the blocking D->H backup (see " + "sglang:hicache_backup_duration_seconds for the copy alone).", labelnames=labels.keys(), buckets=bucket_eviction_duration, ) self.eviction_num_tokens = Counter( name="sglang:evicted_tokens_total", - documentation="The number of tokens evicted from GPU to CPU.", + documentation="The number of device KV token slots freed by " + "eviction, regardless of whether the data was backed up to host " + "(see sglang:hicache_backup_tokens_total) or destroyed (see " + "sglang:hicache_dropped_tokens_total).", labelnames=labels.keys(), ) @@ -2011,15 +2075,63 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin): self.load_back_num_tokens = Counter( name="sglang:load_back_tokens_total", - documentation="The number of tokens loaded from CPU to GPU.", + documentation="The number of tokens loaded back from local host " + "DRAM (L2) to GPU, by host pool (kv, swa, mamba, ...).", + labelnames=list(labels.keys()) + ["pool"], + ) + + self.backup_duration_seconds = Histogram( + name="sglang:hicache_backup_duration_seconds", + documentation="Time taken to back up KV cache from GPU to local " + "host DRAM (L2) in seconds, per merged write op. Covers all D->H " + "backups regardless of --hicache-write-policy. Distinct from the " + "host-to-storage (L3) sglang:backuped_tokens_total.", labelnames=labels.keys(), + buckets=bucket_backup_duration, + ) + + self.backup_num_bytes = Counter( + name="sglang:hicache_backup_bytes_total", + documentation="Bytes backed up from GPU to local host DRAM (L2), " + "all pools combined, including draft/sidecar transfers that the " + "token counter excludes. Divided by the rate of " + "hicache_backup_duration_seconds_sum, gives the achieved D->H " + "bandwidth while transferring.", + labelnames=labels.keys(), + ) + + self.load_back_num_bytes = Counter( + name="sglang:load_back_bytes_total", + 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.", + labelnames=labels.keys(), + ) + + self.backup_num_tokens = Counter( + name="sglang:hicache_backup_tokens_total", + documentation="The number of tokens backed up from GPU to local " + "host DRAM (L2), by host pool (kv, swa, mamba, ...). Covers all " + "D->H backups regardless of --hicache-write-policy. Distinct from " + "the host-to-storage (L3) sglang:backuped_tokens_total.", + labelnames=list(labels.keys()) + ["pool"], + ) + + 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).", + labelnames=list(labels.keys()) + ["reason", "pool"], ) def increment_eviction_num_tokens(self, num_tokens: int) -> None: self.eviction_num_tokens.labels(**self.labels).inc(num_tokens) - def increment_load_back_num_tokens(self, num_tokens: int) -> None: - self.load_back_num_tokens.labels(**self.labels).inc(num_tokens) + def increment_load_back_num_tokens(self, num_tokens: int, pool: str) -> None: + self.load_back_num_tokens.labels(**self.labels, pool=pool).inc(num_tokens) def observe_eviction_duration(self, duration_seconds: float) -> None: self.eviction_duration_seconds.labels(**self.labels).observe(duration_seconds) @@ -2027,6 +2139,23 @@ class RadixCacheMetricsCollector(_StatLoggerDIMixin): def observe_load_back_duration(self, duration_seconds: float) -> None: self.load_back_duration_seconds.labels(**self.labels).observe(duration_seconds) + def increment_backup_num_tokens(self, num_tokens: int, pool: str) -> None: + self.backup_num_tokens.labels(**self.labels, pool=pool).inc(num_tokens) + + def increment_backup_num_bytes(self, num_bytes: int) -> None: + self.backup_num_bytes.labels(**self.labels).inc(num_bytes) + + def increment_load_back_num_bytes(self, num_bytes: int) -> None: + self.load_back_num_bytes.labels(**self.labels).inc(num_bytes) + + def observe_backup_duration(self, duration_seconds: float) -> None: + self.backup_duration_seconds.labels(**self.labels).observe(duration_seconds) + + def increment_dropped_tokens(self, num_tokens: int, reason: str, pool: str) -> None: + self.hicache_dropped_tokens.labels(**self.labels, reason=reason, pool=pool).inc( + num_tokens + ) + class EncoderMetricsCollector(_StatLoggerDIMixin): """Metrics collector for the EPD encoder server (--encoder-only).""" diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index c11332c70..baf3b74f6 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -91,6 +91,8 @@ class TestPrefillAdder(CustomTestCase): req.sampling_params = SimpleNamespace(max_new_tokens=max_new_tokens) req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time) req.retracted_stain = False + req.host_hit_length = 0 + req.storage_hit_length = 0 req.finished.return_value = False req.needs_host_load_back.return_value = False return req diff --git a/test/registered/unit/mem_cache/test_hicache_load_back_timing.py b/test/registered/unit/mem_cache/test_hicache_load_back_timing.py index e056a07a5..b92e63475 100644 --- a/test/registered/unit/mem_cache/test_hicache_load_back_timing.py +++ b/test/registered/unit/mem_cache/test_hicache_load_back_timing.py @@ -65,6 +65,7 @@ class TestLoadBackDurationMetric(CustomTestCase): node_ids=[1, 2], num_tokens=1024, timing_enabled=True, + num_tokens_by_pool={"kv": 1024}, ) stub = object.__new__(HiRadixCache) stub.cache_controller = SimpleNamespace(ack_load_queue=[ack]) @@ -77,7 +78,7 @@ class TestLoadBackDurationMetric(CustomTestCase): stub.loading_check() stub.metrics_collector.increment_load_back_num_tokens.assert_called_once_with( - 1024 + num_tokens=1024, pool="kv" ) stub.metrics_collector.observe_load_back_duration.assert_called_once() (observed,), _ = stub.metrics_collector.observe_load_back_duration.call_args @@ -100,6 +101,7 @@ class TestLoadBackDurationMetric(CustomTestCase): node_ids=[7], num_tokens=512, timing_enabled=False, + num_tokens_by_pool={"kv": 512}, ) stub = object.__new__(HiRadixCache) stub.cache_controller = SimpleNamespace(ack_load_queue=[ack]) @@ -112,7 +114,7 @@ class TestLoadBackDurationMetric(CustomTestCase): stub.loading_check() stub.metrics_collector.increment_load_back_num_tokens.assert_called_once_with( - 512 + num_tokens=512, pool="kv" ) stub.metrics_collector.observe_load_back_duration.assert_not_called() self.assertEqual(stub.cache_controller.ack_load_queue, []) 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 8c79c3565..57c0353bd 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 @@ -139,6 +139,9 @@ def _cpu_per_layer_pf_lf_copy( class _FakeEvent: + def __init__(self, enable_timing=False): + self.enable_timing = enable_timing + def record(self): pass @@ -156,6 +159,15 @@ class _FakeDeviceModule: class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): + def setUp(self): + # start_writing probes timing support via a module-cached check; + # clear it on both sides so results from (or against) the fake + # device module never leak across tests. + manager_cache_controller._timing_events_supported.cache_clear() + + def tearDown(self): + manager_cache_controller._timing_events_supported.cache_clear() + def _patched_transfers(self, src_registry=None, module=MEMORY_POOL_HOST_MODULE): staged_side_effect = None if src_registry is not None: @@ -678,6 +690,10 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): class FakeHostGroup: layout = "page_first" can_use_write_back_jit = True + anchor_entry = SimpleNamespace( + name=PoolName.KV, host_pool=SimpleNamespace(size_per_token=2) + ) + entry_map = {} def backup_from_device_all_layer( self, @@ -718,8 +734,13 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): ) ) - with mock.patch.object( - hybrid_cache_controller, "device_module", _FakeDeviceModule + with ( + mock.patch.object( + hybrid_cache_controller, "device_module", _FakeDeviceModule + ), + mock.patch.object( + manager_cache_controller, "device_module", _FakeDeviceModule + ), ): controller.start_writing() @@ -733,6 +754,10 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): class FakeHostGroup: layout = "page_first" can_use_write_back_jit = False + anchor_entry = SimpleNamespace( + name=PoolName.KV, host_pool=SimpleNamespace(size_per_token=2) + ) + entry_map = {} def backup_from_device_all_layer( self, @@ -770,8 +795,13 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): return_value=(op.host_indices, op.device_indices, op.pool_transfers) ) - with mock.patch.object( - hybrid_cache_controller, "device_module", _FakeDeviceModule + with ( + mock.patch.object( + hybrid_cache_controller, "device_module", _FakeDeviceModule + ), + mock.patch.object( + manager_cache_controller, "device_module", _FakeDeviceModule + ), ): controller.start_writing() @@ -785,6 +815,7 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): class FakeHostPool: layout = "page_first" can_use_write_back_jit = True + size_per_token = 2 def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend @@ -825,6 +856,7 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): class FakeHostPool: layout = "page_first" can_use_write_back_jit = False + size_per_token = 2 def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend