[KV Offload] Split cpu_cache_usage_perc into write/read usage gauges (#47666)
Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Srinivas Krovvidi committed
Jul 14, 2026 at 12:13 UTC
cdaa40d2a8c9297349669fc6e7db6dcb282df0b8
4 files changed
+80
tests/v1/kv_offload/cpu/test_manager.py
+49
@@ -92,6 +92,21 @@ def verify_load_output(
92
assert np.array_equal(expected_array, prepare_load_output.block_ids)
93
94
95
+def check_split_usage_stats(
96
+ manager: CPUOffloadingManager, write: float, read: float, total: float
97
+):
98
+ stats = manager.get_stats()
99
+ assert stats is not None
100
+ reduced = stats.reduce()
101
+ assert reduced[CPUOffloadingMetrics.CPU_CACHE_WRITE_USAGE_PERC] == pytest.approx(
102
+ write
103
+ )
104
+ assert reduced[CPUOffloadingMetrics.CPU_CACHE_READ_USAGE_PERC] == pytest.approx(
105
+ read
106
+ )
107
+ assert reduced[CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC] == pytest.approx(total)
108
+
109
+
110
def verify_events(
111
events: Iterable[OffloadingEvent],
112
expected_stores: tuple[set[int], ...] = (),
@@ -302,6 +317,40 @@ def test_cpu_manager_reports_allocation_size_on_eviction_failure():
317
assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 1
318
319
320
+def test_cpu_manager_reports_cache_write_and_read_usage_gauges():
321
+ manager = make_cpu_manager(num_blocks=4)
322
+
323
+ # Store path: pins write usage until complete_store.
324
+ manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
325
+ check_split_usage_stats(manager, write=0.5, read=0.0, total=0.5)
326
+
327
+ manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
328
+ check_split_usage_stats(manager, write=0.0, read=0.0, total=0.0)
329
+
330
+ # Load path: pins read usage until complete_load.
331
+ assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT
332
+ manager.prepare_load(to_keys([1]), _EMPTY_REQ_CTX)
333
+ check_split_usage_stats(manager, write=0.0, read=0.25, total=0.25)
334
+
335
+ manager.complete_load(to_keys([1]), _EMPTY_REQ_CTX)
336
+ check_split_usage_stats(manager, write=0.0, read=0.0, total=0.0)
337
+
338
+ # Concurrent write + read pins are both reflected and additive.
339
+ manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX)
340
+ manager.prepare_load(to_keys([2]), _EMPTY_REQ_CTX)
341
+ check_split_usage_stats(manager, write=0.5, read=0.25, total=0.75)
342
+
343
+
344
+def test_cpu_manager_clears_write_usage_after_failed_store():
345
+ manager = make_cpu_manager(num_blocks=4)
346
+
347
+ manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
348
+ check_split_usage_stats(manager, write=0.5, read=0.0, total=0.5)
349
+
350
+ manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX, success=False)
351
+ check_split_usage_stats(manager, write=0.0, read=0.0, total=0.0)
352
+
353
+
354
def test_cpu_manager():
355
"""
356
Tests CPUOffloadingManager with lru policy.
vllm/v1/kv_offload/cpu/common.py
+2
@@ -7,6 +7,8 @@ class CPUOffloadingMetrics:
7
STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
8
CPU_CACHE_USAGE_PERC = "vllm:kv_offload_cpu_cache_usage_perc"
9
CPU_ALLOCATION_SIZE = "vllm:kv_offload_cpu_allocation_size"
10
+ CPU_CACHE_WRITE_USAGE_PERC = "vllm:kv_offload_cpu_cache_write_usage_perc"
11
+ CPU_CACHE_READ_USAGE_PERC = "vllm:kv_offload_cpu_cache_read_usage_perc"
12
13
14
class CPULoadStoreSpec(BlockIDsLoadStoreSpec):
vllm/v1/kv_offload/cpu/manager.py
+15
@@ -66,6 +66,8 @@ class CPUOffloadingManager(OffloadingManager):
66
self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks)
67
# Track the number of blocks in the cache that are evictable. i.e. ref_cnt 0.
68
self._num_evictable_cache_blocks: int = 0
69
+ # Track blocks with an in-flight store (ref_cnt -1, not yet completed).
70
+ self._num_write_pending_blocks: int = 0
71
72
self.store_threshold: int = store_threshold
73
self.max_tracker_size: int = max_tracker_size
@@ -229,6 +231,7 @@ class CPUOffloadingManager(OffloadingManager):
231
232
for key, block in zip(keys_to_store, blocks):
233
self._policy.insert(key, block)
234
+ self._num_write_pending_blocks += len(keys_to_store)
235
236
# build store specs for allocated blocks
237
store_spec = self._get_load_store_spec(keys_to_store, blocks)
@@ -253,6 +256,7 @@ class CPUOffloadingManager(OffloadingManager):
256
block = self._policy.get(key)
257
if block is not None and not block.is_ready:
258
block.ref_cnt = 0
259
+ self._num_write_pending_blocks -= 1
260
self._num_evictable_cache_blocks += 1
261
self._policy.mark_evictable(key)
262
stored_keys.append(key)
@@ -260,6 +264,7 @@ class CPUOffloadingManager(OffloadingManager):
264
for key in keys:
265
block = self._policy.get(key)
266
if block is not None and not block.is_ready:
267
+ self._num_write_pending_blocks -= 1
268
self._policy.remove(key)
269
self._free_block(block)
270
@@ -281,6 +286,7 @@ class CPUOffloadingManager(OffloadingManager):
286
# can begin, preventing a cross-direction data race on reused offload block IDs.
287
self._policy.clear()
288
self._num_evictable_cache_blocks = 0
289
+ self._num_write_pending_blocks = 0
290
291
self._free_list.clear()
292
self._num_allocated_blocks = 0
@@ -309,6 +315,15 @@ class CPUOffloadingManager(OffloadingManager):
315
)
316
self.allocation_sizes_in_current_batch.clear()
317
318
+ write_usage = (
319
+ self._num_write_pending_blocks / self._num_blocks
320
+ if self._num_blocks > 0
321
+ else 0.0
322
+ )
323
+ read_usage = max(usage - write_usage, 0.0)
324
+ stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_WRITE_USAGE_PERC, write_usage)
325
+ stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_READ_USAGE_PERC, read_usage)
326
+
327
if self.store_threshold >= 2:
328
stats.increase_counter(
329
CPUOffloadingMetrics.STORES_SKIPPED,
vllm/v1/kv_offload/cpu/spec.py
+14
@@ -39,6 +39,20 @@ class CPUOffloadingSpec(OffloadingSpec):
39
"dropped due to insufficient capacity."
40
),
41
),
42
+ CPUOffloadingMetrics.CPU_CACHE_WRITE_USAGE_PERC: OffloadingGaugeMetadata(
43
+ documentation=(
44
+ "Fraction of CPU KV-cache space currently pinned by "
45
+ "in-flight stores that have not yet "
46
+ "completed (0.0 = idle, 1.0 = saturated)."
47
+ ),
48
+ ),
49
+ CPUOffloadingMetrics.CPU_CACHE_READ_USAGE_PERC: OffloadingGaugeMetadata(
50
+ documentation=(
51
+ "Fraction of CPU KV-cache space currently pinned by "
52
+ "in-flight loads that have not yet "
53
+ "completed (0.0 = idle, 1.0 = saturated)."
54
+ ),
55
+ ),
56
CPUOffloadingMetrics.CPU_ALLOCATION_SIZE: OffloadingHistogramMetadata(
57
documentation=(
58
"Histogram of the number of CPU blocks requested by each "