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 0388e9a0c..5c3b3d1ce 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 @@ -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