improve

Signed-off-by: Cody Yu <hao.yu.cody@gmail.com>

Author: comaniac

vllm/v1/core/kv_cache_manager.py

@@ -73,10 +73,7 @@ def __init__(
  List[KVCacheBlock]] = defaultdict(list)
 
  # Prefix cache metrics.
- self.prefix_caching_metrics: PrefixCachingMetrics = {
- "query_total": 0,
- "query_hit": 0,
- }
+ self.prefix_caching_metrics = PrefixCachingMetrics()
 
  @property
  def usage(self) -> float:
@@ -88,21 +85,14 @@ def usage(self) -> float:
  return 1.0 - (self.free_block_queue.num_free_blocks /
  self.num_gpu_blocks)
 
- def get_and_reset_prefix_cache_hit_rate(self) -> float:
-  """Get the overall hit rate of prefix caching and reset
- the metrics.
+ @property
+ def prefix_cache_hit_rate(self) -> float:
+ """Get the prefix caching hit rate.
 
  Returns:
- The hit rate of prefix caching (between 0.0 and 1.0).
+ The prefix caching hit rate.
  """
- hit_rate = 0.0
- if self.prefix_caching_metrics["query_total"] > 0:
- hit_rate = self.prefix_caching_metrics[
- "query_hit"] / self.prefix_caching_metrics["query_total"]
-
- self.prefix_caching_metrics["query_hit"] = 0
- self.prefix_caching_metrics["query_total"] = 0
- return hit_rate
+ return self.prefix_caching_metrics.hit_rate
 
  def get_computed_blocks(
  self, request: Request) -> Tuple[List[KVCacheBlock], int]:
@@ -139,8 +129,10 @@ def get_computed_blocks(
  else:
  break
 
- self.prefix_caching_metrics["query_total"] += len(block_hashes)
- self.prefix_caching_metrics["query_hit"] += len(computed_blocks)
+ self.prefix_caching_metrics.add_request_query(
+ num_queries=len(block_hashes),
+ num_hits=len(computed_blocks),
+ )
 
  # NOTE(woosuk): Since incomplete blocks are not eligible for
  # sharing, `num_computed_tokens` is always a multiple of
@@ -306,6 +298,9 @@ def reset_prefix_cache(self) -> bool:
  for block in self.block_pool:
  block.reset_hash()
 
+ # Reset the prefix caching metrics.
+ self.prefix_caching_metrics.reset()
+
  logger.info("Successfully reset prefix cache")
  return True
 

vllm/v1/core/kv_cache_utils.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 """KV-Cache Utilities."""
+from collections import deque
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any, List, NamedTuple, Optional, Tuple, TypedDict
+from typing import Any, List, NamedTuple, Optional, Tuple
 
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
@@ -28,14 +29,52 @@ class BlockHashType(NamedTuple):
  extra_keys: Optional[Any] = None
 
 
-class PrefixCachingMetrics(TypedDict):
- """Metrics for prefix caching."""
+class PrefixCachingMetrics:
+ """Metrics for prefix caching with a hit rate of the most recent N requests.
 
- query_total: int
- """The total number of queries."""
+ Args:
+ interval: The number of the most recent requests to aggregate.
+ Defaults to 1000.
+ """
+
+ def __init__(self, interval: int = 1000):
+ self.interval = interval
+ self.aggregated_query_total = 0
+ self.aggregated_query_hit = 0
+ self.request_queries: deque[Tuple[int, int]] = deque()
 
- query_hit: int
- """The number of queries that hit the prefix cache."""
+ def add_request_query(self, num_queries: int, num_hits: int):
+ """Add a request to the metrics. This function is called when
+ a new request is being scheduled and is looking for computed blocks.
+ When there are more than `interval` requests, the oldest request
+ is removed from the metrics.
+
+ Args:
+ num_queries: The number of queries in the request.
+ num_hits: The number of hits in the request.
+ """
+
+ self.request_queries.append((num_queries, num_hits))
+ if len(self.request_queries) > self.interval:
+ old_num_queries, old_num_hits = self.request_queries.popleft()
+ self.aggregated_query_total -= old_num_queries
+ self.aggregated_query_hit -= old_num_hits
+
+ self.aggregated_query_total += num_queries
+ self.aggregated_query_hit += num_hits
+
+ def reset(self):
+ """Reset the metrics."""
+ self.aggregated_query_total = 0
+ self.aggregated_query_hit = 0
+ self.request_queries.clear()
+
+ @property
+ def hit_rate(self) -> float:
+ """Calculate the hit rate for the past N requests."""
+ if self.aggregated_query_total == 0:
+ return 0.0
+ return self.aggregated_query_hit / self.aggregated_query_total
 
 
 @dataclass

vllm/v1/core/scheduler.py

@@ -554,7 +554,7 @@ def make_stats(self) -> SchedulerStats:
  num_waiting_reqs=len(self.waiting),
  gpu_cache_usage=self.kv_cache_manager.usage,
  gpu_prefix_cache_hit_rate=self.kv_cache_manager.
- get_and_reset_prefix_cache_hit_rate(),
+ prefix_cache_hit_rate,
  )