test

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

Author: comaniac

tests/v1/core/test_kv_cache_utils.py

@@ -5,7 +5,7 @@
 from vllm.multimodal.inputs import MultiModalKwargs
 from vllm.sampling_params import SamplingParams
 from vllm.v1.core.kv_cache_utils import (BlockHashType, FreeKVCacheBlockQueue,
-                                         KVCacheBlock,
+                                         KVCacheBlock, PrefixCachingMetrics,
                                          generate_block_hash_extra_keys,
                                          hash_block_tokens,
                                          hash_request_tokens)
@@ -277,3 +277,34 @@ def test_hash_request_tokens_no_mm_inputs():
     assert block_hashes[0].extra_keys is None
     assert block_hashes[1].token_ids == (3, 4, 5)
     assert block_hashes[1].extra_keys is None
+
+
+def test_metrics():
+    """
+    Test the prefix caching metrics.
+    """
+    metrics = PrefixCachingMetrics(interval=5)
+    assert metrics.hit_rate == 0.0
+
+    metrics.add_request_query(20, 9)
+    # 9 / 20 = 0.45
+    assert metrics.hit_rate == 0.45
+
+    for _ in range(4):
+        metrics.add_request_query(20, 4)
+
+    # 25 / 100 = 0.25
+    assert metrics.hit_rate == 0.25
+
+    metrics.add_request_query(10, 2)
+
+    # Remove (20, 9) and add (10, 2): 18 / 90 = 0.2
+    assert metrics.aggregated_query_total == 90
+    assert metrics.aggregated_query_hit == 18
+    assert metrics.hit_rate == 0.2
+
+    metrics.reset()
+    assert metrics.hit_rate == 0.0
+    assert metrics.aggregated_query_total == 0
+    assert metrics.aggregated_query_hit == 0
+    assert not metrics.query_queue

vllm/v1/core/kv_cache_manager.py

@@ -72,8 +72,8 @@ def __init__(
         self.req_to_blocks: DefaultDict[str,
                                         List[KVCacheBlock]] = defaultdict(list)
 
-        # Prefix cache metrics.
-        self.prefix_caching_metrics = PrefixCachingMetrics()
+        # Prefix cache metrics. TODO: Make the interval configurable.
+        self.prefix_caching_metrics = PrefixCachingMetrics(interval=1000)
 
     @property
     def usage(self) -> float:

vllm/v1/core/kv_cache_utils.py

@@ -39,9 +39,11 @@ class PrefixCachingMetrics:
 
     def __init__(self, interval: int = 1000):
         self.interval = interval
+        # The current aggregated query total and hit.
         self.aggregated_query_total = 0
         self.aggregated_query_hit = 0
-        self.request_queries: deque[Tuple[int, int]] = deque()
+        # A deque of (num_queries, num_hits) for the most recent requests.
+        self.query_queue: deque[Tuple[int, int]] = deque()
 
     def add_request_query(self, num_queries: int, num_hits: int):
         """Add a request to the metrics. This function is called when
@@ -54,9 +56,9 @@ def add_request_query(self, num_queries: int, num_hits: int):
             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.query_queue.append((num_queries, num_hits))
+        if len(self.query_queue) > self.interval:
+            old_num_queries, old_num_hits = self.query_queue.popleft()
             self.aggregated_query_total -= old_num_queries
             self.aggregated_query_hit -= old_num_hits
 
@@ -67,7 +69,7 @@ def reset(self):
         """Reset the metrics."""
         self.aggregated_query_total = 0
         self.aggregated_query_hit = 0
-        self.request_queries.clear()
+        self.query_queue.clear()
 
     @property
     def hit_rate(self) -> float: