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 5, 03:30 UTC
[kv_offload+HMA][3/N]: Remove block_size from KVEvents (#36644)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Apr 15, 2026 at 11:53 UTC
235e1f930a29f491e06fa89c2536576a71922d73
6 files changed
+10
-38
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+3
-7
@@ -124,12 +124,8 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
124
return [BlockHash(str(i).encode()) for i in int_hashes]
125
126
def take_events() -> Iterable[OffloadingEvent]:
127
- yield OffloadingEvent(
128
- keys=to_keys([1, 2, 3]), block_size=16, medium="A", removed=False
129
- )
130
- yield OffloadingEvent(
131
- keys=to_keys([4, 5, 6]), block_size=32, medium="B", removed=True
132
- )
127
+ yield OffloadingEvent(keys=to_keys([1, 2, 3]), medium="A", removed=False)
128
+ yield OffloadingEvent(keys=to_keys([4, 5, 6]), medium="B", removed=True)
129
130
runner.manager.take_events.side_effect = take_events
131
events = list(runner.scheduler_connector.take_events())
@@ -137,7 +133,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
133
event = events[0]
134
assert isinstance(event, BlockStored)
135
assert event.block_hashes == to_hashes([1, 2, 3])
140
- assert event.block_size == 16
136
+ assert event.block_size == 0
137
assert event.medium == "A"
138
assert event.token_ids == []
139
assert event.parent_block_hash is None
tests/v1/kv_offload/test_cpu_manager.py
+6
-20
@@ -59,7 +59,6 @@ def verify_load_output(
59
60
def verify_events(
61
events: Iterable[OffloadingEvent],
62
- block_size: int,
62
expected_stores: tuple[set[int], ...] = (),
63
expected_evictions: tuple[set[int], ...] = (),
64
):
@@ -67,7 +66,6 @@ def verify_events(
66
evictions: list[set[OffloadKey]] = []
67
for event in events:
68
assert event.medium == CPULoadStoreSpec.medium()
70
- assert event.block_size == block_size
69
if event.removed:
70
evictions.append(set(event.keys))
71
else:
@@ -98,9 +96,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
96
candidate to make room for [3, 4, 5]
97
- After complete_store([2, 3, 4, 5]), block 2 must still be present.
98
"""
101
- block_size = 256
99
manager = CPUOffloadingManager(
103
- block_size=block_size,
100
num_blocks=4,
101
cache_policy=eviction_policy,
102
enable_events=True,
@@ -138,10 +134,9 @@ def test_cpu_manager():
134
"""
135
Tests CPUOffloadingManager with lru policy.
136
"""
141
- # initialize a CPU backend with a capacity of 4 blocks
142
- block_size = 256
137
+ # initialize a CPU manager with a capacity of 4 blocks
138
cpu_manager = CPUOffloadingManager(
144
- block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
139
+ num_blocks=4, cache_policy="lru", enable_events=True
140
)
141
142
# prepare store [1, 2]
@@ -163,9 +158,7 @@ def test_cpu_manager():
158
159
# complete store [1, 2]
160
cpu_manager.complete_store(to_keys([1, 2]))
166
- verify_events(
167
- cpu_manager.take_events(), block_size=block_size, expected_stores=({1, 2},)
168
- )
161
+ verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
162
163
# lookup [1, 2]
164
assert cpu_manager.lookup(to_keys([1])) == 1
@@ -184,9 +177,7 @@ def test_cpu_manager():
177
)
178
179
# verify eviction event
187
- verify_events(
188
- cpu_manager.take_events(), block_size=block_size, expected_evictions=({1},)
189
- )
180
+ verify_events(cpu_manager.take_events(), expected_evictions=({1},))
181
182
# prepare store with no space
183
assert cpu_manager.prepare_store(to_keys([1, 6])) is None
@@ -241,7 +232,6 @@ def test_cpu_manager():
232
233
verify_events(
234
cpu_manager.take_events(),
244
- block_size=block_size,
235
expected_stores=({3, 4, 5}, {6, 7, 8}),
236
expected_evictions=({2, 3, 4}, {8}),
237
)
@@ -254,7 +244,6 @@ class TestARCPolicy:
244
self, num_blocks: int = 4, enable_events: bool = True
245
) -> tuple[CPUOffloadingManager, ARCCachePolicy]:
246
manager = CPUOffloadingManager(
257
- block_size=256,
247
num_blocks=num_blocks,
248
cache_policy="arc",
249
enable_events=enable_events,
@@ -289,9 +278,7 @@ class TestARCPolicy:
278
279
# complete store [1, 2]
280
cpu_manager.complete_store(to_keys([1, 2]))
292
- verify_events(
293
- cpu_manager.take_events(), block_size=256, expected_stores=({1, 2},)
294
- )
281
+ verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
282
283
# lookup [1, 2]
284
assert cpu_manager.lookup(to_keys([1])) == 1
@@ -547,9 +534,8 @@ def test_filter_reused_manager():
534
"""
535
Tests FilterReusedOffloadingManager with a CPUOffloadingManager.
536
"""
550
- block_size = 256
537
lru_manager = CPUOffloadingManager(
552
- block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
538
+ num_blocks=4, cache_policy="lru", enable_events=True
539
)
540
541
manager = FilterReusedOffloadingManager(
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+1
-1
@@ -424,7 +424,7 @@ class OffloadingConnectorScheduler:
424
parent_block_hash=None,
425
token_ids=[],
426
lora_id=None,
427
- block_size=event.block_size,
427
+ block_size=0,
428
medium=event.medium,
429
lora_name=None,
430
)
vllm/v1/kv_offload/abstract.py
-1
@@ -79,7 +79,6 @@ class PrepareStoreOutput:
79
@dataclass
80
class OffloadingEvent:
81
keys: list[OffloadKey]
82
- block_size: int
82
medium: str
83
# True if blocks are removed, False if stored
84
removed: bool
vllm/v1/kv_offload/cpu/manager.py
-4
@@ -33,12 +33,10 @@ class CPUOffloadingManager(OffloadingManager):
33
34
def __init__(
35
self,
36
- block_size: int,
36
num_blocks: int,
37
cache_policy: Literal["lru", "arc"] = "lru",
38
enable_events: bool = False,
39
):
41
- self.block_size: int = block_size
40
self.medium: str = CPULoadStoreSpec.medium()
41
self._num_blocks: int = num_blocks
42
self._num_allocated_blocks: int = 0
@@ -145,7 +143,6 @@ class CPUOffloadingManager(OffloadingManager):
143
self.events.append(
144
OffloadingEvent(
145
keys=to_evict,
148
- block_size=self.block_size,
146
medium=self.medium,
147
removed=True,
148
)
@@ -188,7 +185,6 @@ class CPUOffloadingManager(OffloadingManager):
185
self.events.append(
186
OffloadingEvent(
187
keys=stored_keys,
191
- block_size=self.block_size,
188
medium=self.medium,
189
removed=False,
190
)
vllm/v1/kv_offload/cpu/spec.py
-5
@@ -60,12 +60,7 @@ class CPUOffloadingSpec(OffloadingSpec):
60
kv_events_config is not None and kv_events_config.enable_kv_cache_events
61
)
62
63
- assert len(self.gpu_block_size) == 1
64
- gpu_block_size = self.gpu_block_size[0]
65
- offloaded_block_size = gpu_block_size * self.block_size_factor
66
-
63
self._manager = CPUOffloadingManager(
68
- block_size=offloaded_block_size,
64
num_blocks=self.num_blocks,
65
cache_policy=self.eviction_policy, # type: ignore[arg-type]
66
enable_events=enable_events,