diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index ff3376ca9..0fedbcf3d 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -261,6 +261,7 @@ class HiCacheController: storage_backend_extra_config: Optional[dict] = None, pp_rank: int = 0, pp_size: int = 1, + enable_storage_metrics: bool = False, ): self.tp_group = tp_group self.mem_pool_device_allocator = token_to_kv_pool_allocator @@ -279,6 +280,7 @@ class HiCacheController: self.storage_backend_type = None self.pp_rank = pp_rank self.pp_size = pp_size + self.enable_storage_metrics = enable_storage_metrics # Default storage page IO functions (may be overridden by attach). self.page_get_func = self._generic_page_get @@ -610,6 +612,7 @@ class HiCacheController: pp_rank=self.pp_rank, pp_size=self.pp_size, is_mla_model=is_mla_backend, + enable_storage_metrics=self.enable_storage_metrics, is_page_first_layout=self.mem_pool_host.layout == "page_first", model_name=model_name, tp_lcm_size=tp_lcm_size, diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 02680db04..e76cfebd4 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -51,6 +51,7 @@ class HiCacheStorageConfig: pp_rank: int pp_size: int is_mla_model: bool + enable_storage_metrics: bool is_page_first_layout: bool model_name: Optional[str] tp_lcm_size: Optional[int] = None diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index e614b7743..041318bf2 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -134,6 +134,7 @@ class HiRadixCache(RadixCache): storage_backend_extra_config=extra_config, pp_rank=self.pp_rank, pp_size=self.pp_size, + enable_storage_metrics=self.enable_storage_metrics, ) self._apply_storage_runtime_config( storage_backend=server_args.hicache_storage_backend, diff --git a/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py b/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py index b20296842..9aa82892d 100644 --- a/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py +++ b/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py @@ -170,6 +170,7 @@ class HiCacheHF3FS(HiCacheStorage): is_mla_model: bool = False, is_page_first_layout: bool = False, use_mock_client: bool = False, + enable_storage_metrics: bool = False, ): self.rank = rank self.file_path = file_path @@ -183,6 +184,7 @@ class HiCacheHF3FS(HiCacheStorage): self.metadata_client = metadata_client self.is_mla_model = is_mla_model self.is_page_first_layout = is_page_first_layout + self.enable_storage_metrics = enable_storage_metrics self.numel = self.bytes_per_page // self.dtype.itemsize self.num_pages = self.file_size // self.bytes_per_page self.skip_backup = False @@ -339,6 +341,7 @@ class HiCacheHF3FS(HiCacheStorage): is_mla_model=is_mla_model, is_page_first_layout=is_page_first_layout, use_mock_client=use_mock_client, + enable_storage_metrics=storage_config.enable_storage_metrics, ) def _batch_get( @@ -376,10 +379,12 @@ class HiCacheHF3FS(HiCacheStorage): end_time = time.perf_counter() ionum = len(batch_indices) - self.prefetch_pgs.append(ionum) - self.prefetch_bandwidth.append( - ionum / (end_time - start_time) * self.gb_per_page - ) + + if self.enable_storage_metrics: + self.prefetch_pgs.append(ionum) + self.prefetch_bandwidth.append( + ionum / (end_time - start_time) * self.gb_per_page + ) results = [False] * len(keys) for batch_index, read_result in zip(batch_indices, read_results): @@ -446,8 +451,12 @@ class HiCacheHF3FS(HiCacheStorage): end_time = time.perf_counter() ionum = len(batch_indices) - self.backup_pgs.append(ionum) - self.backup_bandwidth.append(ionum / (end_time - start_time) * self.gb_per_page) + + if self.enable_storage_metrics: + self.backup_pgs.append(ionum) + self.backup_bandwidth.append( + ionum / (end_time - start_time) * self.gb_per_page + ) written_keys_to_confirm = [] results = [index[0] for index in indices] diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 892ab0d46..2c815fd7e 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -388,11 +388,13 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): self.warmup() logger.info("Mooncake store warmup successfully.") + self.enable_storage_metrics = False if storage_config is not None: self.is_mla_backend = storage_config.is_mla_model self.local_rank = storage_config.tp_rank self.pp_rank = storage_config.pp_rank self.pp_size = storage_config.pp_size + self.enable_storage_metrics = storage_config.enable_storage_metrics else: self.is_mla_backend = False self.local_rank = 0 @@ -581,9 +583,19 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): keys = [f"{prefix}_{key}" for key in keys] key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices) + + start_time = time.perf_counter() get_results = self._get_batch_zero_copy_impl( key_strs, buffer_ptrs, buffer_sizes ) + end_time = time.perf_counter() + + if self.enable_storage_metrics: + self.prefetch_pgs.append(len(keys)) + self.prefetch_bandwidth.append( + len(keys) / (end_time - start_time) * self.gb_per_page + ) + return self._batch_postprocess(get_results, is_set_operate=False) def batch_set_v1( @@ -616,9 +628,18 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): # Only set non-existing keys to storage if len(set_keys) > 0: + start_time = time.perf_counter() put_results = self._put_batch_zero_copy_impl( set_keys, set_buffer_ptrs, set_buffer_sizes ) + end_time = time.perf_counter() + + if self.enable_storage_metrics: + self.backup_pgs.append(len(set_keys)) + self.backup_bandwidth.append( + len(set_keys) / (end_time - start_time) * self.gb_per_page + ) + for i in range(len(set_indices)): set_results[set_indices[i]] = put_results[i] @@ -681,10 +702,11 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): ) end_time = time.perf_counter() - self.backup_pgs.append(len(keys)) - self.backup_bandwidth.append( - len(keys) / (end_time - start_time) * self.gb_per_page - ) + if self.enable_storage_metrics: + self.backup_pgs.append(len(set_keys)) + self.backup_bandwidth.append( + len(set_keys) / (end_time - start_time) * self.gb_per_page + ) for i in range(len(set_indices)): if put_result[i] == 0: @@ -731,10 +753,11 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): else: key_multiplier = 2 - self.prefetch_pgs.append(len(keys)) - self.prefetch_bandwidth.append( - len(keys) / (end_time - start_time) * self.gb_per_page - ) + if self.enable_storage_metrics: + self.prefetch_pgs.append(len(keys)) + self.prefetch_bandwidth.append( + len(keys) / (end_time - start_time) * self.gb_per_page + ) for i in range(len(keys)): if get_result[i] < 0: