| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """KV cache metrics tracking.""" |
| 4 | |
| 5 | import random |
| 6 | import time |
| 7 | from collections import deque |
| 8 | from typing import TYPE_CHECKING |
| 9 | |
| 10 | if TYPE_CHECKING: |
| 11 | from vllm.v1.core.kv_cache_utils import KVCacheBlock |
| 12 | |
| 13 | from vllm.v1.metrics.stats import KVCacheEvictionEvent |
| 14 | |
| 15 | |
| 16 | class BlockMetricsState: |
| 17 | """Tracks lifecycle metrics for a single KV cache block.""" |
| 18 | |
| 19 | def __init__(self): |
| 20 | now_ns = time.monotonic_ns() |
| 21 | self.birth_time_ns = now_ns |
| 22 | self.last_access_ns = now_ns |
| 23 | # Bounded to prevent unbounded growth if a block is accessed many times. |
| 24 | self.access_history: deque[int] = deque(maxlen=4) |
| 25 | |
| 26 | def record_access(self) -> None: |
| 27 | now_ns = time.monotonic_ns() |
| 28 | self.last_access_ns = now_ns |
| 29 | self.access_history.append(now_ns) |
| 30 | |
| 31 | def get_lifetime_seconds(self) -> float: |
| 32 | now_ns = time.monotonic_ns() |
| 33 | return (now_ns - self.birth_time_ns) / 1e9 |
| 34 | |
| 35 | def get_idle_time_seconds(self) -> float: |
| 36 | now_ns = time.monotonic_ns() |
| 37 | return (now_ns - self.last_access_ns) / 1e9 |
| 38 | |
| 39 | def get_reuse_gaps_seconds(self) -> list[float]: |
| 40 | if len(self.access_history) < 2: |
| 41 | return [] |
| 42 | history = list(self.access_history) |
| 43 | return [(history[i] - history[i - 1]) / 1e9 for i in range(1, len(history))] |
| 44 | |
| 45 | |
| 46 | class KVCacheMetricsCollector: |
| 47 | """Collects KV cache residency metrics with sampling.""" |
| 48 | |
| 49 | def __init__(self, sample_rate: float = 0.01): |
| 50 | assert 0 < sample_rate <= 1.0, ( |
| 51 | f"sample_rate must be in (0, 1.0], got {sample_rate}" |
| 52 | ) |
| 53 | self.sample_rate = sample_rate |
| 54 | |
| 55 | self.block_metrics: dict[int, BlockMetricsState] = {} |
| 56 | |
| 57 | self._eviction_events: list[KVCacheEvictionEvent] = [] |
| 58 | |
| 59 | def should_sample_block(self) -> bool: |
| 60 | return random.random() < self.sample_rate |
| 61 | |
| 62 | def on_block_allocated(self, block: "KVCacheBlock") -> None: |
| 63 | if self.should_sample_block(): |
| 64 | self.block_metrics[block.block_id] = BlockMetricsState() |
| 65 | |
| 66 | def on_block_accessed(self, block: "KVCacheBlock") -> None: |
| 67 | metrics = self.block_metrics.get(block.block_id) |
| 68 | if metrics: |
| 69 | metrics.record_access() |
| 70 | |
| 71 | def on_block_evicted(self, block: "KVCacheBlock") -> None: |
| 72 | metrics = self.block_metrics.pop(block.block_id, None) |
| 73 | if not metrics: |
| 74 | return |
| 75 | |
| 76 | lifetime = metrics.get_lifetime_seconds() |
| 77 | idle_time = metrics.get_idle_time_seconds() |
| 78 | reuse_gaps = tuple(metrics.get_reuse_gaps_seconds()) |
| 79 | |
| 80 | self._eviction_events.append( |
| 81 | KVCacheEvictionEvent( |
| 82 | lifetime_seconds=lifetime, |
| 83 | idle_seconds=idle_time, |
| 84 | reuse_gaps_seconds=reuse_gaps, |
| 85 | ) |
| 86 | ) |
| 87 | |
| 88 | def reset(self) -> None: |
| 89 | """Clear all state on cache reset.""" |
| 90 | self.block_metrics.clear() |
| 91 | self._eviction_events.clear() |
| 92 | |
| 93 | def drain_events(self) -> list[KVCacheEvictionEvent]: |
| 94 | events = self._eviction_events |
| 95 | self._eviction_events = [] |
| 96 | return events |