From 726a72080efc14e117a2703f47404cc8333ef551 Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Sun, 5 Jul 2026 19:18:47 -0500 Subject: [PATCH] Split cpu_cache_usage_perc into write/read usage gauges Report separate gauges for CPU cache space pinned by in-flight stores (GPU-to-CPU) versus in-flight loads (CPU-to-GPU), in addition to the existing combined usage gauge. This makes it possible to tell whether cache pressure is coming from writes or reads. Signed-off-by: srinivas_oo7 --- tests/v1/kv_offload/cpu/test_manager.py | 49 +++++++++++++++++++++++++ vllm/v1/kv_offload/cpu/common.py | 2 + vllm/v1/kv_offload/cpu/manager.py | 15 ++++++++ vllm/v1/kv_offload/cpu/spec.py | 14 +++++++ 4 files changed, 80 insertions(+) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index e16d6d378ea2..97ed5168f744 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -92,6 +92,21 @@ def verify_load_output( assert np.array_equal(expected_array, prepare_load_output.block_ids) +def check_split_usage_stats( + manager: CPUOffloadingManager, write: float, read: float, total: float +): + stats = manager.get_stats() + assert stats is not None + reduced = stats.reduce() + assert reduced[CPUOffloadingMetrics.CPU_CACHE_WRITE_USAGE_PERC] == pytest.approx( + write + ) + assert reduced[CPUOffloadingMetrics.CPU_CACHE_READ_USAGE_PERC] == pytest.approx( + read + ) + assert reduced[CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC] == pytest.approx(total) + + def verify_events( events: Iterable[OffloadingEvent], expected_stores: tuple[set[int], ...] = (), @@ -302,6 +317,40 @@ def test_cpu_manager_reports_allocation_size_on_eviction_failure(): assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 1 +def test_cpu_manager_reports_cache_write_and_read_usage_gauges(): + manager = make_cpu_manager(num_blocks=4) + + # Store path: pins write usage until complete_store. + manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) + check_split_usage_stats(manager, write=0.5, read=0.0, total=0.5) + + manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) + check_split_usage_stats(manager, write=0.0, read=0.0, total=0.0) + + # Load path: pins read usage until complete_load. + assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT + manager.prepare_load(to_keys([1]), _EMPTY_REQ_CTX) + check_split_usage_stats(manager, write=0.0, read=0.25, total=0.25) + + manager.complete_load(to_keys([1]), _EMPTY_REQ_CTX) + check_split_usage_stats(manager, write=0.0, read=0.0, total=0.0) + + # Concurrent write + read pins are both reflected and additive. + manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX) + manager.prepare_load(to_keys([2]), _EMPTY_REQ_CTX) + check_split_usage_stats(manager, write=0.5, read=0.25, total=0.75) + + +def test_cpu_manager_clears_write_usage_after_failed_store(): + manager = make_cpu_manager(num_blocks=4) + + manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) + check_split_usage_stats(manager, write=0.5, read=0.0, total=0.5) + + manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX, success=False) + check_split_usage_stats(manager, write=0.0, read=0.0, total=0.0) + + def test_cpu_manager(): """ Tests CPUOffloadingManager with lru policy. diff --git a/vllm/v1/kv_offload/cpu/common.py b/vllm/v1/kv_offload/cpu/common.py index 29f95e1b4076..7875d81501b2 100644 --- a/vllm/v1/kv_offload/cpu/common.py +++ b/vllm/v1/kv_offload/cpu/common.py @@ -7,6 +7,8 @@ class CPUOffloadingMetrics: STORES_SKIPPED = "vllm:kv_offload_stores_skipped" CPU_CACHE_USAGE_PERC = "vllm:kv_offload_cpu_cache_usage_perc" CPU_ALLOCATION_SIZE = "vllm:kv_offload_cpu_allocation_size" + CPU_CACHE_WRITE_USAGE_PERC = "vllm:kv_offload_cpu_cache_write_usage_perc" + CPU_CACHE_READ_USAGE_PERC = "vllm:kv_offload_cpu_cache_read_usage_perc" class CPULoadStoreSpec(BlockIDsLoadStoreSpec): diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index e7416bf63bd2..297c94aabb08 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -66,6 +66,8 @@ def __init__( self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks) # Track the number of blocks in the cache that are evictable. i.e. ref_cnt 0. self._num_evictable_cache_blocks: int = 0 + # Track blocks with an in-flight store (ref_cnt -1, not yet completed). + self._num_write_pending_blocks: int = 0 self.store_threshold: int = store_threshold self.max_tracker_size: int = max_tracker_size @@ -229,6 +231,7 @@ def prepare_store( for key, block in zip(keys_to_store, blocks): self._policy.insert(key, block) + self._num_write_pending_blocks += len(keys_to_store) # build store specs for allocated blocks store_spec = self._get_load_store_spec(keys_to_store, blocks) @@ -253,6 +256,7 @@ def complete_store( block = self._policy.get(key) if block is not None and not block.is_ready: block.ref_cnt = 0 + self._num_write_pending_blocks -= 1 self._num_evictable_cache_blocks += 1 self._policy.mark_evictable(key) stored_keys.append(key) @@ -260,6 +264,7 @@ def complete_store( for key in keys: block = self._policy.get(key) if block is not None and not block.is_ready: + self._num_write_pending_blocks -= 1 self._policy.remove(key) self._free_block(block) @@ -281,6 +286,7 @@ def reset_cache(self) -> None: # can begin, preventing a cross-direction data race on reused offload block IDs. self._policy.clear() self._num_evictable_cache_blocks = 0 + self._num_write_pending_blocks = 0 self._free_list.clear() self._num_allocated_blocks = 0 @@ -309,6 +315,15 @@ def get_stats(self) -> OffloadingConnectorStats | None: ) self.allocation_sizes_in_current_batch.clear() + write_usage = ( + self._num_write_pending_blocks / self._num_blocks + if self._num_blocks > 0 + else 0.0 + ) + read_usage = max(usage - write_usage, 0.0) + stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_WRITE_USAGE_PERC, write_usage) + stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_READ_USAGE_PERC, read_usage) + if self.store_threshold >= 2: stats.increase_counter( CPUOffloadingMetrics.STORES_SKIPPED, diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 26ea3728191f..4ad6974857cf 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -39,6 +39,20 @@ def build_metric_definitions( "dropped due to insufficient capacity." ), ), + CPUOffloadingMetrics.CPU_CACHE_WRITE_USAGE_PERC: OffloadingGaugeMetadata( + documentation=( + "Fraction of CPU KV-cache space currently pinned by " + "in-flight stores that have not yet " + "completed (0.0 = idle, 1.0 = saturated)." + ), + ), + CPUOffloadingMetrics.CPU_CACHE_READ_USAGE_PERC: OffloadingGaugeMetadata( + documentation=( + "Fraction of CPU KV-cache space currently pinned by " + "in-flight loads that have not yet " + "completed (0.0 = idle, 1.0 = saturated)." + ), + ), CPUOffloadingMetrics.CPU_ALLOCATION_SIZE: OffloadingHistogramMetadata( documentation=( "Histogram of the number of CPU blocks requested by each "