Fix mooncake store write/read bandwidth logs (#18294)

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,

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

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,

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]

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: