A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 23:00 UTC
feat(kv-offload): Strategy A — StoreReusedOffloadingManager gates CPU stores on reuse frequency (#35342)
Signed-off-by: srinivas_oo7 <Sriusa4414@gmail.com> Signed-off-by: Sriusa4414@gmail.com Signed-off-by: Srinivasoo7 <158864704+Srinivasoo7@users.noreply.github.com> Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: Srinivasoo7 <158864704+Srinivasoo7@users.noreply.github.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Srinivasoo7 committed
Mar 10, 2026 at 09:43 UTC
106ff69c4eb4921d33341a96b9c3d6db9d12ba76
3 files changed
+184
tests/v1/kv_offload/test_cpu_manager.py
+49
@@ -544,3 +544,52 @@ def test_arc_manager_full_scenario():
544
# verify events
545
events = list(arc_manager.take_events())
546
assert len(events) > 0 # should have store and eviction events
547
+
548
+
549
+def test_filter_reused_manager():
550
+ """
551
+ Tests FilterReusedOffloadingManager with a CPUBackend.
552
+ """
553
+ block_size = 256
554
+ cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
555
+ lru_manager = LRUOffloadingManager(cpu_backend, enable_events=True)
556
+
557
+ from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
558
+
559
+ manager = FilterReusedOffloadingManager(
560
+ backing=lru_manager, store_threshold=2, max_tracker_size=3
561
+ )
562
+
563
+ # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet
564
+ assert manager.lookup(to_hashes([1, 2])) == 0
565
+
566
+ # prepare store [1, 2] -> should be filtered
567
+ prepare_store_output = manager.prepare_store(to_hashes([1, 2]))
568
+ assert prepare_store_output is not None
569
+ assert prepare_store_output.block_hashes_to_store == []
570
+
571
+ # Lookup [1] -> 2nd time, eligible now
572
+ assert manager.lookup(to_hashes([1])) == 0
573
+
574
+ # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered
575
+ prepare_store_output = manager.prepare_store(to_hashes([1, 2]))
576
+ assert prepare_store_output is not None
577
+ assert prepare_store_output.block_hashes_to_store == to_hashes([1])
578
+
579
+ # Lookup [3, 4] -> 1st time
580
+ # (evicts [2] from tracker since max_size is 3 and tracker has [1])
581
+ assert manager.lookup(to_hashes([3, 4])) == 0
582
+ # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4])
583
+ assert to_hashes([2])[0] not in manager.counts
584
+
585
+ # Lookup [2] again -> (this adds [2] back to the tracker as 1st time)
586
+ assert manager.lookup(to_hashes([2])) == 0
587
+ # Verify [2] was re-added with count=1 (not eligible yet)
588
+ assert manager.counts.get(to_hashes([2])[0]) == 1
589
+
590
+ # prepare store [2] -> should still be filtered out since count was reset
591
+ prepare_store_output = manager.prepare_store(to_hashes([2]))
592
+ assert prepare_store_output is not None
593
+ assert prepare_store_output.block_hashes_to_store == []
594
+
595
+ manager.complete_store(to_hashes([1]))
vllm/v1/kv_offload/cpu.py
+15
@@ -13,6 +13,7 @@ from vllm.v1.kv_offload.arc_manager import ARCOffloadingManager
13
from vllm.v1.kv_offload.backends.cpu import CPUBackend
14
from vllm.v1.kv_offload.lru_manager import LRUOffloadingManager
15
from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
16
+from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
17
from vllm.v1.kv_offload.spec import OffloadingSpec
18
from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers
19
from vllm.v1.kv_offload.worker.worker import OffloadingHandler
@@ -83,6 +84,20 @@ class CPUOffloadingSpec(OffloadingSpec):
84
f"Unknown eviction policy: {self.eviction_policy}. "
85
f"Supported policies: lru, arc"
86
)
87
+
88
+ # store_threshold: how many times a block must appear in lookup()
89
+ # before it is eligible for CPU offloading. Values < 2 disable
90
+ # filtering (a threshold of 1 equals no filter; 0 is the default).
91
+ store_threshold = int(self.extra_config.get("store_threshold", 0))
92
+ if store_threshold >= 2:
93
+ max_tracker_size = int(
94
+ self.extra_config.get("max_tracker_size", 64_000)
95
+ )
96
+ self._manager = FilterReusedOffloadingManager(
97
+ backing=self._manager,
98
+ store_threshold=store_threshold,
99
+ max_tracker_size=max_tracker_size,
100
+ )
101
return self._manager
102
103
def get_handlers(
vllm/v1/kv_offload/reuse_manager.py
new
+120
@@ -0,0 +1,120 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Reuse-frequency gating for CPU KV-cache offload stores.
5
+
6
+FilterReusedOffloadingManager — OffloadingManager decorator that skips
7
+ storing blocks that have not yet been seen enough times.
8
+"""
9
+
10
+from collections import OrderedDict
11
+from collections.abc import Iterable
12
+
13
+from vllm.v1.core.kv_cache_utils import BlockHash
14
+from vllm.v1.kv_offload.abstract import (
15
+ LoadStoreSpec,
16
+ OffloadingEvent,
17
+ OffloadingManager,
18
+ PrepareStoreOutput,
19
+)
20
+
21
+
22
+class FilterReusedOffloadingManager(OffloadingManager):
23
+ """An :class:`OffloadingManager` decorator that skips storing blocks
24
+ whose reuse frequency is below *store_threshold*.
25
+
26
+ All methods are delegated to the *backing* manager. Two methods are
27
+ intercepted:
28
+
29
+ * ``lookup`` — records each visited block hash in an internal LRU counter.
30
+ * ``prepare_store`` — filters out block hashes that have not yet
31
+ crossed the threshold *before* calling the backing
32
+ ``prepare_store``.
33
+
34
+ Args:
35
+ backing: The underlying ``OffloadingManager`` to delegate to.
36
+ store_threshold: A block must be seen at least this many times in
37
+ ``lookup()`` before it is eligible for offloading. Must be >= 2
38
+ (a value of 1 would be equivalent to no filtering).
39
+ max_tracker_size: Maximum entries in the internal tracker's LRU table.
40
+ """
41
+
42
+ def __init__(
43
+ self,
44
+ backing: OffloadingManager,
45
+ store_threshold: int = 2,
46
+ max_tracker_size: int = 64_000,
47
+ ):
48
+ if store_threshold < 2:
49
+ raise ValueError(
50
+ "FilterReusedOffloadingManager store_threshold must be >= 2, "
51
+ f"got {store_threshold}"
52
+ )
53
+ if max_tracker_size < 1:
54
+ raise ValueError(
55
+ "FilterReusedOffloadingManager max_tracker_size must be >= 1, "
56
+ f"got {max_tracker_size}"
57
+ )
58
+ self._backing = backing
59
+ self.store_threshold = store_threshold
60
+ self.max_tracker_size = max_tracker_size
61
+ # Ordered so we can evict the LRU entry in O(1).
62
+ self.counts: OrderedDict[BlockHash, int] = OrderedDict()
63
+
64
+ # ------------------------------------------------------------------
65
+ # Intercepted methods
66
+ # ------------------------------------------------------------------
67
+
68
+ def lookup(self, block_hashes: Iterable[BlockHash]) -> int | None:
69
+ """Record each hash, then delegate lookup to backing manager."""
70
+ block_hashes = list(block_hashes)
71
+ for block_hash in block_hashes:
72
+ if block_hash in self.counts:
73
+ self.counts.move_to_end(block_hash)
74
+ self.counts[block_hash] += 1
75
+ else:
76
+ if len(self.counts) >= self.max_tracker_size:
77
+ self.counts.popitem(last=False) # evict LRU
78
+ self.counts[block_hash] = 1
79
+ return self._backing.lookup(block_hashes)
80
+
81
+ def prepare_store(
82
+ self, block_hashes: Iterable[BlockHash]
83
+ ) -> PrepareStoreOutput | None:
84
+ """Filter out blocks below threshold, then delegate to backing.
85
+
86
+ Filtering is evaluated *before* calling the backing manager's
87
+ ``prepare_store`` so that blocks that would be skipped do not
88
+ consume any CPU offload capacity.
89
+ """
90
+ block_hashes = list(block_hashes)
91
+ eligible = [
92
+ bh for bh in block_hashes if self.counts.get(bh, 0) >= self.store_threshold
93
+ ]
94
+
95
+ # Delegate to the backing manager with only the eligible hashes.
96
+ # Passing an empty list is intentional and safe — both
97
+ # LRUOffloadingManager and ARCOffloadingManager handle it correctly,
98
+ # returning a PrepareStoreOutput with empty lists.
99
+ return self._backing.prepare_store(eligible)
100
+
101
+ # ------------------------------------------------------------------
102
+ # Delegated methods
103
+ # ------------------------------------------------------------------
104
+
105
+ def prepare_load(self, block_hashes: Iterable[BlockHash]) -> LoadStoreSpec:
106
+ return self._backing.prepare_load(block_hashes)
107
+
108
+ def touch(self, block_hashes: Iterable[BlockHash]) -> None:
109
+ return self._backing.touch(block_hashes)
110
+
111
+ def complete_load(self, block_hashes: Iterable[BlockHash]) -> None:
112
+ return self._backing.complete_load(block_hashes)
113
+
114
+ def complete_store(
115
+ self, block_hashes: Iterable[BlockHash], success: bool = True
116
+ ) -> None:
117
+ return self._backing.complete_store(block_hashes, success)
118
+
119
+ def take_events(self) -> Iterable[OffloadingEvent]:
120
+ return self._backing.take_events()