Add mooncake store read/write bandwidth logs (#10598)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
Co-authored-by: Lei Huang <lei.huang@novita.ai>

python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py

@@ -17,6 +17,7 @@ from sglang.srt.mem_cache.hicache_storage import (
     HiCacheStorageExtraInfo,
 )
 from sglang.srt.mem_cache.memory_pool_host import HostKVCache, HostTensorAllocator
+from sglang.srt.metrics.collector import StorageMetrics
 
 DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024  # 16 MB
 SETUP_TIMEOUT = 600  # 10min
@@ -346,6 +347,12 @@ class MooncakeStore(HiCacheStorage):
                 self.mha_suffix = f"{self.local_rank}"
                 self.mla_suffix = ""
 
+            self.gb_per_page = None
+            self.prefetch_pgs = []
+            self.backup_pgs = []
+            self.prefetch_bandwidth = []
+            self.backup_bandwidth = []
+
         except ValueError as e:
             logger.error("Configuration loading failed: %s", e)
             raise
@@ -414,6 +421,9 @@ class MooncakeStore(HiCacheStorage):
             logger.error("Failed to register buffer to Mooncake Store: %s", err)
             raise TypeError("Mooncake Store Register Buffer Error.") from err
 
+        bytes_per_page = mem_pool_host.get_ksize_per_token() * mem_pool_host.page_size
+        self.gb_per_page = bytes_per_page / (1 << 30)
+
     def _get_mha_buffer_meta(self, keys, indices):
         ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices)
         key_list = []
@@ -566,9 +576,17 @@ class MooncakeStore(HiCacheStorage):
                 set_target_sizes.append(target_sizes[i])
                 set_indices.append(i)
         # Only set non-existing keys to storage
+        start_time = time.perf_counter()
         put_result = self._put_batch_zero_copy_impl(
             set_keys, set_target_locations, set_target_sizes
         )
+        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
+        )
+
         for i in range(len(set_indices)):
             if put_result[i] == 0:
                 exist_result[set_indices[i]] = 1
@@ -602,13 +620,23 @@ class MooncakeStore(HiCacheStorage):
         assert len(keys) == len(target_locations) == len(target_sizes)
         if len(keys) == 0:
             return 0
+
+        start_time = time.perf_counter()
         get_result = self._get_batch_zero_copy_impl(
             keys, target_locations, target_sizes
         )
+        end_time = time.perf_counter()
+
         if self.is_mla_backend:
             key_multiplier = 1
         else:
             key_multiplier = 2
+
+        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:
                 return i // key_multiplier
@@ -657,3 +685,15 @@ class MooncakeStore(HiCacheStorage):
 
     def _batch_exist(self, key_strs: List[str]) -> List[int]:
         return self.store.batch_is_exist(key_strs)
+
+    def get_stats(self):
+        storage_metrics = StorageMetrics()
+        storage_metrics.prefetch_pgs.extend(self.prefetch_pgs)
+        storage_metrics.backup_pgs.extend(self.backup_pgs)
+        storage_metrics.prefetch_bandwidth.extend(self.prefetch_bandwidth)
+        storage_metrics.backup_bandwidth.extend(self.backup_bandwidth)
+        self.prefetch_pgs.clear()
+        self.backup_pgs.clear()
+        self.prefetch_bandwidth.clear()
+        self.backup_bandwidth.clear()
+        return storage_metrics