diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index 6321efb2e94d..f750c558f416 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -99,6 +99,8 @@ def __init__( storage_backend=server_args.hicache_storage_backend, model_name=server_args.served_model_name, storage_backend_extra_config=hicache_storage_backend_extra_config, + enable_metrics=server_args.enable_metrics, + extra_metric_labels=server_args.extra_metric_labels, ) self.ongoing_offload = {} @@ -219,17 +221,15 @@ def check_offload_progress(self): def _check_offload_progress(self, finish_count): """Check the progress of offload from device to host.""" while finish_count > 0: - _, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0) - finish_event.synchronize() - for ack_id in ack_list: - ( - req, - host_indices, - incremental_tokens, - start_time, - start, - end, - ) = self.ongoing_offload.pop(ack_id) + ack = self.cache_controller.ack_write_queue.pop(0) + ack.finish_event.synchronize() + matched = False + for ack_id in ack.node_ids: + entry = self.ongoing_offload.pop(ack_id, None) + if entry is None: + continue + matched = True + req, host_indices, incremental_tokens, start_time, start, end = entry self._mark_offload_finished(req.rid) prior_hash = ( @@ -247,6 +247,11 @@ def _check_offload_progress(self, finish_count): state = self.offloaded_state.get(req.rid) start_offset = state.prefill_len if state is not None else start self._release_finished_req(req, start_offset) + if matched: + self.cache_controller.record_l1_l2_transfer_complete( + direction="offload", + ack=ack, + ) finish_count -= 1 def _release_finished_req(self, req: Req, start_offset: int): diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index ca276bc11479..5dd24300aeac 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -28,6 +28,9 @@ PoolName, PoolTransfer, ) +from sglang.srt.observability.metrics_collector import ( + HiCacheL1L2TransferMetricsCollector, +) if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator @@ -56,8 +59,17 @@ class LayerLoadingEvent: def __init__(self, num_layers: int): self._num_layers = num_layers - self.load_events = [device_module.Event() for _ in range(num_layers)] - self.start_event = device_module.Event() # start event on controller stream + # The last layer's event doubles as finish_event; together with + # start_event it feeds elapsed_time() in record_l1_l2_transfer_complete, + # which requires both events to be created with enable_timing=True. + # Other per-layer events stay timing-free (only used for stream waits). + self.load_events = [ + device_module.Event(enable_timing=(i == num_layers - 1)) + for i in range(num_layers) + ] + self.start_event = device_module.Event( + enable_timing=True + ) # start event on controller stream def complete(self, layer_index: int): assert 0 <= layer_index < self._num_layers @@ -148,6 +160,15 @@ class HiCacheAck(NamedTuple): finish_event: device_module.Event node_ids: List[int] + # Number of KV token slots moved by this merged operation. + token_count: int + + # Number of KV blocks moved. For HiCache this should be token_count // page_size. + block_count: int + + # Estimated total bytes moved for this operation. + byte_count: int + class StorageOperation: counter = 0 @@ -225,6 +246,8 @@ def __init__( model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, enable_storage_metrics: bool = False, + enable_metrics: bool = False, + extra_metric_labels: Optional[dict[str, str]] = None, ): self.tp_group = tp_group self.attn_cp_group = attn_cp_group @@ -246,6 +269,34 @@ def __init__( self.storage_backend = None self.storage_backend_type = None self.enable_storage_metrics = enable_storage_metrics + self._warned_unknown_host_pool_bytes = False + + # init L1/L2 transfer metrics collection (device-host transfers triggered by write/load). + self.enable_l1_l2_transfer_metrics = enable_metrics + self.hicache_l1_l2_transfer_metrics_collector = None + + self.hicache_l1_l2_transfer_totals = { + "offload": { + "events": 0, + "blocks": 0, + "bytes": 0, + "xfer_us": 0, + }, + "onboard": { + "events": 0, + "blocks": 0, + "bytes": 0, + "xfer_us": 0, + }, + } + + if self.enable_l1_l2_transfer_metrics: + # pp_rank/pp_size are otherwise only set in _generate_storage_config, + # which runs only when a storage backend is attached. Derive them here + # so the transfer-metrics labels are available without L3 storage. + self.pp_rank = get_pipeline_model_parallel_rank() + self.pp_size = get_pipeline_model_parallel_world_size() + self._init_l1_l2_transfer_metrics(extra_metric_labels) # Draft KV pool support (best-effort piggyback on target L2/L3 ops). self.has_draft = False @@ -689,9 +740,12 @@ def start_writing(self) -> None: ) self.write_queue.clear() - start_event = device_module.Event() - finish_event = device_module.Event() + # enable_timing so record_l1_l2_transfer_complete can use elapsed_time() + # for the real device transfer duration (host fallback returns None). + start_event = device_module.Event(enable_timing=True) + finish_event = device_module.Event(enable_timing=True) + token_count = int(host_indices.numel()) start_event.record() with device_module.stream(self.write_stream): start_event.wait(self.write_stream) @@ -714,7 +768,14 @@ def start_writing(self) -> None: 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( + self._make_hicache_ack( + start_event=start_event, + finish_event=finish_event, + node_ids=op.node_ids, + token_count=token_count, + ) + ) def load( self, @@ -766,6 +827,8 @@ def start_loading(self) -> int: ) self.load_queue.clear() producer_event = self.layer_done_counter.events[producer_id] + + token_count = int(host_indices.numel()) producer_event.start_event.record() with device_module.stream(self.load_stream): @@ -796,10 +859,11 @@ def start_loading(self) -> int: device_indices.record_stream(self.load_stream) self.ack_load_queue.append( - HiCacheAck( + self._make_hicache_ack( start_event=producer_event.start_event, finish_event=producer_event.finish_event, node_ids=op.node_ids, + token_count=token_count, ) ) return producer_id @@ -1206,3 +1270,229 @@ def backup_thread_func(self): except Empty: continue + + def _init_l1_l2_transfer_metrics( + self, + extra_metric_labels: Optional[dict[str, str]] = None, + ) -> None: + """Initialize Prometheus metrics for L1<->L2 transfer accounting.""" + + from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + ) + from sglang.srt.layers.dp_attention import ( + get_attention_dp_rank, + get_attention_tp_rank, + get_attention_tp_size, + is_dp_attention_enabled, + ) + + if is_dp_attention_enabled(): + tp_rank = get_attention_tp_rank() + tp_size = get_attention_tp_size() + dp_rank = get_attention_dp_rank() + else: + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + dp_rank = 0 + + attn_cp_rank, attn_cp_size = self.get_attn_cp_rank_and_size() + + labels = { + "tp_rank": str(tp_rank), + "tp_size": str(tp_size), + "dp_rank": str(dp_rank), + "pp_rank": str(self.pp_rank), + "pp_size": str(self.pp_size), + "attn_cp_rank": str(attn_cp_rank), + "attn_cp_size": str(attn_cp_size), + "io_backend": str(self.io_backend), + } + + if extra_metric_labels: + labels.update({k: str(v) for k, v in extra_metric_labels.items()}) + + self.hicache_l1_l2_transfer_metrics_collector = ( + HiCacheL1L2TransferMetricsCollector(labels) + ) + + def _host_pool_size_per_token(self, pool_name=None) -> int: + """Return size_per_token for a specific named pool, or the whole pool for plain caches.""" + if hasattr(self.mem_pool_host, "entry_map") and pool_name is not None: + entry = self.mem_pool_host.entry_map.get(pool_name) + if entry is None: + return 0 + return int(getattr(entry.host_pool, "size_per_token", 0) or 0) + size_per_token = getattr(self.mem_pool_host, "size_per_token", None) + if size_per_token is None: + self._warn_unknown_host_pool_bytes(self.mem_pool_host) + return 0 + return int(size_per_token) + + def _transfer_index_count(self, transfer) -> int: + """Return the number of index slots in a PoolTransfer (host preferred, device fallback).""" + if transfer.host_indices is not None: + return int(transfer.host_indices.numel()) + if transfer.device_indices is not None: + return int(transfer.device_indices.numel()) + return 0 + + def _warn_unknown_host_pool_bytes(self, host_pool) -> None: + if self._warned_unknown_host_pool_bytes: + return + self._warned_unknown_host_pool_bytes = True + logger.warning( + "Unable to estimate HiCache L1/L2 transfer bytes for host pool type %s: " + "missing size_per_token.", + type(host_pool).__name__, + ) + + def _estimate_l1_l2_transfer_bytes( + self, + token_count: int, + pool_transfers: Optional[List[PoolTransfer]] = None, + ) -> int: + """Estimate total bytes moved for one L1<->L2 transfer operation. + + For hybrid caches (Mamba, SWA, DSA sidecars) each sidecar pool may use + a different index count than the primary KV token count, so per-transfer + index counts are used rather than assuming every pool moves token_count slots. + """ + total = int(token_count) * self._host_pool_size_per_token(PoolName.KV) + + for transfer in pool_transfers or []: + total += self._transfer_index_count( + transfer + ) * self._host_pool_size_per_token(transfer.name) + + if self.has_draft and self.mem_pool_host_draft is not None: + draft_size_per_token = getattr( + self.mem_pool_host_draft, "size_per_token", None + ) + if draft_size_per_token is None: + self._warn_unknown_host_pool_bytes(self.mem_pool_host_draft) + else: + total += int(token_count) * int(draft_size_per_token) + + return total + + def _make_hicache_ack( + self, + *, + start_event: device_module.Event, + finish_event: device_module.Event, + node_ids: List[int], + token_count: int, + pool_transfers: Optional[List[PoolTransfer]] = None, + ) -> HiCacheAck: + block_count = (int(token_count) + int(self.page_size) - 1) // int( + self.page_size + ) + byte_count = self._estimate_l1_l2_transfer_bytes(token_count, pool_transfers) + + return HiCacheAck( + start_event=start_event, + finish_event=finish_event, + node_ids=node_ids, + token_count=int(token_count), + block_count=block_count, + byte_count=byte_count, + ) + + def _transfer_elapsed_us(self, ack: HiCacheAck) -> Optional[int]: + """Return device-event transfer duration in microseconds, when available.""" + try: + return max(0, int(ack.start_event.elapsed_time(ack.finish_event) * 1000)) + except Exception: + return None + + def record_l1_l2_transfer_complete( + self, + *, + direction: str, + ack: HiCacheAck, + ) -> None: + """Record logs and Prometheus metrics after a transfer ack completes. + + direction: + - "offload": L1 -> L2 + - "onboard": L2 -> L1 + """ + should_log = logger.isEnabledFor(logging.DEBUG) + should_record_metrics = ( + self.hicache_l1_l2_transfer_metrics_collector is not None + ) + if not should_log and not should_record_metrics: + return + + if direction == "offload": + action = "Offload" + src = "sglang_hicache::L1" + dst = "sglang_hicache::L2" + elif direction == "onboard": + action = "Onboard" + src = "sglang_hicache::L2" + dst = "sglang_hicache::L1" + else: + raise ValueError(f"Unknown HiCache L1/L2 transfer direction: {direction}") + + xfer_us = self._transfer_elapsed_us(ack) + + if should_log: + ts_us = time.time_ns() // 1000 + logger.debug( + "%s transfer complete ts_us=%d blocks=%d bytes=%d xfer_us=%s " + "bandwidth=%.2fGB/s " + 'src="%s" dst="%s"', + action, + ts_us, + ack.block_count, + ack.byte_count, + xfer_us if xfer_us is not None else "unavailable", + ack.byte_count * 0.001 / xfer_us if xfer_us else 0, + src, + dst, + ) + + if should_record_metrics: + self.hicache_l1_l2_transfer_metrics_collector.record_transfer( + direction=direction, + src=src, + dst=dst, + blocks=ack.block_count, + bytes_=ack.byte_count, + xfer_us=xfer_us, + ) + + # Accumulate totals whenever transfer observation is active. This keeps + # cumulative debug logs consistent if DEBUG logging is enabled after metrics + # collection has already been running, while preserving the fully-disabled + # early-exit path above. + totals = self.hicache_l1_l2_transfer_totals[direction] + totals["events"] += 1 + totals["blocks"] += ack.block_count + totals["bytes"] += ack.byte_count + if xfer_us is not None: + totals["xfer_us"] += xfer_us + + if should_log: + logger.debug( + '%s transfer cumulative direction="%s" total_events=%d ' + "total_blocks=%d total_bytes=%d total_xfer_us=%d " + "bandwidth=%.2fGB/s cumulative " + 'src="%s" dst="%s"', + action, + direction, + totals["events"], + totals["blocks"], + totals["bytes"], + totals["xfer_us"], + ( + totals["bytes"] * 0.001 / totals["xfer_us"] + if totals["xfer_us"] > 0 + else 0 + ), + src, + dst, + ) 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 381af58be2a7..3a4d12d3dce9 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -151,6 +151,8 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): prefetch_threshold=prefetch_threshold, load_cache_event=self.load_cache_event, enable_storage_metrics=self.enable_storage_metrics, + enable_metrics=params.enable_metrics, + extra_metric_labels=self.extra_metric_labels, attn_cp_group=params.attn_cp_cache_group, attn_tp_group=params.attn_tp_cache_group, ) @@ -379,19 +381,34 @@ def _inc_hit_count(self, node: TreeNode, chunked=False): # write to host if the node is not backuped self.write_backup(node) + def _finish_write_through_ack(self, ack_id: int, *, release_lock: bool) -> bool: + backuped_node = self.ongoing_write_through.pop(ack_id, None) + if backuped_node is None: + return False + + self._record_store_event(backuped_node, medium=StorageMedium.CPU) + if release_lock: + self.dec_lock_ref(backuped_node) + if self.enable_storage: + self.write_backup_storage(backuped_node) + return True + def writing_check(self, write_back=False): if write_back: # blocking till all write back complete while len(self.ongoing_write_through) > 0: - for _, finish_event, ack_list in self.cache_controller.ack_write_queue: - finish_event.synchronize() - for ack_id in ack_list: - backuped_node = self.ongoing_write_through.pop(ack_id) - self._record_store_event( - backuped_node, medium=StorageMedium.CPU + for ack in self.cache_controller.ack_write_queue: + ack.finish_event.synchronize() + matched = False + for ack_id in ack.node_ids: + matched |= self._finish_write_through_ack( + ack_id, release_lock=False + ) + if matched: + self.cache_controller.record_l1_l2_transfer_complete( + direction="offload", + ack=ack, ) - if self.enable_storage: - self.write_backup_storage(backuped_node) self.cache_controller.ack_write_queue.clear() assert len(self.ongoing_write_through) == 0 return @@ -401,8 +418,8 @@ def writing_check(self, write_back=False): # independently (no cross-rank sync). finish_count = 0 if len(self.ongoing_write_through) > 0: - for _, finish_event, ack_list in self.cache_controller.ack_write_queue: - if not finish_event.query(): + for ack in self.cache_controller.ack_write_queue: + if not ack.finish_event.query(): break finish_count += 1 @@ -416,23 +433,33 @@ def writing_check(self, write_back=False): finish_count = int(queue_size.item()) while finish_count > 0: - _, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0) - finish_event.synchronize() - for ack_id in ack_list: - backuped_node = self.ongoing_write_through.pop(ack_id) - self._record_store_event(backuped_node, medium=StorageMedium.CPU) - self.dec_lock_ref(backuped_node) - if self.enable_storage: - self.write_backup_storage(backuped_node) + ack = self.cache_controller.ack_write_queue.pop(0) + ack.finish_event.synchronize() + + matched = False + for ack_id in ack.node_ids: + matched |= self._finish_write_through_ack(ack_id, release_lock=True) + if matched: + self.cache_controller.record_l1_l2_transfer_complete( + direction="offload", + ack=ack, + ) finish_count -= 1 def loading_check(self): # Every rank must enter the all_reduce below; ongoing_load_back can # diverge across ranks. finish_count = 0 - for _, finish_event, ack_list in self.cache_controller.ack_load_queue: - if not finish_event.query(): + for ack in self.cache_controller.ack_load_queue: + if not ack.finish_event.query(): + # the KV cache loading is still ongoing break + + self.cache_controller.record_l1_l2_transfer_complete( + direction="onboard", + ack=ack, + ) + finish_count += 1 queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu") diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 92e842ad51cb..2d2650835b5f 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -136,6 +136,8 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): extra_config=extra_config, prefetch_threshold=prefetch_threshold, enable_storage_metrics=self.enable_storage_metrics, + enable_metrics=self._enable_metrics_flag, + extra_metric_labels=self.extra_metric_labels, load_cache_event=self.load_cache_event, attn_cp_group=self.attn_cp_group, attn_tp_group=self.attn_tp_group, @@ -157,6 +159,8 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): model_name=server_args.served_model_name, storage_backend_extra_config=extra_config, enable_storage_metrics=self.enable_storage_metrics, + enable_metrics=params.enable_metrics, + extra_metric_labels=self.extra_metric_labels, ) self._apply_storage_runtime_config( storage_backend=server_args.hicache_storage_backend, @@ -797,7 +801,7 @@ def _track_write_through_node(self, node: TreeNode, backup_len: int) -> None: def _replace_pending_write_through_node( self, old_node: TreeNode, new_nodes: List[TreeNode] ) -> None: - ack_id = old_node.write_through_pending_id + ack_id = getattr(old_node, "write_through_pending_id", None) if ack_id is None: return @@ -822,17 +826,21 @@ def _replace_pending_write_through_node( node.write_through_pending_id = ack_id self.ongoing_write_through[ack_id] = (lock_node, backup_len, updated_nodes) - def _finish_write_through_ack(self, ack_id: int, *, release_lock: bool) -> None: - lock_node, backup_len, publish_nodes = self.ongoing_write_through.pop(ack_id) + def _finish_write_through_ack(self, ack_id: int, *, release_lock: bool) -> bool: + pending = self.ongoing_write_through.pop(ack_id, None) + if pending is None: + return False + + lock_node, backup_len, publish_nodes = pending for node in publish_nodes: - if node.write_through_pending_id == ack_id: + if getattr(node, "write_through_pending_id", None) == ack_id: node.write_through_pending_id = None - # DMA confirmed -- block is now on host. self._record_store_event(node, medium=StorageMedium.CPU) if self.enable_storage: self.write_backup_storage(lock_node, backup_len) if release_lock: self.dec_lock_ref(lock_node) + return True def write_backup_storage(self, node: TreeNode, backup_len: Optional[int] = None): # Recover pre-split data via walk-and-concat if node was split. @@ -910,10 +918,18 @@ def writing_check(self, write_back=False): if write_back: # blocking till all write back complete while len(self.ongoing_write_through) > 0: - for _, finish_event, ack_list in self.cache_controller.ack_write_queue: - finish_event.synchronize() - for ack_id in ack_list: - self._finish_write_through_ack(ack_id, release_lock=False) + for ack in self.cache_controller.ack_write_queue: + ack.finish_event.synchronize() + matched = False + for ack_id in ack.node_ids: + matched |= self._finish_write_through_ack( + ack_id, release_lock=False + ) + if matched: + self.cache_controller.record_l1_l2_transfer_complete( + direction="offload", + ack=ack, + ) self.cache_controller.ack_write_queue.clear() assert len(self.ongoing_write_through) == 0 return @@ -924,40 +940,51 @@ def writing_check(self, write_back=False): # sequence and deadlocks under TP > 1. (Matches UnifiedRadixCache.) finish_count = 0 if self.pp_rank == 0: - for _, finish_event, ack_list in self.cache_controller.ack_write_queue: - if not finish_event.query(): + for ack in self.cache_controller.ack_write_queue: + if not ack.finish_event.query(): break finish_count += 1 finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu") self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN) - finish_count = finish_count_tensor.item() + finish_count = int(finish_count_tensor.item()) if finish_count > 0: logger.debug(f"Process {finish_count} write back operations") while finish_count > 0: - _, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0) - finish_event.synchronize() - for ack_id in ack_list: - self._finish_write_through_ack(ack_id, release_lock=True) + ack = self.cache_controller.ack_write_queue.pop(0) + ack.finish_event.synchronize() + + matched = False + for ack_id in ack.node_ids: + matched |= self._finish_write_through_ack(ack_id, release_lock=True) + if matched: + self.cache_controller.record_l1_l2_transfer_complete( + direction="offload", + ack=ack, + ) finish_count -= 1 def loading_check(self): finish_count = 0 if self.pp_rank == 0: - for _, finish_event, ack_list in self.cache_controller.ack_load_queue: - if not finish_event.query(): + for ack in self.cache_controller.ack_load_queue: + if not ack.finish_event.query(): break finish_count += 1 finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu") self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN) - finish_count = finish_count_tensor.item() + finish_count = int(finish_count_tensor.item()) if finish_count > 0: logger.debug(f"Process {finish_count} load operations") while finish_count > 0: - _, finish_event, ack_list = self.cache_controller.ack_load_queue.pop(0) - finish_event.synchronize() - for ack_id in ack_list: + ack = self.cache_controller.ack_load_queue.pop(0) + ack.finish_event.synchronize() + self.cache_controller.record_l1_l2_transfer_complete( + direction="onboard", + ack=ack, + ) + for ack_id in ack.node_ids: end_node = self.ongoing_load_back.pop(ack_id) self.dec_lock_ref(end_node) finish_count -= 1 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 011a877e79ff..40df95c8730e 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 @@ -11,9 +11,6 @@ import torch from sglang.srt.managers.cache_controller import CacheOperation as BaseCacheOperation -from sglang.srt.managers.cache_controller import ( - HiCacheAck, -) from sglang.srt.managers.cache_controller import ( HiCacheController as BaseHiCacheController, ) @@ -172,6 +169,8 @@ def __init__( storage_backend_extra_config: Optional[dict] = None, transfer_layer_num: Optional[int] = None, enable_storage_metrics: bool = False, + enable_metrics: bool = False, + extra_metric_labels: Optional[dict[str, str]] = None, ): startup_storage_backend = storage_backend self.extra_host_mem_release_queues: dict[PoolName, Queue[torch.Tensor]] = {} @@ -191,6 +190,8 @@ def __init__( model_name=model_name, storage_backend_extra_config=storage_backend_extra_config, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) # Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)), # not just the full attention layers reported by full_kv_pool. @@ -408,8 +409,12 @@ def start_writing(self) -> None: self.move_hybrid_indices(op) ) self.write_queue.clear() - start_event = device_module.Event() - finish_event = device_module.Event() + # enable_timing so record_l1_l2_transfer_complete can use elapsed_time() + # for the real device transfer duration (host fallback returns None). + start_event = device_module.Event(enable_timing=True) + finish_event = device_module.Event(enable_timing=True) + + token_count = int(host_indices.numel()) start_event.record() with device_module.stream(self.write_stream): start_event.wait(self.write_stream) @@ -434,7 +439,15 @@ def start_writing(self) -> None: device_indices, resolved_pool_transfers, ) - self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) + self.ack_write_queue.append( + self._make_hicache_ack( + start_event=start_event, + finish_event=finish_event, + node_ids=op.node_ids, + token_count=token_count, + pool_transfers=resolved_pool_transfers, + ) + ) def load( self, @@ -489,6 +502,8 @@ def start_loading(self) -> int: ) self.load_queue.clear() producer_event = self.layer_done_counter.events[producer_id] + + token_count = int(host_indices.numel()) producer_event.start_event.record() with device_module.stream(self.load_stream): producer_event.start_event.wait(self.load_stream) @@ -521,10 +536,12 @@ def start_loading(self) -> int: resolved_pool_transfers, ) self.ack_load_queue.append( - HiCacheAck( - producer_event.start_event, - producer_event.finish_event, - op.node_ids, + self._make_hicache_ack( + start_event=producer_event.start_event, + finish_event=producer_event.finish_event, + node_ids=op.node_ids, + token_count=token_count, + pool_transfers=resolved_pool_transfers, ) ) return producer_id diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index daaf05e63c12..86b9c82da8fd 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -119,6 +119,8 @@ def build_kv_only_stack( model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, enable_storage_metrics: bool = False, + enable_metrics: bool = False, + extra_metric_labels: Optional[dict[str, str]] = None, ) -> tuple[HostPoolGroup, HybridCacheController]: transfer_layer_num = len(full_layer_mapping) kv_host_pool = build_kv_host_pool( @@ -156,6 +158,8 @@ def build_kv_only_stack( storage_backend_extra_config=storage_backend_extra_config, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) return host_pool_group, cache_controller @@ -182,6 +186,8 @@ def build_hybrid_swa_stack( model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, enable_storage_metrics: bool = False, + enable_metrics: bool = False, + extra_metric_labels: Optional[dict[str, str]] = None, ) -> tuple[HostPoolGroup, HybridCacheController]: transfer_layer_num = len(full_layer_mapping | swa_layer_mapping) kv_host_pool = build_kv_host_pool( @@ -238,6 +244,8 @@ def build_hybrid_swa_stack( storage_backend_extra_config=storage_backend_extra_config, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) return host_pool_group, cache_controller @@ -285,6 +293,8 @@ def build_deepseek_v4_hicache_stack( model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, enable_storage_metrics: bool = False, + enable_metrics: bool = False, + extra_metric_labels: Optional[dict[str, str]] = None, ) -> tuple[HostPoolGroup, HybridCacheController]: # TODO(hzh0425): Support PP for deepseek v4 with hicache transfer_layer_num = kvcache.end_layer - kvcache.start_layer @@ -475,6 +485,8 @@ def build_deepseek_v4_hicache_stack( storage_backend_extra_config=storage_backend_extra_config, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) return host_pool_group, cache_controller @@ -501,6 +513,8 @@ def build_hybrid_mamba_stack( model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, enable_storage_metrics: bool = False, + enable_metrics: bool = False, + extra_metric_labels: Optional[dict[str, str]] = None, ) -> tuple[HostPoolGroup, HybridCacheController]: transfer_layer_num = len(full_layer_mapping | mamba_layer_mapping) mamba_allocator = params.req_to_token_pool.mamba_allocator @@ -556,6 +570,8 @@ def build_hybrid_mamba_stack( storage_backend_extra_config=storage_backend_extra_config, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) return host_pool_group, cache_controller @@ -581,6 +597,8 @@ def build_anchor_sidecar_stack( model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, enable_storage_metrics: bool = False, + enable_metrics: bool = False, + extra_metric_labels: Optional[dict[str, str]] = None, ) -> tuple[HostPoolGroup, HybridCacheController]: transfer_layer_num = len(full_layer_mapping) kv_host_pool = build_kv_host_pool( @@ -626,6 +644,8 @@ def build_anchor_sidecar_stack( storage_backend_extra_config=storage_backend_extra_config, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) return host_pool_group, cache_controller @@ -669,6 +689,8 @@ def build( prefetch_threshold: int = 256, model_name: Optional[str] = None, enable_storage_metrics: bool = False, + enable_metrics: bool = False, + extra_metric_labels: Optional[dict[str, str]] = None, ) -> StackBuildResult: raise NotImplementedError @@ -699,6 +721,8 @@ def build( prefetch_threshold=256, model_name=None, enable_storage_metrics=False, + enable_metrics=False, + extra_metric_labels=None, ): from sglang.srt.mem_cache.base_prefix_cache import EvictParams @@ -719,6 +743,8 @@ def build( model_name=model_name, storage_backend_extra_config=storage_backend_extra_config, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) sidecars = [ SidecarPoolSpec( @@ -777,6 +803,8 @@ def build( prefetch_threshold=256, model_name=None, enable_storage_metrics=False, + enable_metrics=False, + extra_metric_labels=None, ): from sglang.srt.mem_cache.base_prefix_cache import EvictParams @@ -803,6 +831,8 @@ def build( model_name=model_name, storage_backend_extra_config=storage_backend_extra_config, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) return StackBuildResult( host_pool_group=host_pool_group, @@ -853,6 +883,8 @@ def build( prefetch_threshold=256, model_name=None, enable_storage_metrics=False, + enable_metrics=False, + extra_metric_labels=None, ): from sglang.srt.mem_cache.base_prefix_cache import EvictParams @@ -878,6 +910,8 @@ def build( model_name=model_name, storage_backend_extra_config=storage_backend_extra_config, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) return StackBuildResult( host_pool_group=host_pool_group, @@ -914,6 +948,8 @@ def build( prefetch_threshold=256, model_name=None, enable_storage_metrics=False, + enable_metrics=False, + extra_metric_labels=None, ): from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool @@ -944,6 +980,8 @@ def build( model_name=model_name, storage_backend_extra_config=storage_backend_extra_config, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) return StackBuildResult( host_pool_group=host_pool_group, @@ -995,6 +1033,8 @@ def build( prefetch_threshold=256, model_name=None, enable_storage_metrics=False, + enable_metrics=False, + extra_metric_labels=None, ): from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool @@ -1018,6 +1058,8 @@ def build( model_name=model_name, storage_backend_extra_config=storage_backend_extra_config, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) return StackBuildResult( host_pool_group=host_pool_group, @@ -1116,6 +1158,8 @@ def attach_hybrid_pool_to_unified_cache( prefetch_threshold=storage_prefetch_threshold, model_name=server_args.served_model_name, enable_storage_metrics=cache._enable_metrics_flag, + enable_metrics=cache._enable_metrics_flag, + extra_metric_labels=cache.extra_metric_labels, ) _apply_stack_result(cache, kvcache, params, result) except Exception: @@ -1131,6 +1175,8 @@ def attach_hybrid_dsa_pool_to_hiradix_cache( extra_config: dict, prefetch_threshold: int, enable_storage_metrics: bool, + enable_metrics: bool = False, + extra_metric_labels: Optional[dict[str, str]] = None, load_cache_event, attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, @@ -1167,6 +1213,8 @@ def attach_hybrid_dsa_pool_to_hiradix_cache( model_name=server_args.served_model_name, storage_backend_extra_config=extra_config, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) radix_cache.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV) radix_cache.token_to_kv_pool_host = host_pool_group @@ -1190,6 +1238,8 @@ def attach_hybrid_pool_to_mamba_cache( prefetch_threshold: int, load_cache_event, enable_storage_metrics: bool = False, + enable_metrics: bool = False, + extra_metric_labels: Optional[dict[str, str]] = None, attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, ) -> None: @@ -1223,6 +1273,8 @@ def attach_hybrid_pool_to_mamba_cache( model_name=server_args.served_model_name, storage_backend_extra_config=extra_config, enable_storage_metrics=enable_storage_metrics, + enable_metrics=enable_metrics, + extra_metric_labels=extra_metric_labels, ) mamba_cache.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV) mamba_cache.mamba_pool_host = host_pool_group.get_pool(PoolName.MAMBA) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index fa26bbee505b..19b38b8e18dc 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -92,6 +92,7 @@ def __init__(self, tree_components: tuple[ComponentType, ...], priority: int = 0 self.hash_value = None self.hit_count = 0 self.priority = priority + self.write_through_pending_id: Optional[int] = None self.lru_prev: list[UnifiedTreeNode | None] = [None] * ( _NUM_COMPONENT_TYPES * 2 ) @@ -100,7 +101,6 @@ def __init__(self, tree_components: tuple[ComponentType, ...], priority: int = 0 ) self.id = UnifiedTreeNode.counter UnifiedTreeNode.counter += 1 - self.write_through_pending_id: Optional[int] = None def component(self, component_type: ComponentType) -> ComponentData: return self.component_data[component_type] @@ -1643,8 +1643,12 @@ def _replace_pending_write_through_node( updated_nodes, ) - def _finish_write_through_ack(self, ack_id: int) -> None: - lock_node, lock_params, publish_nodes = self.ongoing_write_through.pop(ack_id) + def _finish_write_through_ack(self, ack_id: int) -> bool: + pending = self.ongoing_write_through.pop(ack_id, None) + if pending is None: + return False + + lock_node, lock_params, publish_nodes = pending for node in publish_nodes: if node.write_through_pending_id == ack_id: node.write_through_pending_id = None @@ -1656,6 +1660,7 @@ def _finish_write_through_ack(self, ack_id: int) -> None: # suffix; the prefix fragment must be persisted as well. for node in publish_nodes: self.write_backup_storage(node) + return True def load_back( self, @@ -2349,11 +2354,16 @@ def writing_check(self, write_back: bool = False) -> None: if write_back: # Blocking: wait for all pending write-backs while self.ongoing_write_through: - for _, finish_event, ack_list in cc.ack_write_queue: - finish_event.synchronize() - for ack_id in ack_list: - if ack_id in self.ongoing_write_through: - self._finish_write_through_ack(ack_id) + for ack in cc.ack_write_queue: + ack.finish_event.synchronize() + matched = False + for ack_id in ack.node_ids: + matched |= self._finish_write_through_ack(ack_id) + if matched: + cc.record_l1_l2_transfer_complete( + direction="offload", + ack=ack, + ) cc.ack_write_queue.clear() assert len(self.ongoing_write_through) == 0 return @@ -2362,21 +2372,27 @@ def writing_check(self, write_back: bool = False) -> None: # diverge across ranks (e.g. write_backup returning 0 on a subset). finish_count = 0 if self.pp_rank == 0: - for _, finish_event, ack_list in cc.ack_write_queue: - if not finish_event.query(): + for ack in cc.ack_write_queue: + if not ack.finish_event.query(): break finish_count += 1 - finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu") self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN) - finish_count = finish_count_tensor.item() + finish_count = int(finish_count_tensor.item()) # Process completed acks while finish_count > 0: - _, finish_event, ack_list = cc.ack_write_queue.pop(0) - finish_event.synchronize() - for ack_id in ack_list: - self._finish_write_through_ack(ack_id) + ack = cc.ack_write_queue.pop(0) + ack.finish_event.synchronize() + + matched = False + for ack_id in ack.node_ids: + matched |= self._finish_write_through_ack(ack_id) + if matched: + self.cache_controller.record_l1_l2_transfer_complete( + direction="offload", + ack=ack, + ) finish_count -= 1 def loading_check(self) -> None: @@ -2388,18 +2404,22 @@ def loading_check(self) -> None: # diverge across ranks. finish_count = 0 if self.pp_rank == 0: - for _, finish_event, ack_list in cc.ack_load_queue: - if not finish_event.query(): + for ack in cc.ack_load_queue: + if not ack.finish_event.query(): break finish_count += 1 finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu") self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN) - finish_count = finish_count_tensor.item() + finish_count = int(finish_count_tensor.item()) while finish_count > 0: - _, finish_event, ack_list = cc.ack_load_queue.pop(0) - finish_event.synchronize() - for ack_id in ack_list: + ack = cc.ack_load_queue.pop(0) + ack.finish_event.synchronize() + self.cache_controller.record_l1_l2_transfer_complete( + direction="onboard", + ack=ack, + ) + for ack_id in ack.node_ids: node, lock_params, host_lock_params = self.ongoing_load_back.pop(ack_id) self.dec_lock_ref(node, lock_params) self.dec_host_lock_ref(node, host_lock_params) diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index d1a9cca04e28..2e7d170d7b8c 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -1826,6 +1826,96 @@ def log_storage_metrics(self, storage_metrics: Optional[StorageMetrics] = None): self._log_histogram(self.histogram_backup_bandwidth, v) +class HiCacheL1L2TransferMetricsCollector(_StatLoggerDIMixin): + """Prometheus metrics for HiCache L1<->L2 KV block transfers. + + L1: device/GPU KV cache. + L2: host/CPU KV cache. + """ + + _metric_cache: dict[ + tuple[tuple[str, ...], type, type], + tuple[Any, Any, Any], + ] = {} + + def __init__(self, labels: Optional[dict[str, str]] = None): + from prometheus_client import Counter as _PromCounter + from prometheus_client import Histogram as _PromHistogram + + Counter = self._counter_cls or _PromCounter + Histogram = self._histogram_cls or _PromHistogram + + self.labels = labels or {} + labelnames = list(self.labels.keys()) + ["direction", "src", "dst"] + cache_key = (tuple(labelnames), Counter, Histogram) + + cached_metrics = self._metric_cache.get(cache_key) + if cached_metrics is None: + cached_metrics = ( + Counter( + "sglang:hicache_l1_l2_transfer_blocks_total", + "Total number of KV cache blocks transferred between HiCache L1 and L2.", + labelnames=labelnames, + ), + Counter( + "sglang:hicache_l1_l2_transfer_bytes_total", + "Total number of KV cache bytes transferred between HiCache L1 and L2.", + labelnames=labelnames, + ), + Histogram( + "sglang:hicache_l1_l2_transfer_duration_us", + "Observed device-event duration in microseconds for one completed HiCache L1<->L2 KV block transfer.", + labelnames=labelnames, + buckets=( + 100, + 250, + 500, + 1_000, + 2_500, + 5_000, + 10_000, + 25_000, + 50_000, + 100_000, + 250_000, + 500_000, + 1_000_000, + 2_500_000, + 5_000_000, + ), + ), + ) + self._metric_cache[cache_key] = cached_metrics + + ( + self.transfer_blocks_total, + self.transfer_bytes_total, + self.transfer_duration_us, + ) = cached_metrics + + def record_transfer( + self, + *, + direction: str, + src: str, + dst: str, + blocks: int, + bytes_: int, + xfer_us: Optional[int], + ) -> None: + metric_labels = { + **self.labels, + "direction": direction, + "src": src, + "dst": dst, + } + + self.transfer_blocks_total.labels(**metric_labels).inc(blocks) + self.transfer_bytes_total.labels(**metric_labels).inc(bytes_) + if xfer_us is not None: + self.transfer_duration_us.labels(**metric_labels).observe(xfer_us) + + class ExpertDispatchCollector(_StatLoggerDIMixin): def __init__(self, ep_size: int) -> None: from prometheus_client import Histogram as _PromHistogram diff --git a/test/registered/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/disaggregation/test_specv2_kvcache_offloading.py index 296138eaa673..91a4ba6fd8b7 100644 --- a/test/registered/disaggregation/test_specv2_kvcache_offloading.py +++ b/test/registered/disaggregation/test_specv2_kvcache_offloading.py @@ -8,6 +8,7 @@ """ import unittest +from types import SimpleNamespace from unittest.mock import MagicMock import torch @@ -91,6 +92,10 @@ def synchronize(self): pass +def _ack(node_id: int): + return SimpleNamespace(finish_event=_FinishedEvent(), node_ids=[node_id]) + + class TestReleaseFinishedReq(unittest.TestCase): """Tests for _release_finished_req overallocation cleanup.""" @@ -293,7 +298,7 @@ def test_unfinished_offload_ack_does_not_free_incremental_slots(self): 8, ) manager.cache_controller = MagicMock() - manager.cache_controller.ack_write_queue = [(None, _FinishedEvent(), [7])] + manager.cache_controller.ack_write_queue = [_ack(7)] manager._trigger_backup = MagicMock(return_value="last_hash") manager._check_offload_progress(1) @@ -327,7 +332,7 @@ def test_offload_kv_cache_tracks_inflight_write_until_ack(self): self.assertEqual(manager.offloaded_state[req.rid].inc_len, 4) manager.cache_controller.write.assert_called_once() - manager.cache_controller.ack_write_queue = [(None, _FinishedEvent(), [1])] + manager.cache_controller.ack_write_queue = [_ack(1)] manager._trigger_backup = MagicMock(return_value="last_hash") manager._check_offload_progress(1) @@ -370,7 +375,7 @@ def test_finished_offload_ack_waits_for_other_inflight_writes(self): 8, ) manager.cache_controller = MagicMock() - manager.cache_controller.ack_write_queue = [(None, _FinishedEvent(), [8])] + manager.cache_controller.ack_write_queue = [_ack(8)] manager._trigger_backup = MagicMock(return_value="last_hash") manager._check_offload_progress(1) @@ -400,7 +405,7 @@ def test_finished_request_releases_all_committed_slots_after_last_offload_ack( 12, ) manager.cache_controller = MagicMock() - manager.cache_controller.ack_write_queue = [(None, _FinishedEvent(), [9])] + manager.cache_controller.ack_write_queue = [_ack(9)] manager._trigger_backup = MagicMock(return_value="last_hash") manager._check_offload_progress(1) diff --git a/test/registered/unit/mem_cache/test_hicache_l1l2_metrics.py b/test/registered/unit/mem_cache/test_hicache_l1l2_metrics.py new file mode 100644 index 000000000000..e06e7af9c13d --- /dev/null +++ b/test/registered/unit/mem_cache/test_hicache_l1l2_metrics.py @@ -0,0 +1,317 @@ +"""Unit tests for HiCache L1/L2 transfer metrics. + +All tests run on CPU; no GPU device or distributed process group is required. +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + +import unittest +from unittest.mock import MagicMock + +from sglang.srt.managers.cache_controller import HiCacheAck, HiCacheController +from sglang.srt.observability.metrics_collector import ( + HiCacheL1L2TransferMetricsCollector, +) + +# ── Test doubles ────────────────────────────────────────────────────────────── + + +class _RecordingCounter: + """prometheus_client.Counter stand-in that records every call.""" + + def __init__(self, *args, **kwargs): + self.inc_calls: list = [] + + def labels(self, **kwargs): + return self + + def inc(self, amount=1): + self.inc_calls.append(amount) + + +class _RecordingHistogram(_RecordingCounter): + """prometheus_client.Histogram stand-in.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.observe_calls: list = [] + + def observe(self, value): + self.observe_calls.append(value) + + +# ── Controller stub ─────────────────────────────────────────────────────────── + + +class _ControllerStub: + """Minimal stand-in that lets HiCacheController unbound methods run on CPU.""" + + def __init__(self): + self.hicache_l1_l2_transfer_metrics_collector = None + self.hicache_l1_l2_transfer_totals = { + "offload": {"events": 0, "blocks": 0, "bytes": 0, "xfer_us": 0}, + "onboard": {"events": 0, "blocks": 0, "bytes": 0, "xfer_us": 0}, + } + + _transfer_elapsed_us = HiCacheController._transfer_elapsed_us + record_l1_l2_transfer_complete = HiCacheController.record_l1_l2_transfer_complete + + +# ── Ack factory ─────────────────────────────────────────────────────────────── + + +def _make_ack(*, block_count=4, byte_count=65536) -> HiCacheAck: + return HiCacheAck( + start_event=MagicMock(), + finish_event=MagicMock(), + node_ids=[1, 2], + token_count=64, + block_count=block_count, + byte_count=byte_count, + ) + + +# ══════════════════════════════════════════════════════════════════════════════ +# HiCacheL1L2TransferMetricsCollector — call-chain (DI doubles) +# ══════════════════════════════════════════════════════════════════════════════ + + +class TestHiCacheL1L2TransferCollectorRecordTransfer(unittest.TestCase): + class _Collector(HiCacheL1L2TransferMetricsCollector): + _counter_cls = _RecordingCounter + _histogram_cls = _RecordingHistogram + + def setUp(self): + HiCacheL1L2TransferMetricsCollector._metric_cache.clear() + + def tearDown(self): + HiCacheL1L2TransferMetricsCollector._metric_cache.clear() + + def _make_collector(self): + return self._Collector(labels={"tp_rank": "0", "io_backend": "test"}) + + def test_record_offload_increments_blocks_counter(self): + c = self._make_collector() + c.record_transfer( + direction="offload", + src="sglang_hicache::L1", + dst="sglang_hicache::L2", + blocks=8, + bytes_=131072, + xfer_us=None, + ) + self.assertEqual(c.transfer_blocks_total.inc_calls, [8]) + + def test_record_with_xfer_us_observes_histogram(self): + c = self._make_collector() + c.record_transfer( + direction="offload", + src="sglang_hicache::L1", + dst="sglang_hicache::L2", + blocks=2, + bytes_=32768, + xfer_us=5000, + ) + self.assertEqual(c.transfer_duration_us.observe_calls, [5000]) + + +# ══════════════════════════════════════════════════════════════════════════════ +# HiCacheController.record_l1_l2_transfer_complete +# ══════════════════════════════════════════════════════════════════════════════ + + +class TestRecordL1L2TransferComplete(unittest.TestCase): + def _make_ctrl(self) -> _ControllerStub: + ctrl = _ControllerStub() + ctrl.hicache_l1_l2_transfer_metrics_collector = MagicMock() + ctrl._transfer_elapsed_us = MagicMock(return_value=2000) + return ctrl + + def test_offload_totals_accumulate_over_two_calls(self): + ctrl = self._make_ctrl() + ack = _make_ack(block_count=4, byte_count=65536) + ctrl.record_l1_l2_transfer_complete(direction="offload", ack=ack) + ctrl.record_l1_l2_transfer_complete(direction="offload", ack=ack) + + totals = ctrl.hicache_l1_l2_transfer_totals["offload"] + self.assertEqual(totals["events"], 2) + self.assertEqual(totals["blocks"], 8) + self.assertEqual(totals["bytes"], 131072) + + def test_collector_called_with_correct_args_offload(self): + ctrl = self._make_ctrl() + ctrl._transfer_elapsed_us = MagicMock(return_value=5000) + ack = _make_ack(block_count=3, byte_count=49152) + ctrl.record_l1_l2_transfer_complete(direction="offload", ack=ack) + + ctrl.hicache_l1_l2_transfer_metrics_collector.record_transfer.assert_called_once_with( + direction="offload", + src="sglang_hicache::L1", + dst="sglang_hicache::L2", + blocks=3, + bytes_=49152, + xfer_us=5000, + ) + + def test_collector_called_with_correct_args_onboard(self): + ctrl = self._make_ctrl() + ctrl._transfer_elapsed_us = MagicMock(return_value=3000) + ack = _make_ack(block_count=6, byte_count=98304) + ctrl.record_l1_l2_transfer_complete(direction="onboard", ack=ack) + + ctrl.hicache_l1_l2_transfer_metrics_collector.record_transfer.assert_called_once_with( + direction="onboard", + src="sglang_hicache::L2", + dst="sglang_hicache::L1", + blocks=6, + bytes_=98304, + xfer_us=3000, + ) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Prometheus text exposition — real prometheus_client, isolated registry +# ══════════════════════════════════════════════════════════════════════════════ + + +class TestHiCacheL1L2TransferCollectorPrometheusOutput(unittest.TestCase): + """Verifies metric names, label values, and sample values in real Prometheus + text output using an isolated CollectorRegistry per test.""" + + def setUp(self): + import prometheus_client + from prometheus_client import CollectorRegistry + + self.registry = CollectorRegistry() + _reg = self.registry + _PC = prometheus_client.Counter + _PH = prometheus_client.Histogram + + # Each setUp() defines new class objects, giving a unique _metric_cache + # key so every test gets fresh Counter/Histogram instances. + class _BoundCounter: + def __init__(self, *args, **kwargs): + kwargs["registry"] = _reg + self._inner = _PC(*args, **kwargs) + + def labels(self, **kwargs): + return self._inner.labels(**kwargs) + + class _BoundHistogram: + def __init__(self, *args, **kwargs): + kwargs["registry"] = _reg + self._inner = _PH(*args, **kwargs) + + def labels(self, **kwargs): + return self._inner.labels(**kwargs) + + class _TestCollector(HiCacheL1L2TransferMetricsCollector): + _counter_cls = _BoundCounter + _histogram_cls = _BoundHistogram + + self._TestCollector = _TestCollector + HiCacheL1L2TransferMetricsCollector._metric_cache.clear() + + def tearDown(self): + HiCacheL1L2TransferMetricsCollector._metric_cache.clear() + + def _make_collector(self): + return self._TestCollector(labels={"tp_rank": "0", "io_backend": "nixl"}) + + def _exposition(self) -> str: + from prometheus_client import generate_latest + + return generate_latest(self.registry).decode("utf-8") + + def _find_metric_line( + self, exposition: str, name_fragment: str, **label_filters + ) -> str: + """Return the first non-comment line matching name_fragment and all + label_filters, or fail the test.""" + for line in exposition.splitlines(): + if line.startswith("#"): + continue + if name_fragment not in line: + continue + if all(f'{k}="{v}"' in line for k, v in label_filters.items()): + return line + self.fail( + f"No metric line for fragment={name_fragment!r} " + f"labels={label_filters} in:\n{exposition}" + ) + + def test_blocks_counter_value(self): + c = self._make_collector() + c.record_transfer( + direction="offload", + src="sglang_hicache::L1", + dst="sglang_hicache::L2", + blocks=7, + bytes_=114688, + xfer_us=None, + ) + line = self._find_metric_line( + self._exposition(), + "hicache_l1_l2_transfer_blocks_total", + direction="offload", + ) + self.assertTrue(line.endswith(" 7.0"), f"Expected value 7.0 in: {line!r}") + + def test_bytes_counter_value(self): + c = self._make_collector() + c.record_transfer( + direction="onboard", + src="sglang_hicache::L2", + dst="sglang_hicache::L1", + blocks=1, + bytes_=32768, + xfer_us=None, + ) + line = self._find_metric_line( + self._exposition(), + "hicache_l1_l2_transfer_bytes_total", + direction="onboard", + ) + self.assertTrue( + line.endswith(" 32768.0"), f"Expected value 32768.0 in: {line!r}" + ) + + def test_duration_histogram_sum_when_xfer_us_provided(self): + c = self._make_collector() + c.record_transfer( + direction="offload", + src="sglang_hicache::L1", + dst="sglang_hicache::L2", + blocks=1, + bytes_=1024, + xfer_us=2500, + ) + line = self._find_metric_line( + self._exposition(), + "hicache_l1_l2_transfer_duration_us_sum", + direction="offload", + ) + self.assertTrue(line.endswith(" 2500.0"), f"Expected sum 2500.0 in: {line!r}") + + def test_duration_histogram_count_when_xfer_us_provided(self): + c = self._make_collector() + c.record_transfer( + direction="offload", + src="sglang_hicache::L1", + dst="sglang_hicache::L2", + blocks=1, + bytes_=1024, + xfer_us=1000, + ) + line = self._find_metric_line( + self._exposition(), + "hicache_l1_l2_transfer_duration_us_count", + direction="offload", + ) + self.assertTrue(line.endswith(" 1.0"), f"Expected count 1.0 in: {line!r}") + + +if __name__ == "__main__": + unittest.main() 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 67bcf88f504b..17d6fbc360ab 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 @@ -497,8 +497,8 @@ def _load_back_node(self, tree, node): self.assertTrue(loaded) producer_id = tree.ready_to_load_host_cache() self.assertNotEqual(producer_id, -1) - for _, finish_event, _ in list(tree.cache_controller.ack_load_queue): - finish_event.synchronize() + for ack in list(tree.cache_controller.ack_load_queue): + ack.finish_event.synchronize() tree.loading_check() def test_kv_events_store_and_remove_full_blocks(self): @@ -2792,8 +2792,8 @@ def _load_back_node(self, tree, node): self.assertTrue(loaded) producer_id = tree.ready_to_load_host_cache() self.assertNotEqual(producer_id, -1) - for _, finish_event, _ in list(tree.cache_controller.ack_load_queue): - finish_event.synchronize() + for ack in list(tree.cache_controller.ack_load_queue): + ack.finish_event.synchronize() tree.loading_check() return node.component_data[ComponentType.FULL].value @@ -3294,8 +3294,8 @@ def _release_ongoing_load_back_locks(self, tree): def _finish_pending_loads(self, tree): producer_id = tree.ready_to_load_host_cache() self.assertNotEqual(producer_id, -1) - for _, finish_event, _ in list(tree.cache_controller.ack_load_queue): - finish_event.synchronize() + for ack in list(tree.cache_controller.ack_load_queue): + ack.finish_event.synchronize() tree.loading_check() def _match_tokens_for_chain(self, chain):