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 3, 22:45 UTC
[kv_offload] Move `FilterReusedOffloadingManager` logic to `CPUOffloadingManager` (#41727)
Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Martin Hickey committed
May 11, 2026 at 16:09 UTC
4955990f1bd91e5d6070a56ad3ed9684a23e5146
4 files changed
+37
-148
tests/v1/kv_offload/cpu/test_manager.py
+7
-8
@@ -17,7 +17,6 @@ from vllm.v1.kv_offload.base import (
17
from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
18
from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
19
from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
20
-from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
20
21
22
def make_req_context(kv_transfer_params: dict | None = None) -> ReqContext:
@@ -565,14 +564,14 @@ class TestARCPolicy:
564
565
def test_filter_reused_manager():
566
"""
568
- Tests FilterReusedOffloadingManager with a CPUOffloadingManager.
567
+ Tests CPUOffloadingManager reuse filtering (store_threshold=2).
568
"""
570
- lru_manager = CPUOffloadingManager(
571
- num_blocks=4, cache_policy="lru", enable_events=True
572
- )
573
-
574
- manager = FilterReusedOffloadingManager(
575
- backing=lru_manager, store_threshold=2, max_tracker_size=3
569
+ manager = CPUOffloadingManager(
570
+ num_blocks=4,
571
+ cache_policy="lru",
572
+ enable_events=True,
573
+ store_threshold=2,
574
+ max_tracker_size=3,
575
)
576
577
# Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet
vllm/v1/kv_offload/cpu/manager.py
+20
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections import OrderedDict
4
from collections.abc import Collection, Iterable
5
from typing import Literal
6
@@ -37,6 +38,8 @@ class CPUOffloadingManager(OffloadingManager):
38
num_blocks: int,
39
cache_policy: Literal["lru", "arc"] = "lru",
40
enable_events: bool = False,
41
+ store_threshold: int = 1,
42
+ max_tracker_size: int = 64_000,
43
):
44
self.medium: str = CPULoadStoreSpec.medium()
45
self._num_blocks: int = num_blocks
@@ -50,6 +53,13 @@ class CPUOffloadingManager(OffloadingManager):
53
f"Supported: {list(_CACHE_POLICIES)}"
54
)
55
self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks)
56
+ self.store_threshold: int = store_threshold
57
+ self.max_tracker_size: int = max_tracker_size
58
+
59
+ # Number of block references. It is ordered so can evict the LRU entry in O(1).
60
+ self.counts: OrderedDict[OffloadKey, int] | None = (
61
+ OrderedDict() if store_threshold >= 2 else None
62
+ )
63
64
# --- block pool ---
65
@@ -85,6 +95,14 @@ class CPUOffloadingManager(OffloadingManager):
95
# --- OffloadingManager interface ---
96
97
def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
98
+ if self.counts is not None:
99
+ if key in self.counts:
100
+ self.counts.move_to_end(key)
101
+ self.counts[key] += 1
102
+ else:
103
+ if len(self.counts) >= self.max_tracker_size:
104
+ self.counts.popitem(last=False)
105
+ self.counts[key] = 1
106
block = self._policy.get(key)
107
if block is None:
108
return False
@@ -123,6 +141,8 @@ class CPUOffloadingManager(OffloadingManager):
141
keys: Collection[OffloadKey],
142
req_context: ReqContext,
143
) -> PrepareStoreOutput | None:
144
+ if self.counts is not None:
145
+ keys = [k for k in keys if self.counts.get(k, 0) >= self.store_threshold]
146
# filter out blocks that are already stored
147
keys_to_store = [k for k in keys if self._policy.get(k) is None]
148
vllm/v1/kv_offload/cpu/spec.py
+10
-15
@@ -15,7 +15,6 @@ from vllm.v1.kv_offload.base import (
15
from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
16
from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers
17
from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
18
-from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
18
from vllm.v1.kv_offload.worker.worker import OffloadingHandler
19
20
@@ -61,25 +60,21 @@ class CPUOffloadingSpec(OffloadingSpec):
60
kv_events_config is not None and kv_events_config.enable_kv_cache_events
61
)
62
63
+ # store_threshold: how many times a block must appear in lookup()
64
+ # before it is eligible for CPU offloading. Values < 2 disable
65
+ # filtering (a threshold of 1 equals no filter; 0 is the default).
66
+ store_threshold = int(self.extra_config.get("store_threshold", 0))
67
+
68
+ # Maximum entries in the internal tracker's LRU table.
69
+ max_tracker_size = int(self.extra_config.get("max_tracker_size", 64_000))
70
+
71
self._manager = CPUOffloadingManager(
72
num_blocks=self.num_blocks,
73
cache_policy=self.eviction_policy, # type: ignore[arg-type]
74
enable_events=enable_events,
75
+ store_threshold=store_threshold,
76
+ max_tracker_size=max_tracker_size,
77
)
69
-
70
- # store_threshold: how many times a block must appear in lookup()
71
- # before it is eligible for CPU offloading. Values < 2 disable
72
- # filtering (a threshold of 1 equals no filter; 0 is the default).
73
- store_threshold = int(self.extra_config.get("store_threshold", 0))
74
- if store_threshold >= 2:
75
- max_tracker_size = int(
76
- self.extra_config.get("max_tracker_size", 64_000)
77
- )
78
- self._manager = FilterReusedOffloadingManager(
79
- backing=self._manager,
80
- store_threshold=store_threshold,
81
- max_tracker_size=max_tracker_size,
82
- )
78
return self._manager
79
80
def get_handlers(
vllm/v1/kv_offload/reuse_manager.py
deleted
-125
@@ -1,125 +0,0 @@
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 Collection, Iterable
12
-
13
-from vllm.v1.kv_offload.base import (
14
- LoadStoreSpec,
15
- OffloadingEvent,
16
- OffloadingManager,
17
- OffloadKey,
18
- PrepareStoreOutput,
19
- ReqContext,
20
-)
21
-
22
-
23
-class FilterReusedOffloadingManager(OffloadingManager):
24
- """An :class:`OffloadingManager` decorator that skips storing blocks
25
- whose reuse frequency is below *store_threshold*.
26
-
27
- All methods are delegated to the *backing* manager. Two methods are
28
- intercepted:
29
-
30
- * ``prepare_store`` — filters out keys that have not yet
31
- * ``lookup`` — records the visited key in an internal LRU
32
- counter, then delegates to the backing manager.
33
- crossed the threshold *before* calling the backing
34
- ``prepare_store``.
35
-
36
- Args:
37
- backing: The underlying ``OffloadingManager`` to delegate to.
38
- store_threshold: A block must be seen at least this many times in
39
- ``lookup()`` before it is eligible for offloading. Must be >= 2
40
- (a value of 1 would be equivalent to no filtering).
41
- max_tracker_size: Maximum entries in the internal tracker's LRU table.
42
- """
43
-
44
- def __init__(
45
- self,
46
- backing: OffloadingManager,
47
- store_threshold: int = 2,
48
- max_tracker_size: int = 64_000,
49
- ):
50
- if store_threshold < 2:
51
- raise ValueError(
52
- "FilterReusedOffloadingManager store_threshold must be >= 2, "
53
- f"got {store_threshold}"
54
- )
55
- if max_tracker_size < 1:
56
- raise ValueError(
57
- "FilterReusedOffloadingManager max_tracker_size must be >= 1, "
58
- f"got {max_tracker_size}"
59
- )
60
- self._backing = backing
61
- self.store_threshold = store_threshold
62
- self.max_tracker_size = max_tracker_size
63
- # Ordered so we can evict the LRU entry in O(1).
64
- self.counts: OrderedDict[OffloadKey, int] = OrderedDict()
65
-
66
- # ------------------------------------------------------------------
67
- # Intercepted methods
68
- # ------------------------------------------------------------------
69
-
70
- def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
71
- """Record the key, then delegate lookup to backing manager."""
72
- if key in self.counts:
73
- self.counts.move_to_end(key)
74
- self.counts[key] += 1
75
- else:
76
- if len(self.counts) >= self.max_tracker_size:
77
- self.counts.popitem(last=False) # evict LRU
78
- self.counts[key] = 1
79
- return self._backing.lookup(key, req_context)
80
-
81
- def prepare_store(
82
- self, keys: Collection[OffloadKey], req_context: ReqContext
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
- eligible = [
91
- key for key in keys if self.counts.get(key, 0) >= self.store_threshold
92
- ]
93
-
94
- # Passing an empty list is intentional and safe — CPUOffloadingManager
95
- # handles it correctly, returning a PrepareStoreOutput with empty lists.
96
- # Delegate to the backing manager with only the eligible keys.
97
- return self._backing.prepare_store(eligible, req_context)
98
-
99
- # ------------------------------------------------------------------
100
- # Delegated methods
101
- # ------------------------------------------------------------------
102
-
103
- def prepare_load(
104
- self, keys: Collection[OffloadKey], req_context: ReqContext
105
- ) -> LoadStoreSpec:
106
- return self._backing.prepare_load(keys, req_context)
107
-
108
- def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:
109
- return self._backing.touch(keys, req_context)
110
-
111
- def complete_load(
112
- self, keys: Collection[OffloadKey], req_context: ReqContext
113
- ) -> None:
114
- return self._backing.complete_load(keys, req_context)
115
-
116
- def complete_store(
117
- self,
118
- keys: Collection[OffloadKey],
119
- req_context: ReqContext,
120
- success: bool = True,
121
- ) -> None:
122
- return self._backing.complete_store(keys, req_context, success)
123
-
124
- def take_events(self) -> Iterable[OffloadingEvent]:
125
- return self._backing.take_events()