[KV Offload] Define clean backend configuration boundary (#48150)
Signed-off-by: Change72 <changg@nvidia.com> Signed-off-by: Chang Guo <cguo51@asu.edu> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Or Ozeri <or@ozery.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Chang Guo committed
Jul 16, 2026 at 03:27 UTC
a9531edfa63197dc083cf60be85e777c64fbfb60
27 files changed
+1275
-843
tests/v1/kv_connector/unit/offloading_connector/test_events.py
+25
-22
@@ -8,6 +8,9 @@ import torch
8
from tests.v1.kv_connector.unit.utils import create_vllm_config
9
from vllm.config import KVEventsConfig, KVTransferConfig
10
from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
11
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
12
+ build_offloading_config,
13
+)
14
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import (
15
OffloadingEventGroupSpec,
16
OffloadingEventsTracker,
@@ -70,15 +73,15 @@ def _group_config(
73
*,
74
group_idx: int = 0,
75
block_size: int = 4,
73
- block_size_factor: int = 1,
74
- sliding_window_size_in_blocks: int | None = None,
76
+ blocks_per_chunk: int = 1,
77
+ sliding_window_size_in_chunks: int | None = None,
78
) -> GroupOffloadConfig:
79
return GroupOffloadConfig(
80
group_idx=group_idx,
78
- gpu_block_size=block_size,
79
- offloaded_block_size=block_size * block_size_factor,
80
- hash_block_size_factor=block_size_factor,
81
- sliding_window_size_in_blocks=sliding_window_size_in_blocks,
81
+ tokens_per_block=block_size,
82
+ tokens_per_chunk=block_size * blocks_per_chunk,
83
+ hashes_per_chunk=blocks_per_chunk,
84
+ sliding_window_size_in_chunks=sliding_window_size_in_chunks,
85
kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC,
86
)
87
@@ -90,7 +93,7 @@ def _record_chunks(
93
num_chunks: int,
94
) -> list[OffloadKey]:
95
keys: list[OffloadKey] = []
93
- hbf = group_config.hash_block_size_factor
96
+ hbf = group_config.hashes_per_chunk
97
for chunk_idx in range(num_chunks):
98
tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1]
99
assert tail_hash is not None
@@ -149,14 +152,14 @@ def test_take_events_publishes_routable_block_stored():
152
153
def test_take_events_factor_gt_1_chunk_store_and_remove():
154
block_size = 4
152
- block_size_factor = 3
155
+ blocks_per_chunk = 3
156
tracker = _tracker()
157
group_config = _group_config(
155
- block_size=block_size, block_size_factor=block_size_factor
158
+ block_size=block_size, blocks_per_chunk=blocks_per_chunk
159
)
160
req = _request(
161
block_hashes=[_hash(i) for i in range(6)],
159
- token_count=block_size * block_size_factor * 2,
162
+ token_count=block_size * blocks_per_chunk * 2,
163
)
164
keys = _record_chunks(tracker, req, group_config, num_chunks=2)
165
@@ -169,17 +172,17 @@ def test_take_events_factor_gt_1_chunk_store_and_remove():
172
expected_chunk_hashes = [
173
_wire_hash(_hash(i))
174
for i in range(
172
- chunk_idx * block_size_factor,
173
- (chunk_idx + 1) * block_size_factor,
175
+ chunk_idx * blocks_per_chunk,
176
+ (chunk_idx + 1) * blocks_per_chunk,
177
)
178
]
179
assert event.block_hashes == expected_chunk_hashes
180
assert event.block_size == block_size
178
- assert len(event.token_ids) == block_size * block_size_factor
181
+ assert len(event.token_ids) == block_size * blocks_per_chunk
182
if chunk_idx == 0:
183
assert event.parent_block_hash is None
184
else:
182
- assert event.parent_block_hash == _wire_hash(_hash(block_size_factor - 1))
185
+ assert event.parent_block_hash == _wire_hash(_hash(blocks_per_chunk - 1))
186
expected_hashes.extend(expected_chunk_hashes)
187
188
assert len(tracker._pending_event_metadata) == 2
@@ -194,12 +197,12 @@ def test_take_events_factor_gt_1_chunk_store_and_remove():
197
198
199
def test_take_events_factor_gt_1_store_is_order_independent():
197
- block_size_factor = 3
200
+ blocks_per_chunk = 3
201
tracker = _tracker()
199
- group_config = _group_config(block_size_factor=block_size_factor)
202
+ group_config = _group_config(blocks_per_chunk=blocks_per_chunk)
203
req = _request(
204
block_hashes=[_hash(i) for i in range(6)],
202
- token_count=4 * block_size_factor * 2,
205
+ token_count=4 * blocks_per_chunk * 2,
206
)
207
keys = _record_chunks(tracker, req, group_config, num_chunks=2)
208
unknown_key = make_offload_key(_hash(12345), 0)
@@ -244,7 +247,7 @@ def test_take_events_opt_out_keeps_placeholders():
247
248
def test_record_store_skips_sliding_window_group():
249
tracker = _tracker()
247
- group_config = _group_config(sliding_window_size_in_blocks=2)
250
+ group_config = _group_config(sliding_window_size_in_chunks=2)
251
req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
252
keys = _record_chunks(tracker, req, group_config, num_chunks=3)
253
@@ -258,8 +261,8 @@ def test_record_store_skips_sliding_window_group():
261
262
def test_take_events_groups_removed_hashes_by_kv_group():
263
tracker = _tracker()
261
- group0_config = _group_config(group_idx=0, block_size_factor=2)
262
- group1_config = _group_config(group_idx=1, block_size_factor=2)
264
+ group0_config = _group_config(group_idx=0, blocks_per_chunk=2)
265
+ group1_config = _group_config(group_idx=1, blocks_per_chunk=2)
266
req0 = _request(block_hashes=[_hash(0), _hash(1)], token_count=8)
267
req1 = _request(block_hashes=[_hash(10), _hash(11)], token_count=8)
268
key0 = _record_chunks(tracker, req0, group0_config, num_chunks=1)[0]
@@ -293,7 +296,7 @@ def test_take_events_supports_restore_after_eviction():
296
assert not tracker._pending_event_metadata
297
298
req.all_token_ids = [5, 6, 7, 8]
296
- tracker.record_store(req, group_config, offload_block_idx=0, offload_key=key)
299
+ tracker.record_store(req, group_config, chunk_idx=0, offload_key=key)
300
301
second_store = list(tracker.take_events([_stored_event([key])]))
302
assert len(second_store) == 1
@@ -351,4 +354,4 @@ def test_tiering_rejects_self_describing_kv_events():
354
)
355
356
with pytest.raises(ValueError, match="TieringOffloadingSpec"):
354
- TieringOffloadingSpec(vllm_config, kv_cache_config)
357
+ TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+167
-169
@@ -105,20 +105,20 @@ def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
105
@pytest.mark.parametrize("async_scheduling", [True, False])
106
def test_offloading_connector(request_runner, async_scheduling: bool):
107
block_size = 4
108
- block_size_factor = 3
109
- offloaded_block_size = block_size * block_size_factor
108
+ blocks_per_chunk = 3
109
+ tokens_per_chunk = block_size * blocks_per_chunk
110
num_gpu_blocks = 100
111
112
runner = request_runner(
113
block_size=block_size,
114
num_gpu_blocks=num_gpu_blocks,
115
async_scheduling=async_scheduling,
116
- block_size_factor=block_size_factor,
116
+ blocks_per_chunk=blocks_per_chunk,
117
)
118
119
# 3 blocks, store just the middle block (skip first and last)
120
# blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
121
- runner.new_request(token_ids=[0] * offloaded_block_size * 3)
121
+ runner.new_request(token_ids=[0] * tokens_per_chunk * 3)
122
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
123
generate_store_output(list(keys)[1:2])
124
)
@@ -126,7 +126,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
126
127
# add block missing 1 token -> no offload
128
runner.run(
129
- decoded_tokens=[0] * (offloaded_block_size - 1),
129
+ decoded_tokens=[0] * (tokens_per_chunk - 1),
130
expected_stored=(3, 4, 5),
131
)
132
runner.manager.touch.assert_not_called()
@@ -141,7 +141,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
141
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
142
generate_store_output([])
143
)
144
- runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
144
+ runner.run(decoded_tokens=[0] * (tokens_per_chunk + 1))
145
146
# 1 more block (+ token for kicking off offloading)
147
# now check touch was called with all 6 blocks
@@ -149,7 +149,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
149
generate_store_output(keys)
150
)
151
runner.run(
152
- decoded_tokens=[0] * (offloaded_block_size + 1),
152
+ decoded_tokens=[0] * (tokens_per_chunk + 1),
153
expected_stored=(15, 16, 17),
154
)
155
runner.manager.touch.assert_called()
@@ -160,7 +160,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
160
runner.run(decoded_tokens=[EOS_TOKEN_ID])
161
162
# create a new request differing only on the last token
163
- runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
163
+ runner.new_request(token_ids=[0] * (tokens_per_chunk * 6 - 1) + [1])
164
runner.run(decoded_tokens=[0])
165
runner.manager.touch.assert_called()
166
block_hashes2 = list(runner.manager.touch.call_args.args[0])
@@ -173,12 +173,12 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
173
# terminate request
174
runner.run(
175
decoded_tokens=[EOS_TOKEN_ID],
176
- expected_stored=tuple(range(6 * block_size_factor)),
176
+ expected_stored=tuple(range(6 * blocks_per_chunk)),
177
)
178
179
- # full_block_tokens - num_computed_tokens < offloaded_block_size
179
+ # full_block_tokens - num_computed_tokens < tokens_per_chunk
180
runner.new_request(
181
- token_ids=[0] * block_size + [1] * (offloaded_block_size - block_size)
181
+ token_ids=[0] * block_size + [1] * (tokens_per_chunk - block_size)
182
)
183
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
184
generate_store_output([])
@@ -187,7 +187,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
187
runner.manager.lookup.assert_not_called()
188
189
# single block lookup with no hits
190
- runner.new_request(token_ids=[1] * offloaded_block_size)
190
+ runner.new_request(token_ids=[1] * tokens_per_chunk)
191
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
192
generate_store_output([])
193
)
@@ -196,7 +196,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
196
197
# single block lookup with a hit
198
runner.scheduler.reset_prefix_cache()
199
- runner.new_request(token_ids=[0] * offloaded_block_size)
199
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
200
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
201
generate_store_output([])
202
)
@@ -204,9 +204,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
204
runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(0, 1, 2))
205
206
# single block lookup with a hit in a middle block
207
- runner.new_request(
208
- token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
209
- )
207
+ runner.new_request(token_ids=[0] * tokens_per_chunk * 2 + [1] * tokens_per_chunk)
208
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
209
generate_store_output([])
210
)
@@ -217,15 +215,15 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
215
@pytest.mark.parametrize("async_scheduling", [True, False])
216
def test_request_preemption(request_runner, async_scheduling: bool):
217
block_size = 4
220
- block_size_factor = 3
221
- offloaded_block_size = block_size * block_size_factor
218
+ blocks_per_chunk = 3
219
+ tokens_per_chunk = block_size * blocks_per_chunk
220
num_gpu_blocks = 100
221
222
runner = request_runner(
223
block_size=block_size,
224
num_gpu_blocks=num_gpu_blocks,
225
async_scheduling=async_scheduling,
228
- block_size_factor=block_size_factor,
226
+ blocks_per_chunk=blocks_per_chunk,
227
)
228
229
free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
@@ -233,7 +231,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):
231
232
# 2 blocks, store all, without flushing
233
# blocks = [0, 1, 2], [3, 4, 5]
236
- runner.new_request(token_ids=[0] * offloaded_block_size * 2)
234
+ runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
235
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
236
generate_store_output(keys)
237
)
@@ -247,7 +245,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):
245
generate_store_output(keys)
246
)
247
runner.run(
250
- decoded_tokens=[0] * (2 * offloaded_block_size - block_size),
248
+ decoded_tokens=[0] * (2 * tokens_per_chunk - block_size),
249
complete_transfers=False,
250
)
251
@@ -297,14 +295,14 @@ def test_on_request_finished_is_not_deferred_until_store_completion(
295
still arrive afterward for already-submitted transfer jobs.
296
"""
297
block_size = 4
300
- block_size_factor = 3
301
- offloaded_block_size = block_size * block_size_factor
298
+ blocks_per_chunk = 3
299
+ tokens_per_chunk = block_size * blocks_per_chunk
300
301
runner = request_runner(
302
block_size=block_size,
303
num_gpu_blocks=100,
304
async_scheduling=async_scheduling,
307
- block_size_factor=block_size_factor,
305
+ blocks_per_chunk=blocks_per_chunk,
306
)
307
308
# Record the order of per-request connector calls on the (mocked) manager.
@@ -324,10 +322,10 @@ def test_on_request_finished_is_not_deferred_until_store_completion(
322
323
# Decode a couple of blocks, keeping every transfer in flight
324
# (complete_transfers=False) so no store completes while the request runs.
327
- runner.new_request(token_ids=[0] * offloaded_block_size * 2)
325
+ runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
326
runner.run(decoded_tokens=[0], complete_transfers=False)
327
runner.run(
330
- decoded_tokens=[0] * (2 * offloaded_block_size),
328
+ decoded_tokens=[0] * (2 * tokens_per_chunk),
329
complete_transfers=False,
330
)
331
@@ -347,7 +345,7 @@ def test_on_request_finished_is_not_deferred_until_store_completion(
345
runner.run(
346
decoded_tokens=[],
347
complete_transfers=True,
350
- expected_stored=tuple(range(4 * block_size_factor)),
348
+ expected_stored=tuple(range(4 * blocks_per_chunk)),
349
)
350
351
# on_request_finished is issued exactly once.
@@ -364,19 +362,19 @@ def test_on_request_finished_is_not_deferred_until_store_completion(
362
@pytest.mark.parametrize("async_scheduling", [True, False])
363
def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
364
block_size = 4
367
- block_size_factor = 3
368
- offloaded_block_size = block_size * block_size_factor
365
+ blocks_per_chunk = 3
366
+ tokens_per_chunk = block_size * blocks_per_chunk
367
num_gpu_blocks = 100
368
369
runner = request_runner(
370
block_size=block_size,
371
num_gpu_blocks=num_gpu_blocks,
372
async_scheduling=async_scheduling,
375
- block_size_factor=block_size_factor,
373
+ blocks_per_chunk=blocks_per_chunk,
374
)
375
376
# store 1 blocks
379
- runner.new_request(token_ids=[0] * offloaded_block_size)
377
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
378
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
379
generate_store_output(keys)
380
)
@@ -387,7 +385,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
385
386
# start a request to load the first block, but don't complete
387
runner.scheduler.reset_prefix_cache()
390
- runner.new_request(token_ids=[0] * offloaded_block_size)
388
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
389
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
390
runner.run(
391
decoded_tokens=[],
@@ -399,7 +397,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
397
assert transfer_jobs
398
399
# start a new request to load the same first block
402
- runner.new_request(token_ids=[0] * offloaded_block_size)
400
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
401
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
402
runner.run(
403
decoded_tokens=[],
@@ -428,19 +426,19 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
426
@pytest.mark.parametrize("async_scheduling", [True, False])
427
def test_abort_loading_requests(request_runner, async_scheduling: bool):
428
block_size = 4
431
- block_size_factor = 3
432
- offloaded_block_size = block_size * block_size_factor
429
+ blocks_per_chunk = 3
430
+ tokens_per_chunk = block_size * blocks_per_chunk
431
num_gpu_blocks = 100
432
433
runner = request_runner(
434
block_size=block_size,
435
num_gpu_blocks=num_gpu_blocks,
436
async_scheduling=async_scheduling,
439
- block_size_factor=block_size_factor,
437
+ blocks_per_chunk=blocks_per_chunk,
438
)
439
440
# store 1 blocks
443
- runner.new_request(token_ids=[0] * offloaded_block_size)
441
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
442
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
443
generate_store_output(keys)
444
)
@@ -451,7 +449,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
449
450
# start a request to load the first block, but don't complete
451
runner.scheduler.reset_prefix_cache()
454
- runner.new_request(token_ids=[0] * offloaded_block_size)
452
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
453
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
454
runner.run(
455
decoded_tokens=[],
@@ -483,7 +481,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
481
def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bool):
482
block_size = 4
483
num_gpu_blocks = 100
486
- # sliding_window=8 -> 2 offloaded blocks (block_size_factor=1)
484
+ # sliding_window=8 -> 2 offloaded chunks (blocks_per_chunk=1)
485
sliding_window = 8
486
487
kv_cache_groups = [
@@ -518,8 +516,8 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
516
# Verify group configs: group 0 = full attention, group 1 = sliding window
517
kv_group_configs = runner.connector_scheduler.config.kv_group_configs
518
assert len(kv_group_configs) == 2
521
- assert kv_group_configs[0].sliding_window_size_in_blocks is None
522
- assert kv_group_configs[1].sliding_window_size_in_blocks == 2
519
+ assert kv_group_configs[0].sliding_window_size_in_chunks is None
520
+ assert kv_group_configs[1].sliding_window_size_in_chunks == 2
521
522
# Blocks [0, 1, 2] miss
523
runner.new_request(token_ids=[0] * block_size * 3)
@@ -587,16 +585,16 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
585
586
@pytest.mark.parametrize("async_scheduling", [True, False])
587
def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool):
590
- hash_block_size = 4
588
+ tokens_per_hash = 4
589
num_gpu_blocks = 100
590
593
- # Group 0: block_size=12 (offloaded_block_size=12)
594
- # Group 1: block_size=16 (offloaded_block_size=16)
591
+ # Group 0: block_size=12 (tokens_per_chunk=12)
592
+ # Group 1: block_size=16 (tokens_per_chunk=16)
593
kv_cache_groups = [
594
KVCacheGroupSpec(
595
["layer0"],
596
FullAttentionSpec(
599
- block_size=hash_block_size * 3,
597
+ block_size=tokens_per_hash * 3,
598
num_kv_heads=1,
599
head_size=1,
600
dtype=torch.float32,
@@ -605,7 +603,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
603
KVCacheGroupSpec(
604
["layer1"],
605
FullAttentionSpec(
608
- block_size=hash_block_size * 4,
606
+ block_size=tokens_per_hash * 4,
607
num_kv_heads=1,
608
head_size=1,
609
dtype=torch.float32,
@@ -614,7 +612,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
612
]
613
614
runner = request_runner(
617
- block_size=hash_block_size,
615
+ block_size=tokens_per_hash,
616
num_gpu_blocks=num_gpu_blocks,
617
async_scheduling=async_scheduling,
618
kv_cache_groups=kv_cache_groups,
@@ -623,10 +621,10 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
621
# Verify group configs
622
kv_group_configs = runner.connector_scheduler.config.kv_group_configs
623
assert len(kv_group_configs) == 2
626
- assert kv_group_configs[0].gpu_block_size == 12
627
- assert kv_group_configs[0].offloaded_block_size == 12
628
- assert kv_group_configs[1].gpu_block_size == 16
629
- assert kv_group_configs[1].offloaded_block_size == 16
624
+ assert kv_group_configs[0].tokens_per_block == 12
625
+ assert kv_group_configs[0].tokens_per_chunk == 12
626
+ assert kv_group_configs[1].tokens_per_block == 16
627
+ assert kv_group_configs[1].tokens_per_chunk == 16
628
629
# Prompt: 25 tokens, unaligned to both block sizes.
630
# Group 0 blocks: [0, 1], ending_token_offset = 24
@@ -932,20 +930,20 @@ class TestSlidingWindowLookup:
930
@pytest.mark.parametrize("async_scheduling", [True, False])
931
def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool):
932
"""With REQUEST_LEVEL policy, all blocks are stored — including prefix hits."""
935
- gpu_block_size = 4
936
- block_size_factor = 3
937
- offloaded_block_size = gpu_block_size * block_size_factor
933
+ tokens_per_block = 4
934
+ blocks_per_chunk = 3
935
+ tokens_per_chunk = tokens_per_block * blocks_per_chunk
936
num_gpu_blocks = 100
937
938
runner = request_runner(
941
- block_size_factor=block_size_factor,
942
- block_size=gpu_block_size,
939
+ blocks_per_chunk=blocks_per_chunk,
940
+ block_size=tokens_per_block,
941
num_gpu_blocks=num_gpu_blocks,
942
async_scheduling=async_scheduling,
943
)
944
947
- # Store 1 offloaded block (3 GPU blocks) via a normal request.
948
- runner.new_request(token_ids=[0] * offloaded_block_size)
945
+ # Store 1 offloaded chunk (3 GPU blocks) via a normal request.
946
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
947
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
948
generate_store_output(keys)
949
)
@@ -962,14 +960,14 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling
960
policy=OffloadPolicy.REQUEST_LEVEL
961
)
962
965
- # New request with 2 offloaded blocks; first matches what's in CPU.
966
- runner.new_request(token_ids=[0] * offloaded_block_size * 2)
963
+ # New request with 2 offloaded chunks; first matches what's in CPU.
964
+ runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
965
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
966
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
967
generate_store_output(keys)
968
)
969
972
- # Load the first offloaded block from CPU.
970
+ # Load the first offloaded chunk from CPU.
971
runner.run(decoded_tokens=[0], expected_loaded=(0, 1, 2))
972
973
# Store must include ALL 6 GPU blocks (both the loaded prefix and
@@ -989,7 +987,7 @@ def test_loads_do_not_populate_fence_index(request_runner):
987
"""Loads don't populate _block_id_to_pending_jobs (protected by
988
delay_free_blocks while in flight)."""
989
runner = request_runner(
992
- block_size_factor=3,
990
+ blocks_per_chunk=3,
991
block_size=4,
992
num_gpu_blocks=100,
993
async_scheduling=False,
@@ -1008,7 +1006,7 @@ def test_fence_at_update_state_after_alloc(request_runner):
1006
req1 just freed.
1007
"""
1008
runner = request_runner(
1011
- block_size_factor=1,
1009
+ blocks_per_chunk=1,
1010
block_size=4,
1011
num_gpu_blocks=2,
1012
async_scheduling=False,
@@ -1059,7 +1057,7 @@ def test_fence_at_build_store_jobs(request_runner):
1057
reusing a finished request's pending-store block is flushed by
1058
_build_store_jobs's fence."""
1059
runner = request_runner(
1062
- block_size_factor=1,
1060
+ blocks_per_chunk=1,
1061
block_size=4,
1062
num_gpu_blocks=2,
1063
async_scheduling=False,
@@ -1108,16 +1106,16 @@ def test_fence_at_build_store_jobs(request_runner):
1106
def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
1107
"""complete_store fires per-job, not deferred to request finish.
1108
Each call carries only that store's keys."""
1111
- gpu_block_size = 4
1112
- block_size_factor = 3
1113
- offloaded_block_size = gpu_block_size * block_size_factor
1109
+ tokens_per_block = 4
1110
+ blocks_per_chunk = 3
1111
+ tokens_per_chunk = tokens_per_block * blocks_per_chunk
1112
runner = request_runner(
1115
- block_size_factor=block_size_factor,
1116
- block_size=gpu_block_size,
1113
+ blocks_per_chunk=blocks_per_chunk,
1114
+ block_size=tokens_per_block,
1115
num_gpu_blocks=100,
1116
async_scheduling=async_scheduling,
1117
)
1120
- runner.new_request(token_ids=[0] * offloaded_block_size)
1118
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
1119
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1120
generate_store_output(keys)
1121
)
@@ -1131,7 +1129,7 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
1129
1130
# Second store: fires when block 1 is fully populated, with different keys.
1131
runner.run(
1134
- decoded_tokens=[0] * (offloaded_block_size + 1),
1132
+ decoded_tokens=[0] * (tokens_per_chunk + 1),
1133
expected_stored=(3, 4, 5),
1134
)
1135
assert runner.manager.complete_store.call_count == 1
@@ -1148,25 +1146,25 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
1146
def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
1147
"""Validates max_offload_tokens: type coercion, boundary values, and capping.
1148
1151
- Setup: 3 offloaded blocks × 3 GPU blocks each = 9 GPU block offsets (0–8).
1149
+ Setup: 3 offloaded chunks × 3 GPU blocks each = 9 GPU block offsets (0–8).
1150
"""
1153
- gpu_block_size = 4
1154
- block_size_factor = 3
1155
- offloaded_block_size = gpu_block_size * block_size_factor # 12
1151
+ tokens_per_block = 4
1152
+ blocks_per_chunk = 3
1153
+ tokens_per_chunk = tokens_per_block * blocks_per_chunk # 12
1154
num_gpu_blocks = 100
1155
all_offsets = (0, 1, 2, 3, 4, 5, 6, 7, 8)
1156
1157
def make_runner():
1158
return request_runner(
1161
- block_size=gpu_block_size,
1159
+ block_size=tokens_per_block,
1160
num_gpu_blocks=num_gpu_blocks,
1161
async_scheduling=async_scheduling,
1164
- block_size_factor=block_size_factor,
1162
+ blocks_per_chunk=blocks_per_chunk,
1163
)
1164
1165
def setup(r, max_offload_tokens):
1166
r.new_request(
1169
- token_ids=[0] * offloaded_block_size * 3,
1167
+ token_ids=[0] * tokens_per_chunk * 3,
1168
kv_transfer_params={"max_offload_tokens": max_offload_tokens},
1169
)
1170
r.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -1228,9 +1226,9 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
1226
setup(r, 0)
1227
r.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=())
1228
1231
- # positive int cap -> limits offload to first 2 offloaded blocks (offsets 0–5)
1229
+ # positive int cap -> limits offload to first 2 chunks (offsets 0–5)
1230
r = make_runner()
1233
- setup(r, 24) # 24 tokens = 2 offloaded blocks × 12 tokens each
1231
+ setup(r, 24) # 24 tokens = 2 offloaded chunks × 12 tokens each
1232
r.run(
1233
decoded_tokens=[EOS_TOKEN_ID],
1234
expected_stored=(0, 1, 2, 3, 4, 5),
@@ -1242,8 +1240,8 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
1240
def test_offload_prompt_only(request_runner, async_scheduling: bool):
1241
"""offload_prompt_only=True offloads prompt blocks but never decode blocks.
1242
1245
- Setup: a 2-offloaded-block prompt followed by enough decode tokens to fill
1246
- 4 more offloaded blocks. The flag clamps the offloadable token count to the
1243
+ Setup: a 2-chunk prompt followed by enough decode tokens to fill
1244
+ 4 more offloaded chunks. The flag clamps the offloadable token count to the
1245
prompt length, so only the prompt's blocks (GPU offsets 0-5) are ever
1246
eligible for store; the decode blocks (offsets >= 6) are skipped.
1247
@@ -1253,18 +1251,18 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool):
1251
subtleties. The decode steps are still enough for the prompt store to
1252
complete and show up in expected_stored.
1253
"""
1256
- gpu_block_size = 4
1257
- block_size_factor = 3
1258
- offloaded_block_size = gpu_block_size * block_size_factor # 12
1254
+ tokens_per_block = 4
1255
+ blocks_per_chunk = 3
1256
+ tokens_per_chunk = tokens_per_block * blocks_per_chunk # 12
1257
num_prompt_blocks = 2
1258
num_decode_blocks = 4
1259
prompt_offsets = (0, 1, 2, 3, 4, 5)
1260
1261
runner = request_runner(
1264
- block_size=gpu_block_size,
1262
+ block_size=tokens_per_block,
1263
num_gpu_blocks=100,
1264
async_scheduling=async_scheduling,
1267
- block_size_factor=block_size_factor,
1265
+ blocks_per_chunk=blocks_per_chunk,
1266
extra_config_overrides={"offload_prompt_only": True},
1267
)
1268
@@ -1272,9 +1270,9 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool):
1270
generate_store_output(keys)
1271
)
1272
1275
- runner.new_request(token_ids=[0] * offloaded_block_size * num_prompt_blocks)
1273
+ runner.new_request(token_ids=[0] * tokens_per_chunk * num_prompt_blocks)
1274
runner.run(
1277
- decoded_tokens=[0] * (offloaded_block_size * num_decode_blocks),
1275
+ decoded_tokens=[0] * (tokens_per_chunk * num_decode_blocks),
1276
expected_stored=prompt_offsets,
1277
)
1278
@@ -1291,21 +1289,21 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool):
1289
@pytest.mark.parametrize("async_scheduling", [True, False])
1290
def test_reset_cache(request_runner, async_scheduling: bool):
1291
"""reset_cache flushes in-flight loads, calls manager.reset_cache(), resets
1294
- next_stored_block_idx for active requests and clears job tracking."""
1292
+ next_stored_chunk_idx for active requests and clears job tracking."""
1293
block_size = 4
1296
- block_size_factor = 3
1297
- offloaded_block_size = block_size * block_size_factor
1294
+ blocks_per_chunk = 3
1295
+ tokens_per_chunk = block_size * blocks_per_chunk
1296
num_gpu_blocks = 100
1297
1298
runner = request_runner(
1299
block_size=block_size,
1300
num_gpu_blocks=num_gpu_blocks,
1301
async_scheduling=async_scheduling,
1304
- block_size_factor=block_size_factor,
1302
+ blocks_per_chunk=blocks_per_chunk,
1303
)
1304
1307
- # Store 1 offloaded block (3 GPU blocks) to CPU.
1308
- runner.new_request(token_ids=[0] * offloaded_block_size)
1305
+ # Store 1 offloaded chunk (3 GPU blocks) to CPU.
1306
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
1307
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1308
generate_store_output(keys)
1309
)
@@ -1317,7 +1315,7 @@ def test_reset_cache(request_runner, async_scheduling: bool):
1315
# Reset GPU prefix cache then start a request that loads from CPU.
1316
# Leave the load in-flight so that reset_cache must flush it.
1317
runner.scheduler.reset_prefix_cache()
1320
- runner.new_request(token_ids=[0] * offloaded_block_size)
1318
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
1319
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
1320
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1321
generate_store_output([])
@@ -1335,11 +1333,11 @@ def test_reset_cache(request_runner, async_scheduling: bool):
1333
# Record job counter to verify the reset counter is set correctly.
1334
job_counter_before_reset = runner.connector_scheduler._job_counter
1335
1338
- # After update_state_after_alloc, next_stored_block_idx is advanced to
1336
+ # After update_state_after_alloc, next_stored_chunk_idx is advanced to
1337
# skip the loaded prefix; reset_cache must bring it back to 0.
1338
for req_status in runner.connector_scheduler._req_status.values():
1339
for group_state in req_status.group_states:
1342
- assert group_state.next_stored_block_idx > 0
1340
+ assert group_state.next_stored_chunk_idx > 0
1341
1342
# Reset the cache
1343
runner.connector_scheduler.reset_cache()
@@ -1354,18 +1352,18 @@ def test_reset_cache(request_runner, async_scheduling: bool):
1352
# All internal job tracking must be cleared.
1353
assert not runner.connector_scheduler._jobs
1354
assert not runner.connector_scheduler._block_id_to_pending_jobs
1357
- if runner.connector_scheduler._blocks_being_loaded is not None:
1358
- assert not runner.connector_scheduler._blocks_being_loaded
1355
+ if runner.connector_scheduler._chunks_being_loaded is not None:
1356
+ assert not runner.connector_scheduler._chunks_being_loaded
1357
1358
# Job reset counter must equal the job counter so that completions for
1359
# pre-reset jobs arriving from workers are silently discarded.
1360
assert runner.connector_scheduler._stale_job_threshold == job_counter_before_reset
1361
1364
- # next_stored_block_idx must be reset to 0 for every active request so
1362
+ # next_stored_chunk_idx must be reset to 0 for every active request so
1363
# that post-reset stores restart from block 0.
1364
for req_status in runner.connector_scheduler._req_status.values():
1365
for group_state in req_status.group_states:
1368
- assert group_state.next_stored_block_idx == 0
1366
+ assert group_state.next_stored_chunk_idx == 0
1367
1368
1369
@pytest.mark.parametrize("async_scheduling", [True, False])
@@ -1376,14 +1374,14 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
1374
without calling on_request_finished twice.
1375
"""
1376
block_size = 4
1379
- block_size_factor = 3
1380
- offloaded_block_size = block_size * block_size_factor
1377
+ blocks_per_chunk = 3
1378
+ tokens_per_chunk = block_size * blocks_per_chunk
1379
1380
runner = request_runner(
1381
block_size=block_size,
1382
num_gpu_blocks=100,
1383
async_scheduling=async_scheduling,
1386
- block_size_factor=block_size_factor,
1384
+ blocks_per_chunk=blocks_per_chunk,
1385
)
1386
1387
finalized: list[str] = []
@@ -1396,10 +1394,10 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
1394
1395
# Decode a couple of blocks and keep every transfer in flight, so the
1396
# request has pending store jobs.
1399
- runner.new_request(token_ids=[0] * offloaded_block_size * 2)
1397
+ runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
1398
runner.run(decoded_tokens=[0], complete_transfers=False)
1399
runner.run(
1402
- decoded_tokens=[0] * (2 * offloaded_block_size),
1400
+ decoded_tokens=[0] * (2 * tokens_per_chunk),
1401
complete_transfers=False,
1402
)
1403
@@ -1430,7 +1428,7 @@ def test_pending_transfer_defers_prefix_lookup():
1428
With async scheduling, a preempted request's store can be flushed by the
1429
worker before the scheduler consumes its completion. If the request is
1430
re-admitted in that window, the connector should defer it instead of
1433
- looking up offloaded blocks and later asserting when a load is queued while
1431
+ looking up offloaded chunks and later asserting when a load is queued while
1432
the store job is still tracked.
1433
"""
1434
scheduler = object.__new__(OffloadingConnectorScheduler)
@@ -1464,27 +1462,27 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner
1462
re-admission path must defer while the scheduler still tracks the store.
1463
"""
1464
block_size = 4
1467
- block_size_factor = 3
1468
- offloaded_block_size = block_size * block_size_factor
1465
+ blocks_per_chunk = 3
1466
+ tokens_per_chunk = block_size * blocks_per_chunk
1467
1468
runner = request_runner(
1469
block_size=block_size,
1470
num_gpu_blocks=100,
1471
async_scheduling=True,
1474
- block_size_factor=block_size_factor,
1472
+ blocks_per_chunk=blocks_per_chunk,
1473
)
1474
free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
1475
num_free_blocks_empty = free_block_queue.num_free_blocks
1476
1477
req_id = "0"
1480
- runner.new_request(token_ids=[0] * offloaded_block_size * 2)
1478
+ runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
1479
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1480
generate_store_output(keys)
1481
)
1482
1483
runner.run(decoded_tokens=[0], complete_transfers=False)
1484
runner.run(
1487
- decoded_tokens=[0] * (2 * offloaded_block_size - block_size),
1485
+ decoded_tokens=[0] * (2 * tokens_per_chunk - block_size),
1486
complete_transfers=False,
1487
)
1488
@@ -1530,8 +1528,8 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
1528
- Group 0: full attention (MLA-like), block_size=16
1529
- Group 1: SWA, block_size=4, sliding_window=8
1530
1533
- alignment_block_count = 16 / 4 = 4 SWA blocks per alignment segment.
1534
- sliding_window_size_in_blocks = ceil(8 / 4) = 2.
1531
+ alignment_chunk_count = 16 / 4 = 4 SWA blocks per alignment segment.
1532
+ sliding_window_size_in_chunks = ceil(8 / 4) = 2.
1533
Within each segment of 4 SWA blocks, only the trailing 2 are stored.
1534
1535
With 32 tokens (2 full-attn blocks, 8 SWA blocks):
@@ -1574,17 +1572,17 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
1572
kv_cache_groups=kv_cache_groups,
1573
)
1574
1577
- # Verify config: alignment_block_count computed correctly
1575
+ # Verify config: alignment_chunk_count computed correctly
1576
kv_group_configs = runner.connector_scheduler.config.kv_group_configs
1577
assert len(kv_group_configs) == 2
1578
# Group 0: full attention -> no alignment skip
1581
- assert kv_group_configs[0].alignment_block_count is None
1582
- assert kv_group_configs[0].sliding_window_size_in_blocks is None
1583
- assert kv_group_configs[0].offloaded_block_size == full_attn_block_size
1584
- # Group 1: SWA -> alignment_block_count = 16/4 = 4, tail = 2
1585
- assert kv_group_configs[1].alignment_block_count == 4
1586
- assert kv_group_configs[1].sliding_window_size_in_blocks == 2
1587
- assert kv_group_configs[1].offloaded_block_size == swa_block_size
1579
+ assert kv_group_configs[0].alignment_chunk_count is None
1580
+ assert kv_group_configs[0].sliding_window_size_in_chunks is None
1581
+ assert kv_group_configs[0].tokens_per_chunk == full_attn_block_size
1582
+ # Group 1: SWA -> alignment_chunk_count = 16/4 = 4, tail = 2
1583
+ assert kv_group_configs[1].alignment_chunk_count == 4
1584
+ assert kv_group_configs[1].sliding_window_size_in_chunks == 2
1585
+ assert kv_group_configs[1].tokens_per_chunk == swa_block_size
1586
1587
# Send 32 tokens = 2 full-attn blocks (block_size=16) = 8 SWA blocks
1588
# (block_size=4). Decode 1 token to kick off processing (stores are
@@ -1602,9 +1600,9 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
1600
)
1601
runner.run(
1602
decoded_tokens=[EOS_TOKEN_ID],
1605
- # Group 0 (full attn, block_size=16): 2 offloaded blocks
1603
+ # Group 0 (full attn, block_size=16): 2 offloaded chunks
1604
# -> GPU blocks (0, 0) and (0, 1)
1607
- # Group 1 (SWA, block_size=4): 8 offloaded blocks, skip first 2
1605
+ # Group 1 (SWA, block_size=4): 8 offloaded chunks, skip first 2
1606
# per segment of 4:
1607
# Segment 0 (blocks 0-3): skip 0,1 -> store (1, 2), (1, 3)
1608
# Segment 1 (blocks 4-7): skip 4,5 -> store (1, 6), (1, 7)
@@ -1625,7 +1623,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
1623
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2
1624
runner.run(
1625
decoded_tokens=[EOS_TOKEN_ID],
1628
- # Group 0: full prefix lookup hits 2 offloaded blocks
1626
+ # Group 0: full prefix lookup hits 2 offloaded chunks
1627
# -> loads GPU blocks (0, 0), (0, 1)
1628
# Group 1: sliding window lookup finds trailing 2 from last segment
1629
# (blocks 6, 7 which were stored)
@@ -1684,7 +1682,7 @@ def test_stale_sliding_window_block_after_prepare_store_failure(
1682
runner.new_request(token_ids=[0] * block_size * 3)
1683
1684
# First step: prepare_store FAILS -> offloading delayed.
1687
- # next_stored_block_idx stays at 0, block_ids[0] still holds the
1685
+ # next_stored_chunk_idx stays at 0, block_ids[0] still holds the
1686
# original block_id for position 0.
1687
runner.manager.prepare_store.side_effect = lambda keys, req_context: None
1688
runner.run(decoded_tokens=[0])
@@ -1719,19 +1717,19 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool):
1717
"""When skip_reading_prefix_cache=True, the offloading connector must not
1718
load any blocks from CPU even if a matching prefix is cached there."""
1719
block_size = 4
1722
- block_size_factor = 3
1723
- offloaded_block_size = block_size * block_size_factor
1720
+ blocks_per_chunk = 3
1721
+ tokens_per_chunk = block_size * blocks_per_chunk
1722
num_gpu_blocks = 100
1723
1724
runner = request_runner(
1725
block_size=block_size,
1726
num_gpu_blocks=num_gpu_blocks,
1727
async_scheduling=async_scheduling,
1730
- block_size_factor=block_size_factor,
1728
+ blocks_per_chunk=blocks_per_chunk,
1729
)
1730
1731
# Populate the CPU offload cache with one block.
1734
- runner.new_request(token_ids=[0] * offloaded_block_size)
1732
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
1733
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1734
generate_store_output(keys)
1735
)
@@ -1747,7 +1745,7 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool):
1745
# The offloading connector must not load anything from CPU, but must
1746
# still offload the freshly computed blocks (state management intact).
1747
runner.new_request(
1750
- token_ids=[0] * offloaded_block_size,
1748
+ token_ids=[0] * tokens_per_chunk,
1749
skip_reading_prefix_cache=True,
1750
)
1751
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -2219,7 +2217,7 @@ class TestEagle:
2217
def test_full_attn_store_excludes_trailing_decode_block(
2218
self, request_runner, async_scheduling: bool
2219
):
2222
- """Eagle full-attention group excludes the trailing block only while
2220
+ """Eagle full-attention group excludes the trailing chunk only while
2221
decoding.
2222
2223
Setup: 2 groups — group 0 is normal full-attention, group 1 is
@@ -2229,8 +2227,8 @@ class TestEagle:
2227
draft-layer KV is volatile until the next block starts).
2228
"""
2229
block_size = 4
2232
- block_size_factor = 1
2233
- offloaded_block_size = block_size * block_size_factor
2230
+ blocks_per_chunk = 1
2231
+ tokens_per_chunk = block_size * blocks_per_chunk
2232
num_gpu_blocks = 100
2233
2234
kv_cache_groups = [
@@ -2260,7 +2258,7 @@ class TestEagle:
2258
num_gpu_blocks=num_gpu_blocks,
2259
async_scheduling=async_scheduling,
2260
kv_cache_groups=kv_cache_groups,
2263
- block_size_factor=block_size_factor,
2261
+ blocks_per_chunk=blocks_per_chunk,
2262
)
2263
2264
kv_group_configs = runner.connector_scheduler.config.kv_group_configs
@@ -2268,7 +2266,7 @@ class TestEagle:
2266
assert not kv_group_configs[0].is_eagle_group
2267
assert kv_group_configs[1].is_eagle_group
2268
2271
- runner.new_request(token_ids=[0] * offloaded_block_size * 3)
2269
+ runner.new_request(token_ids=[0] * tokens_per_chunk * 3)
2270
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2271
generate_store_output(keys)
2272
)
@@ -2292,7 +2290,7 @@ class TestEagle:
2290
self, request_runner, async_scheduling: bool
2291
):
2292
"""Eagle sliding-window group stores all prompt blocks but excludes
2295
- the trailing block while decoding."""
2293
+ the trailing chunk while decoding."""
2294
block_size = 4
2295
sliding_window = 8
2296
num_gpu_blocks = 100
@@ -2321,7 +2319,7 @@ class TestEagle:
2319
kv_group_configs = runner.connector_scheduler.config.kv_group_configs
2320
assert len(kv_group_configs) == 1
2321
assert kv_group_configs[0].is_eagle_group
2324
- assert kv_group_configs[0].sliding_window_size_in_blocks == 2
2322
+ assert kv_group_configs[0].sliding_window_size_in_chunks == 2
2323
2324
runner.new_request(token_ids=[0] * block_size * 3)
2325
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -2340,8 +2338,8 @@ class TestEagle:
2338
"""An eagle group with a single-block prompt stores it at the end of
2339
prefill: prompt blocks are stable, so no tail is held back."""
2340
block_size = 4
2343
- block_size_factor = 1
2344
- offloaded_block_size = block_size * block_size_factor
2341
+ blocks_per_chunk = 1
2342
+ tokens_per_chunk = block_size * blocks_per_chunk
2343
num_gpu_blocks = 100
2344
2345
kv_cache_groups = [
@@ -2362,10 +2360,10 @@ class TestEagle:
2360
num_gpu_blocks=num_gpu_blocks,
2361
async_scheduling=async_scheduling,
2362
kv_cache_groups=kv_cache_groups,
2365
- block_size_factor=block_size_factor,
2363
+ blocks_per_chunk=blocks_per_chunk,
2364
)
2365
2368
- runner.new_request(token_ids=[0] * offloaded_block_size)
2366
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
2367
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2368
generate_store_output(keys)
2369
)
@@ -2378,18 +2376,18 @@ class TestEagle:
2376
"""Eagle store must not drop interior blocks across prefill chunks.
2377
2378
Regression: the trailing-block exclusion (num_blocks - 1) was applied
2381
- when collecting keys, but next_stored_block_idx advanced by the
2382
- non-decremented count, so the trailing block of every chunked-prefill
2379
+ when collecting keys, but next_stored_chunk_idx advanced by the
2380
+ non-decremented count, so the trailing chunk of every chunked-prefill
2381
chunk was skipped and never re-considered. With the harness chunk budget
2382
(1000 tokens) and block_size 4, a prompt longer than one chunk lost the
2383
block at the chunk boundary, leaving a permanent gap that caps prefix
2384
reuse at the first hole. Only the trailing decode block may be held
2385
back; all other blocks must be stored exactly once (no duplicates from
2388
- next_stored_block_idx regressing at the prefill->decode transition).
2386
+ next_stored_chunk_idx regressing at the prefill->decode transition).
2387
"""
2388
block_size = 4
2391
- block_size_factor = 1
2392
- offloaded_block_size = block_size * block_size_factor
2389
+ blocks_per_chunk = 1
2390
+ tokens_per_chunk = block_size * blocks_per_chunk
2391
num_gpu_blocks = 1000
2392
2393
kv_cache_groups = [
@@ -2409,13 +2407,13 @@ class TestEagle:
2407
num_gpu_blocks=num_gpu_blocks,
2408
async_scheduling=async_scheduling,
2409
kv_cache_groups=kv_cache_groups,
2412
- block_size_factor=block_size_factor,
2410
+ blocks_per_chunk=blocks_per_chunk,
2411
)
2412
assert runner.connector_scheduler.config.kv_group_configs[0].is_eagle_group
2413
2414
# Prompt spans more than one prefill chunk (chunk budget 1000 tokens).
2415
num_blocks = 256
2418
- runner.new_request(token_ids=[0] * offloaded_block_size * num_blocks)
2416
+ runner.new_request(token_ids=[0] * tokens_per_chunk * num_blocks)
2417
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2418
generate_store_output(keys)
2419
)
@@ -2429,7 +2427,7 @@ class TestEagle:
2427
for b in t.gpu_blocks
2428
)
2429
# The stored blocks must be contiguous from 0: no interior block is
2432
- # dropped at a chunk boundary. (The bug left a gap at offloaded block
2430
+ # dropped at a chunk boundary. (The bug left a gap at offloaded chunk
2431
# 249, the tail of the first 1000-token chunk.)
2432
assert offsets == list(range(len(offsets))), (
2433
f"interior hole in stored blocks: {offsets}"
@@ -2439,14 +2437,14 @@ class TestEagle:
2437
def test_full_attn_store_then_load(self, request_runner, async_scheduling: bool):
2438
"""Eagle group constrains load: convergence tightens both groups.
2439
2442
- Store 3 offloaded blocks per group (all prompt blocks, so the eagle
2440
+ Store 3 offloaded chunks per group (all prompt chunks, so the eagle
2441
group stores all 3 as well). Then a new request loads from CPU. The
2442
eagle group pops its trailing hit block on load, tightening the hit
2443
to 2 blocks for both groups.
2444
"""
2445
block_size = 4
2448
- block_size_factor = 1
2449
- offloaded_block_size = block_size * block_size_factor
2446
+ blocks_per_chunk = 1
2447
+ tokens_per_chunk = block_size * blocks_per_chunk
2448
num_gpu_blocks = 100
2449
2450
kv_cache_groups = [
@@ -2476,10 +2474,10 @@ class TestEagle:
2474
num_gpu_blocks=num_gpu_blocks,
2475
async_scheduling=async_scheduling,
2476
kv_cache_groups=kv_cache_groups,
2479
- block_size_factor=block_size_factor,
2477
+ blocks_per_chunk=blocks_per_chunk,
2478
)
2479
2482
- runner.new_request(token_ids=[0] * offloaded_block_size * 3)
2480
+ runner.new_request(token_ids=[0] * tokens_per_chunk * 3)
2481
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2482
generate_store_output(keys)
2483
)
@@ -2497,7 +2495,7 @@ class TestEagle:
2495
2496
runner.scheduler.reset_prefix_cache()
2497
2500
- runner.new_request(token_ids=[0] * offloaded_block_size * 3 + [1])
2498
+ runner.new_request(token_ids=[0] * tokens_per_chunk * 3 + [1])
2499
runner.manager.lookup.return_value = LookupResult.HIT
2500
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2501
generate_store_output([])
@@ -2527,8 +2525,8 @@ def test_request_finished_with_pending_stores_populates_fence(request_runner):
2525
GPU blocks before the store completes.
2526
"""
2527
block_size = 4
2530
- block_size_factor = 1
2531
- offloaded_block_size = block_size * block_size_factor
2528
+ blocks_per_chunk = 1
2529
+ tokens_per_chunk = block_size * blocks_per_chunk
2530
2531
# Use 2 GPU blocks so the second run reuses the same blocks,
2532
# triggering a fence-based flush of the in-flight job from run 1.
@@ -2536,11 +2534,11 @@ def test_request_finished_with_pending_stores_populates_fence(request_runner):
2534
block_size=block_size,
2535
num_gpu_blocks=2,
2536
async_scheduling=False,
2539
- block_size_factor=block_size_factor,
2537
+ blocks_per_chunk=blocks_per_chunk,
2538
)
2539
2540
# 4 prompt tokens → 1 GPU block (block 0)
2543
- runner.new_request(token_ids=[0] * offloaded_block_size)
2541
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
2542
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2543
generate_store_output(keys)
2544
)
@@ -2576,7 +2574,7 @@ def test_request_finished_with_pending_stores_populates_fence(request_runner):
2574
2575
# Run 2: block reuse triggers fence-based flush → cleanup.
2576
runner.scheduler.reset_prefix_cache()
2579
- runner.new_request(token_ids=[0] * offloaded_block_size)
2577
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
2578
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2579
generate_store_output(keys)
2580
)
@@ -2603,26 +2601,26 @@ def test_multiple_in_flight_stores_all_flushed_by_fence(request_runner):
2601
- Run 3: block reuse → both jobs flushed via fence
2602
"""
2603
block_size = 4
2606
- block_size_factor = 1
2607
- offloaded_block_size = block_size * block_size_factor
2604
+ blocks_per_chunk = 1
2605
+ tokens_per_chunk = block_size * blocks_per_chunk
2606
2607
# 4 GPU blocks: block 0 is null, blocks 1-3 are usable.
2608
runner = request_runner(
2609
block_size=block_size,
2610
num_gpu_blocks=4,
2611
async_scheduling=False,
2614
- block_size_factor=block_size_factor,
2612
+ blocks_per_chunk=blocks_per_chunk,
2613
)
2614
2615
# Prompt: 4 tokens → block 1
2618
- runner.new_request(token_ids=[0] * offloaded_block_size)
2616
+ runner.new_request(token_ids=[0] * tokens_per_chunk)
2617
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2618
generate_store_output(keys)
2619
)
2620
2621
# Run 1: 4 decoded tokens → block 2 full → job_0 created for block 1.
2622
runner.run(
2625
- decoded_tokens=[0] * offloaded_block_size,
2623
+ decoded_tokens=[0] * tokens_per_chunk,
2624
complete_transfers=False,
2625
)
2626
assert len(runner.connector_scheduler._jobs) >= 1
@@ -2630,7 +2628,7 @@ def test_multiple_in_flight_stores_all_flushed_by_fence(request_runner):
2628
# Run 2: 4 more tokens + EOS → block 3 full → more jobs created.
2629
# Request finishes → all jobs registered in fence.
2630
runner.run(
2633
- decoded_tokens=[0] * offloaded_block_size + [EOS_TOKEN_ID],
2631
+ decoded_tokens=[0] * tokens_per_chunk + [EOS_TOKEN_ID],
2632
complete_transfers=False,
2633
)
2634
num_jobs = len(runner.connector_scheduler._jobs)
@@ -2638,7 +2636,7 @@ def test_multiple_in_flight_stores_all_flushed_by_fence(request_runner):
2636
2637
# Run 3: block reuse → fence flushes both jobs.
2638
runner.scheduler.reset_prefix_cache()
2641
- runner.new_request(token_ids=[0] * offloaded_block_size * 3)
2639
+ runner.new_request(token_ids=[0] * tokens_per_chunk * 3)
2640
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2641
generate_store_output(keys)
2642
)
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+4
-3
@@ -96,11 +96,12 @@ def _make_worker(kv_cache_config: KVCacheConfig):
96
)
97
98
spec = MagicMock(spec=OffloadingSpec)
99
- spec.kv_cache_config = kv_cache_config
100
- spec.vllm_config = MagicMock()
99
spec.get_worker.return_value = MagicMock()
100
103
- worker = OffloadingConnectorWorker(spec=spec)
101
+ worker = OffloadingConnectorWorker(
102
+ spec=spec,
103
+ kv_cache_config=kv_cache_config,
104
+ )
105
worker.worker = MagicMock()
106
107
return worker, spec
tests/v1/kv_connector/unit/offloading_connector/utils.py
+17
-19
@@ -17,7 +17,6 @@ from vllm import SamplingParams
17
from vllm.config import (
18
KVEventsConfig,
19
KVTransferConfig,
20
- VllmConfig,
20
set_current_vllm_config,
21
)
22
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
@@ -57,6 +56,7 @@ from vllm.v1.kv_offload.base import (
56
TransferResult,
57
make_offload_key,
58
)
59
+from vllm.v1.kv_offload.config import OffloadingConfig
60
from vllm.v1.request import Request
61
from vllm.v1.structured_output import StructuredOutputManager
62
@@ -123,8 +123,8 @@ class MockOffloadingWorker(OffloadingWorker):
123
124
125
class MockOffloadingSpec(OffloadingSpec):
126
- def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
127
- super().__init__(vllm_config, kv_cache_config)
126
+ def __init__(self, config: OffloadingConfig):
127
+ super().__init__(config)
128
129
self.manager = MagicMock(spec=OffloadingManager)
130
self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)
@@ -175,17 +175,17 @@ class RequestRunner:
175
self,
176
block_size: int,
177
num_gpu_blocks: int,
178
- block_size_factor: int = 1,
178
+ blocks_per_chunk: int = 1,
179
async_scheduling: bool = True,
180
kv_cache_groups: list[KVCacheGroupSpec] | None = None,
181
extra_config_overrides: dict[str, Any] | None = None,
182
):
183
- assert block_size_factor == 1 or kv_cache_groups is None, (
184
- "block_size_factor > 1 requires all groups to have the same "
183
+ assert blocks_per_chunk == 1 or kv_cache_groups is None, (
184
+ "blocks_per_chunk > 1 requires all groups to have the same "
185
"block size, so kv_cache_groups must be None (use default group)"
186
)
187
188
- self.block_size_factor: int = block_size_factor
188
+ self.blocks_per_chunk: int = blocks_per_chunk
189
self.block_size: int = block_size
190
self.num_gpu_blocks: int = num_gpu_blocks
191
self.async_scheduling: bool = async_scheduling
@@ -208,8 +208,8 @@ class RequestRunner:
208
# opt-out tests override this to cover the legacy placeholders.
209
"self_describing_kv_events": True,
210
}
211
- if block_size_factor > 1:
212
- extra_config["block_size"] = block_size * block_size_factor
211
+ if blocks_per_chunk > 1:
212
+ extra_config["block_size"] = block_size * blocks_per_chunk
213
if extra_config_overrides:
214
extra_config.update(extra_config_overrides)
215
@@ -313,11 +313,9 @@ class RequestRunner:
313
self.connector_scheduler.config.kv_group_configs,
314
kv_cache_config.kv_cache_groups,
315
):
316
- gpu_block_size = kv_cache_group.kv_cache_spec.block_size
317
- assert group_config.gpu_block_size == gpu_block_size
318
- assert (
319
- group_config.offloaded_block_size == gpu_block_size * block_size_factor
320
- )
316
+ tokens_per_block = kv_cache_group.kv_cache_spec.block_size
317
+ assert group_config.tokens_per_block == tokens_per_block
318
+ assert group_config.tokens_per_chunk == tokens_per_block * blocks_per_chunk
319
320
# extract OffloadingSpec of worker_connector
321
connector_worker = self.worker_connector.connector_worker
@@ -389,7 +387,7 @@ class RequestRunner:
387
for block_id in dst_spec.block_ids:
388
self.flushed_gpu_blocks.add(self.gpu_blocks[block_id.item()])
389
392
- block_size_factor = self.block_size_factor
390
+ blocks_per_chunk = self.blocks_per_chunk
391
392
for src_spec, dst_spec in self.offloading_spec.get_completed_transfers():
393
if isinstance(src_spec, GPULoadStoreSpec):
@@ -412,7 +410,7 @@ class RequestRunner:
410
# list of (offload_key, sub_block_offset)
411
offload_addresses: list[Any] = []
412
for offload_key in offload_spec.offload_keys:
415
- for sub_block_idx in range(block_size_factor):
413
+ for sub_block_idx in range(blocks_per_chunk):
414
offload_addresses.append((offload_key, sub_block_idx))
415
416
assert gpu_spec.block_indices is not None
@@ -426,7 +424,7 @@ class RequestRunner:
424
gpu_block_end_offset = gpu_block_offset + group_size
425
assert gpu_block_end_offset <= len(gpu_blocks)
426
429
- offload_addresses_to_skip = logical_offset % block_size_factor
427
+ offload_addresses_to_skip = logical_offset % blocks_per_chunk
428
offload_addresses_end_offset = (
429
offload_address_offset + offload_addresses_to_skip + group_size
430
)
@@ -651,14 +649,14 @@ def request_runner():
649
block_size,
650
num_gpu_blocks,
651
async_scheduling,
654
- block_size_factor=1,
652
+ blocks_per_chunk=1,
653
kv_cache_groups=None,
654
extra_config_overrides=None,
655
):
656
runner = RequestRunner(
657
block_size=block_size,
658
num_gpu_blocks=num_gpu_blocks,
661
- block_size_factor=block_size_factor,
659
+ blocks_per_chunk=blocks_per_chunk,
660
async_scheduling=async_scheduling,
661
kv_cache_groups=kv_cache_groups,
662
extra_config_overrides=extra_config_overrides,
tests/v1/kv_connector/unit/test_offloading_connector.py
+1
-1
@@ -45,7 +45,7 @@ if current_platform.is_cuda():
45
# Falcon-H1: parallel hybrid (every layer has both attention and SSM).
46
# The mamba and attention groups end up with different GPU block sizes
47
# after page-size unification, so we leave cpu_block_size=None
48
- # (block_size_factor stays 1).
48
+ # (blocks_per_chunk stays 1).
49
("tiiuae/Falcon-H1-0.5B-Instruct", None, None, True),
50
]
51
tests/v1/kv_offload/cpu/test_gpu_worker.py
+20
-20
@@ -23,7 +23,7 @@ from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
23
NUM_GPU_BLOCKS = [64]
24
NUM_CPU_BLOCKS = [256]
25
GPU_PAGE_SIZES = [512, 1024]
26
-BLOCK_SIZE_FACTORS = [1, 3]
26
+BLOCKS_PER_CHUNK_VALUES = [1, 3]
27
NUM_TENSORS = [4]
28
SEEDS = [0]
29
DEVICE_TYPE = current_platform.device_type
@@ -35,7 +35,7 @@ NUM_MAPPINGS_PER_GROUP = [2]
35
@pytest.mark.parametrize("gpu_to_cpu", [True, False])
36
@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
37
@pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)
38
-@pytest.mark.parametrize("block_size_factor", BLOCK_SIZE_FACTORS)
38
+@pytest.mark.parametrize("blocks_per_chunk", BLOCKS_PER_CHUNK_VALUES)
39
@pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)
40
@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
41
@pytest.mark.parametrize("num_tensors", NUM_TENSORS)
@@ -48,7 +48,7 @@ def test_transfer(
48
gpu_to_cpu: bool,
49
num_mappings: int,
50
gpu_page_size_bytes: int,
51
- block_size_factor: int,
51
+ blocks_per_chunk: int,
52
num_gpu_blocks: int,
53
num_cpu_blocks: int,
54
num_tensors: int,
@@ -92,7 +92,7 @@ def test_transfer(
92
mmap_region: SharedOffloadRegion | None = None
93
if use_shared_memory:
94
cpu_page_size = round_up(
95
- gpu_page_size_bytes * num_tensors * block_size_factor,
95
+ gpu_page_size_bytes * num_tensors * blocks_per_chunk,
96
SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT,
97
)
98
mmap_region = SharedOffloadRegion(
@@ -105,25 +105,25 @@ def test_transfer(
105
106
worker = CPUOffloadingWorker(
107
kv_caches=kv_caches,
108
- block_size_factor=block_size_factor,
108
+ blocks_per_chunk=blocks_per_chunk,
109
num_cpu_blocks=num_cpu_blocks,
110
mmap_region=mmap_region,
111
)
112
113
# select block mappings
114
- gpu_blocks = random.sample(range(num_gpu_blocks), num_mappings * block_size_factor)
114
+ gpu_blocks = random.sample(range(num_gpu_blocks), num_mappings * blocks_per_chunk)
115
cpu_blocks = random.sample(range(num_cpu_blocks), num_mappings)
116
117
# expand cpu blocks to gpu-page granularity for uniform comparison:
118
- # each cpu block maps to block_size_factor consecutive sub-blocks
118
+ # each cpu block maps to blocks_per_chunk consecutive sub-blocks
119
cpu_blocks_expanded = [
120
- cpu_block * block_size_factor + j
120
+ cpu_block * blocks_per_chunk + j
121
for cpu_block in cpu_blocks
122
- for j in range(block_size_factor)
122
+ for j in range(blocks_per_chunk)
123
]
124
125
# maybe skip some GPU blocks to test reading/writing from the middle of a CPU block
126
- blocks_to_skip = block_size_factor - 1
126
+ blocks_to_skip = blocks_per_chunk - 1
127
if blocks_to_skip > 0:
128
gpu_blocks = gpu_blocks[blocks_to_skip:]
129
cpu_blocks_expanded = cpu_blocks_expanded[blocks_to_skip:]
@@ -214,7 +214,7 @@ def test_transfer(
214
@pytest.mark.parametrize("gpu_to_cpu", [True, False])
215
@pytest.mark.parametrize("num_mappings_per_group", NUM_MAPPINGS_PER_GROUP)
216
@pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)
217
-@pytest.mark.parametrize("block_size_factor", BLOCK_SIZE_FACTORS)
217
+@pytest.mark.parametrize("blocks_per_chunk", BLOCKS_PER_CHUNK_VALUES)
218
@pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)
219
@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
220
@pytest.mark.parametrize("seed", SEEDS)
@@ -225,7 +225,7 @@ def test_transfer_multi_group(
225
gpu_to_cpu: bool,
226
num_mappings_per_group: int,
227
gpu_page_size_bytes: int,
228
- block_size_factor: int,
228
+ blocks_per_chunk: int,
229
num_gpu_blocks: int,
230
num_cpu_blocks: int,
231
seed: int,
@@ -234,7 +234,7 @@ def test_transfer_multi_group(
234
"""Test transfers with three KV cache groups:
235
- Group 0: aligned transfer with num_mappings_per_group blocks
236
- Group 1: zero blocks (empty group)
237
- - Group 2: unaligned CPU->GPU transfer (logical_offset=block_size_factor-1,
237
+ - Group 2: unaligned CPU->GPU transfer (logical_offset=blocks_per_chunk-1,
238
causing the implementation to skip source sub-blocks) with
239
num_mappings_per_group blocks
240
"""
@@ -275,7 +275,7 @@ def test_transfer_multi_group(
275
276
worker = CPUOffloadingWorker(
277
kv_caches=canonical_kv_caches,
278
- block_size_factor=block_size_factor,
278
+ blocks_per_chunk=blocks_per_chunk,
279
num_cpu_blocks=num_cpu_blocks,
280
)
281
@@ -283,7 +283,7 @@ def test_transfer_multi_group(
283
group_sizes_in_cpu_blocks = [num_mappings_per_group, 0, num_mappings_per_group]
284
285
total_cpu_blocks = sum(group_sizes_in_cpu_blocks)
286
- total_gpu_blocks_needed = total_cpu_blocks * block_size_factor
286
+ total_gpu_blocks_needed = total_cpu_blocks * blocks_per_chunk
287
gpu_blocks_all = random.sample(range(num_gpu_blocks), total_gpu_blocks_needed)
288
cpu_blocks_all = random.sample(range(num_cpu_blocks), total_cpu_blocks)
289
@@ -293,7 +293,7 @@ def test_transfer_multi_group(
293
gpu_offset = 0
294
cpu_offset = 0
295
for size in group_sizes_in_cpu_blocks:
296
- gpu_count = size * block_size_factor
296
+ gpu_count = size * blocks_per_chunk
297
gpu_blocks_per_group.append(gpu_blocks_all[gpu_offset : gpu_offset + gpu_count])
298
cpu_blocks_per_group.append(cpu_blocks_all[cpu_offset : cpu_offset + size])
299
gpu_offset += gpu_count
@@ -302,15 +302,15 @@ def test_transfer_multi_group(
302
# expand cpu blocks to gpu-page granularity
303
cpu_blocks_expanded_per_group = [
304
[
305
- cpu_block * block_size_factor + j
305
+ cpu_block * blocks_per_chunk + j
306
for cpu_block in cpu_blocks
307
- for j in range(block_size_factor)
307
+ for j in range(blocks_per_chunk)
308
]
309
for cpu_blocks in cpu_blocks_per_group
310
]
311
312
# skip sub-blocks from group 2 to test unaligned transfers.
313
- sub_blocks_to_skip = block_size_factor - 1 # e.g. 2 when block_size_factor=3
313
+ sub_blocks_to_skip = blocks_per_chunk - 1 # e.g. 2 when blocks_per_chunk=3
314
if sub_blocks_to_skip > 0:
315
gpu_blocks_per_group[2] = gpu_blocks_per_group[2][
316
sub_blocks_to_skip:-sub_blocks_to_skip
@@ -347,7 +347,7 @@ def test_transfer_multi_group(
347
cpu_blocks_expanded_per_group, gpu_blocks_per_group
348
)
349
]
350
- num_dst_sub_blocks = num_cpu_blocks * block_size_factor
350
+ num_dst_sub_blocks = num_cpu_blocks * blocks_per_chunk
351
else:
352
handler = worker._load_handler
353
src_spec = CPULoadStoreSpec(cpu_blocks)
tests/v1/kv_offload/test_factory.py
+329
-13
@@ -11,17 +11,33 @@ These tests verify:
11
4. Error paths — unregistered specs, missing config, duplicate registration.
12
"""
13
14
+from typing import cast
15
+from unittest.mock import MagicMock, patch
16
+
17
import pytest
18
import torch
19
17
-from vllm.config import KVTransferConfig
20
+from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig
21
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
22
+ build_offloading_config,
23
+)
24
+from vllm.platforms import current_platform
25
from vllm.v1.kv_cache_interface import (
26
FullAttentionSpec,
27
KVCacheConfig,
28
KVCacheGroupSpec,
29
KVCacheTensor,
30
+ MLAAttentionSpec,
31
+ SlidingWindowSpec,
32
)
24
-from vllm.v1.kv_offload.base import OffloadingHistogramMetadata, OffloadingSpec
33
+from vllm.v1.kv_offload.base import (
34
+ CanonicalKVCaches,
35
+ OffloadingHistogramMetadata,
36
+ OffloadingManager,
37
+ OffloadingSpec,
38
+ OffloadingWorker,
39
+)
40
+from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
41
from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
42
from vllm.v1.kv_offload.factory import OffloadingSpecFactory
43
from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
@@ -39,6 +55,17 @@ def restore_registry():
55
OffloadingSpecFactory._registry = original
56
57
58
+def _get_extra_config(config: VllmConfig) -> dict:
59
+ assert config.kv_transfer_config is not None
60
+ return config.kv_transfer_config.kv_connector_extra_config
61
+
62
+
63
+def _create_spec(config: VllmConfig, kv_cache_config: KVCacheConfig) -> OffloadingSpec:
64
+ return OffloadingSpecFactory.create_spec(
65
+ build_offloading_config(config, kv_cache_config)
66
+ )
67
+
68
+
69
def _make_vllm_config(
70
spec_name: str | None = "CPUOffloadingSpec",
71
cpu_bytes_to_use: int | None = None,
@@ -95,6 +122,46 @@ def _make_vllm_config(
122
)
123
124
125
+def _make_layout_vllm_config(
126
+ spec_name: str = "CPUOffloadingSpec",
127
+ cpu_bytes_to_use: int | None = None,
128
+ extra_config: dict | None = None,
129
+ tensor_parallel_size: int = 1,
130
+ pipeline_parallel_size: int = 1,
131
+ prefill_context_parallel_size: int = 1,
132
+ decode_context_parallel_size: int = 1,
133
+) -> VllmConfig:
134
+ config = MagicMock()
135
+ config.cache_config.block_size = 16
136
+ config.cache_config.enable_prefix_caching = True
137
+ config.cache_config.prefix_match_unit = None
138
+ config.cache_config.cache_dtype = torch.float16
139
+ config.model_config.model = "test-model"
140
+ world_size = (
141
+ tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size
142
+ )
143
+ with patch.object(current_platform, "device_count", return_value=world_size):
144
+ config.parallel_config = ParallelConfig(
145
+ tensor_parallel_size=tensor_parallel_size,
146
+ pipeline_parallel_size=pipeline_parallel_size,
147
+ prefill_context_parallel_size=prefill_context_parallel_size,
148
+ decode_context_parallel_size=decode_context_parallel_size,
149
+ )
150
+ config.kv_events_config = None
151
+ config.use_v2_model_runner = False
152
+
153
+ connector_extra_config = dict(extra_config or {})
154
+ connector_extra_config["spec_name"] = spec_name
155
+ if cpu_bytes_to_use is not None:
156
+ connector_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use
157
+ config.kv_transfer_config = KVTransferConfig(
158
+ kv_connector="OffloadingConnector",
159
+ kv_role="kv_both",
160
+ kv_connector_extra_config=connector_extra_config,
161
+ )
162
+ return cast(VllmConfig, config)
163
+
164
+
165
def _make_kv_cache_config():
166
"""Build a minimal KVCacheConfig with one KV cache tensor."""
167
num_blocks = 16
@@ -122,6 +189,78 @@ def _make_kv_cache_config():
189
)
190
191
192
+def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig:
193
+ num_blocks = 4
194
+ if packed:
195
+ kv_cache_tensors = [
196
+ KVCacheTensor(
197
+ size=64,
198
+ shared_by=[layer_name],
199
+ block_stride=16,
200
+ )
201
+ for layer_name in ("layer0", "layer1")
202
+ ]
203
+ else:
204
+ kv_cache_tensors = [
205
+ KVCacheTensor(size=40, shared_by=["layer0"]),
206
+ KVCacheTensor(size=24, shared_by=["layer1"]),
207
+ ]
208
+
209
+ return KVCacheConfig(
210
+ num_blocks=num_blocks,
211
+ kv_cache_tensors=kv_cache_tensors,
212
+ kv_cache_groups=[
213
+ KVCacheGroupSpec(
214
+ ["layer0", "layer1"],
215
+ FullAttentionSpec(
216
+ block_size=16,
217
+ num_kv_heads=1,
218
+ head_size=1,
219
+ dtype=torch.float32,
220
+ ),
221
+ )
222
+ ],
223
+ )
224
+
225
+
226
+def _make_hybrid_kv_cache_config() -> KVCacheConfig:
227
+ return KVCacheConfig(
228
+ num_blocks=4,
229
+ kv_cache_tensors=[
230
+ KVCacheTensor(size=40, shared_by=["full_layer"]),
231
+ KVCacheTensor(size=24, shared_by=["mla_layer"]),
232
+ ],
233
+ kv_cache_groups=[
234
+ KVCacheGroupSpec(
235
+ ["full_layer"],
236
+ FullAttentionSpec(
237
+ block_size=12,
238
+ num_kv_heads=1,
239
+ head_size=1,
240
+ dtype=torch.float32,
241
+ ),
242
+ ),
243
+ KVCacheGroupSpec(
244
+ ["mla_layer"],
245
+ MLAAttentionSpec(
246
+ block_size=16,
247
+ num_kv_heads=1,
248
+ head_size=576,
249
+ dtype=torch.float32,
250
+ ),
251
+ ),
252
+ ],
253
+ )
254
+
255
+
256
+class SingleArgExternalOffloadingSpec(OffloadingSpec):
257
+ def get_manager(self) -> OffloadingManager:
258
+ raise NotImplementedError
259
+
260
+ def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker:
261
+ raise NotImplementedError
262
+
263
+
264
# ---------------------------------------------------------------------------
265
# Pre-registration integrity (CI sentinel)
266
# ---------------------------------------------------------------------------
@@ -154,7 +293,7 @@ def test_tiering_spec_registered():
293
def test_get_spec_cls_returns_registered_class():
294
"""Registered spec_name returns correct class."""
295
config = _make_vllm_config(spec_name="CPUOffloadingSpec")
157
- spec_cls = OffloadingSpecFactory.get_spec_cls(config)
296
+ spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
297
assert spec_cls is CPUOffloadingSpec
298
299
@@ -162,7 +301,7 @@ def test_get_spec_cls_default_to_cpu():
301
"""Default spec_name (absent from config) resolves to CPUOffloadingSpec."""
302
config = _make_vllm_config(spec_name=None)
303
config.kv_transfer_config.kv_connector_extra_config.pop("spec_name", None)
165
- spec_cls = OffloadingSpecFactory.get_spec_cls(config)
304
+ spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
305
assert spec_cls is CPUOffloadingSpec
306
307
@@ -176,16 +315,193 @@ def test_create_cpu_offloading_spec_end_to_end():
315
316
Verifies:
317
- cpu_bytes_to_use validation and num_blocks calculation
179
- - block_size % hash_block_size assertion
318
+ - block_size % tokens_per_hash assertion
319
- spec instance is CPUOffloadingSpec
320
"""
321
config = _make_vllm_config(cpu_bytes_to_use=65536)
322
kv_cache_config = _make_kv_cache_config()
184
- spec = OffloadingSpecFactory.create_spec(config, kv_cache_config)
323
+ spec = _create_spec(config, kv_cache_config)
324
assert isinstance(spec, CPUOffloadingSpec)
325
assert spec.num_blocks > 0
326
327
328
+@pytest.mark.parametrize("packed", [False, True])
329
+def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool):
330
+ cpu_bytes_to_use = 1920
331
+ config = _make_layout_vllm_config(
332
+ cpu_bytes_to_use=cpu_bytes_to_use,
333
+ extra_config={"block_size": 32},
334
+ tensor_parallel_size=3,
335
+ pipeline_parallel_size=2,
336
+ )
337
+
338
+ spec = _create_spec(config, _make_sizing_kv_cache_config(packed))
339
+
340
+ assert isinstance(spec, CPUOffloadingSpec)
341
+ assert spec.cpu_page_size_per_worker == 32
342
+ assert spec.kv_bytes_per_chunk == 192
343
+ assert spec.num_blocks == cpu_bytes_to_use // 192
344
+
345
+
346
+def test_cpu_spec_rejects_partially_packed_tensor_layout():
347
+ config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
348
+ kv_cache_config = _make_sizing_kv_cache_config(packed=False)
349
+ kv_cache_config.kv_cache_tensors[0].block_stride = 16
350
+
351
+ with pytest.raises(AssertionError):
352
+ _create_spec(config, kv_cache_config)
353
+
354
+
355
+def test_cpu_spec_zero_blocks_skips_tensor_layout_validation():
356
+ config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
357
+ kv_cache_config = _make_sizing_kv_cache_config(packed=False)
358
+ kv_cache_config.num_blocks = 0
359
+ kv_cache_config.kv_cache_tensors[0].block_stride = 16
360
+
361
+ spec = _create_spec(config, kv_cache_config)
362
+
363
+ assert isinstance(spec, CPUOffloadingSpec)
364
+ assert spec.cpu_page_size_per_worker == 0
365
+ assert spec.kv_bytes_per_chunk == 0
366
+ assert spec.num_blocks == 0
367
+
368
+
369
+def test_tiering_spec_aligns_row_size():
370
+ alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
371
+ cpu_bytes_to_use = alignment * 3
372
+ config = _make_layout_vllm_config(
373
+ spec_name="TieringOffloadingSpec",
374
+ cpu_bytes_to_use=cpu_bytes_to_use,
375
+ extra_config={"block_size": 32},
376
+ tensor_parallel_size=3,
377
+ pipeline_parallel_size=2,
378
+ )
379
+
380
+ spec = _create_spec(config, _make_sizing_kv_cache_config(packed=False))
381
+
382
+ assert isinstance(spec, TieringOffloadingSpec)
383
+ assert spec.cpu_page_size_per_worker == 32
384
+ assert spec.kv_bytes_per_chunk == alignment
385
+ assert spec.num_blocks == cpu_bytes_to_use // alignment
386
+
387
+
388
+def test_offloading_spec_resolves_prefill_context_parallel_block_sizes():
389
+ config = _make_layout_vllm_config(
390
+ cpu_bytes_to_use=65536,
391
+ extra_config={"block_size": 64},
392
+ prefill_context_parallel_size=2,
393
+ )
394
+
395
+ spec = _create_spec(config, _make_kv_cache_config())
396
+
397
+ assert spec.tokens_per_block == (32,)
398
+ assert spec.tokens_per_hash == 32
399
+ assert spec.blocks_per_chunk == 2
400
+
401
+
402
+def test_offloading_config_preserves_data_parallel_index():
403
+ config = _make_layout_vllm_config()
404
+ config.parallel_config.data_parallel_index = 2
405
+
406
+ offloading_config = build_offloading_config(config, _make_kv_cache_config())
407
+
408
+ assert offloading_config.parallel.data_parallel_index == 2
409
+
410
+
411
+def test_offloading_spec_resolves_heterogeneous_hybrid_block_sizes():
412
+ config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
413
+ config.cache_config.block_size = 4
414
+
415
+ spec = _create_spec(config, _make_hybrid_kv_cache_config())
416
+
417
+ assert spec.tokens_per_block == (12, 16)
418
+ assert spec.tokens_per_hash == 4
419
+ assert spec.blocks_per_chunk == 1
420
+
421
+
422
+def _full_attention_spec(block_size: int = 16) -> FullAttentionSpec:
423
+ return FullAttentionSpec(
424
+ block_size=block_size, num_kv_heads=4, head_size=128, dtype=torch.float32
425
+ )
426
+
427
+
428
+def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool:
429
+ config = _make_layout_vllm_config()
430
+ kv_cache_config = KVCacheConfig(
431
+ num_blocks=0, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups
432
+ )
433
+ offloading_config = build_offloading_config(config, kv_cache_config)
434
+ return offloading_config.parallel.is_parallelism_agnostic
435
+
436
+
437
+def test_parallelism_agnostic_for_single_full_attention_group():
438
+ assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())])
439
+
440
+
441
+@pytest.mark.parametrize(
442
+ "kv_cache_groups",
443
+ [
444
+ # MLA latent KV is replicated per rank, never head-sharded.
445
+ [
446
+ KVCacheGroupSpec(
447
+ ["l0"],
448
+ MLAAttentionSpec(
449
+ block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32
450
+ ),
451
+ )
452
+ ],
453
+ # Sliding window is not full attention.
454
+ [
455
+ KVCacheGroupSpec(
456
+ ["l0"],
457
+ SlidingWindowSpec(
458
+ block_size=16,
459
+ num_kv_heads=4,
460
+ head_size=128,
461
+ dtype=torch.float32,
462
+ sliding_window=128,
463
+ ),
464
+ )
465
+ ],
466
+ # Hybrid model: more than one KV cache group.
467
+ [
468
+ KVCacheGroupSpec(["l0"], _full_attention_spec()),
469
+ KVCacheGroupSpec(["l1"], _full_attention_spec()),
470
+ ],
471
+ ],
472
+)
473
+def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
474
+ assert not _parallelism_agnostic(kv_cache_groups)
475
+
476
+
477
+def test_parallelism_agnostic_disabled_on_v2_model_runner():
478
+ config = _make_layout_vllm_config()
479
+ config.use_v2_model_runner = True
480
+ kv_cache_config = KVCacheConfig(
481
+ num_blocks=0,
482
+ kv_cache_tensors=[],
483
+ kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())],
484
+ )
485
+ offloading_config = build_offloading_config(config, kv_cache_config)
486
+ assert not offloading_config.parallel.is_parallelism_agnostic
487
+
488
+
489
+def test_create_dynamic_spec_receives_translated_config():
490
+ config = _make_layout_vllm_config(
491
+ spec_name="SingleArgExternalOffloadingSpec",
492
+ extra_config={
493
+ "spec_module_path": "tests.v1.kv_offload.test_factory",
494
+ },
495
+ )
496
+ kv_cache_config = _make_kv_cache_config()
497
+ offloading_config = build_offloading_config(config, kv_cache_config)
498
+
499
+ spec = OffloadingSpecFactory.create_spec(offloading_config)
500
+
501
+ assert isinstance(spec, SingleArgExternalOffloadingSpec)
502
+ assert spec.config is offloading_config
503
+
504
+
505
# ---------------------------------------------------------------------------
506
# Dynamic import via spec_module_path
507
# ---------------------------------------------------------------------------
@@ -205,7 +521,7 @@ def test_dynamic_load_via_spec_module_path():
521
config.kv_transfer_config.kv_connector_extra_config["spec_module_path"] = (
522
"vllm.v1.kv_offload.cpu.spec"
523
)
208
- spec_cls = OffloadingSpecFactory.get_spec_cls(config)
524
+ spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
525
assert spec_cls is CPUOffloadingSpec
526
527
@@ -218,12 +534,12 @@ def test_unregistered_spec_without_module_path_raises():
534
"""spec_name not in registry + no spec_module_path → ValueError."""
535
config = _make_vllm_config(spec_name="NonexistentSpec")
536
with pytest.raises(ValueError, match="Unsupported spec type"):
221
- OffloadingSpecFactory.get_spec_cls(config)
537
+ OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
538
539
# create_spec should also fail (calls get_spec_cls internally)
540
kv_cache_config = _make_kv_cache_config()
541
with pytest.raises(ValueError, match="Unsupported spec type"):
226
- OffloadingSpecFactory.create_spec(config, kv_cache_config)
542
+ _create_spec(config, kv_cache_config)
543
544
545
def test_cpu_spec_missing_cpu_bytes_to_use_raises():
@@ -232,7 +548,7 @@ def test_cpu_spec_missing_cpu_bytes_to_use_raises():
548
config.kv_transfer_config.kv_connector_extra_config.pop("cpu_bytes_to_use", None)
549
kv_cache_config = _make_kv_cache_config()
550
with pytest.raises(Exception, match="cpu_bytes_to_use must be specified"):
235
- OffloadingSpecFactory.create_spec(config, kv_cache_config)
551
+ _create_spec(config, kv_cache_config)
552
553
554
def test_duplicate_registration_raises():
@@ -253,7 +569,7 @@ def test_build_metric_definitions_below_threshold():
569
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
570
571
config = _make_vllm_config(store_threshold=1)
256
- spec_cls = OffloadingSpecFactory.get_spec_cls(config)
572
+ spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
573
metrics = spec_cls.build_metric_definitions(
574
config.kv_transfer_config.kv_connector_extra_config
575
)
@@ -266,7 +582,7 @@ def test_build_metric_definitions_allocation_size_histogram():
582
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
583
584
config = _make_vllm_config(store_threshold=0)
269
- spec_cls = OffloadingSpecFactory.get_spec_cls(config)
585
+ spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
586
metrics = spec_cls.build_metric_definitions(
587
config.kv_transfer_config.kv_connector_extra_config
588
)
@@ -291,7 +607,7 @@ def test_build_metric_definitions_returns_counter_at_threshold():
607
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
608
609
config = _make_vllm_config(store_threshold=2)
294
- spec_cls = OffloadingSpecFactory.get_spec_cls(config)
610
+ spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
611
metrics = spec_cls.build_metric_definitions(
612
config.kv_transfer_config.kv_connector_extra_config
613
)
tests/v1/kv_offload/test_file_mapper.py
+94
-132
@@ -4,80 +4,60 @@
4
5
from unittest.mock import MagicMock
6
7
-import torch
8
-
9
-from vllm.v1.kv_cache_interface import (
10
- FullAttentionSpec,
11
- KVCacheGroupSpec,
12
- MLAAttentionSpec,
13
- SlidingWindowSpec,
14
-)
15
-from vllm.v1.kv_offload.base import (
16
- OffloadingSpec,
17
- make_offload_key,
7
+from vllm.v1.kv_offload.base import OffloadingSpec, make_offload_key
8
+from vllm.v1.kv_offload.config import (
9
+ OffloadingCacheConfig,
10
+ OffloadingConfig,
11
+ OffloadingGroupConfig,
12
+ OffloadingModelConfig,
13
+ OffloadingParallelConfig,
14
)
15
from vllm.v1.kv_offload.file_mapper import FileMapper
16
21
-# ---------------------------------------------------------------------------
22
-# Shared mocks (mirrors test_fs_tier.py pattern)
23
-# ---------------------------------------------------------------------------
24
-
25
-_MOCK_VLLM_CONFIG = MagicMock()
26
-_MOCK_VLLM_CONFIG.model_config.model = "test-model"
27
-_MOCK_VLLM_CONFIG.cache_config.block_size = 16
28
-_MOCK_VLLM_CONFIG.cache_config.cache_dtype = "torch.float32"
29
-_MOCK_VLLM_CONFIG.parallel_config.tensor_parallel_size = 1
30
-_MOCK_VLLM_CONFIG.parallel_config.pipeline_parallel_size = 1
31
-_MOCK_VLLM_CONFIG.parallel_config.prefill_context_parallel_size = 1
32
-_MOCK_VLLM_CONFIG.parallel_config.decode_context_parallel_size = 1
33
-_MOCK_VLLM_CONFIG.parallel_config.rank = 0
34
-
35
-_MOCK_KV_CACHE_CONFIG = MagicMock()
36
-_MOCK_KV_CACHE_CONFIG.kv_cache_groups = []
37
-
38
-_MOCK_OFFLOADING_SPEC = MagicMock(spec=OffloadingSpec)
39
-_MOCK_OFFLOADING_SPEC.vllm_config = _MOCK_VLLM_CONFIG
40
-_MOCK_OFFLOADING_SPEC.kv_cache_config = _MOCK_KV_CACHE_CONFIG
41
-_MOCK_OFFLOADING_SPEC.block_size_factor = 1
42
-
43
-
17
# ---------------------------------------------------------------------------
18
# Helper
19
# ---------------------------------------------------------------------------
20
21
22
def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
50
- """Helper to create FileMapper with customizable mock config."""
51
- # Create a copy of the mock config to avoid modifying the global one
52
- mock_vllm_config = MagicMock()
53
- mock_vllm_config.model_config.model = kwargs.get("model_name", "test-model")
54
- mock_vllm_config.cache_config.block_size = kwargs.get("hash_block_size", 16)
55
- mock_vllm_config.cache_config.cache_dtype = (
56
- f"torch.{kwargs.get('dtype', 'float16')}"
57
- )
58
- mock_vllm_config.parallel_config.tensor_parallel_size = kwargs.get("tp_size", 1)
59
- mock_vllm_config.parallel_config.pipeline_parallel_size = kwargs.get("pp_size", 1)
60
- mock_vllm_config.parallel_config.prefill_context_parallel_size = kwargs.get(
61
- "pcp_size", 1
62
- )
63
- mock_vllm_config.parallel_config.decode_context_parallel_size = kwargs.get(
64
- "dcp_size", 1
23
+ """Build a FileMapper from a mocked spec carrying a hand-built config."""
24
+ config = OffloadingConfig(
25
+ groups=tuple(
26
+ OffloadingGroupConfig(
27
+ tokens_per_block=tokens_per_block,
28
+ layer_names=(layer_name,),
29
+ )
30
+ for tokens_per_block, layer_name in kwargs.get("groups", ())
31
+ ),
32
+ worker_kv_bytes_per_block=0,
33
+ enable_kv_cache_events=False,
34
+ extra_config={},
35
+ engine_id="test-engine",
36
+ model=OffloadingModelConfig(
37
+ name=kwargs.get("model_name", "test-model"),
38
+ dtype=kwargs.get("dtype", "float16"),
39
+ ),
40
+ cache=OffloadingCacheConfig(
41
+ tokens_per_hash=kwargs.get("tokens_per_hash", 16),
42
+ blocks_per_chunk=kwargs.get("blocks_per_chunk", 1),
43
+ ),
44
+ parallel=OffloadingParallelConfig(
45
+ rank=kwargs.get("rank", 0),
46
+ world_size=kwargs.get("world_size", 1),
47
+ tp_size=kwargs.get("tp_size", 1),
48
+ pp_size=kwargs.get("pp_size", 1),
49
+ pcp_size=kwargs.get("pcp_size", 1),
50
+ dcp_size=kwargs.get("dcp_size", 1),
51
+ data_parallel_index=0,
52
+ is_parallelism_agnostic=kwargs.get("is_parallelism_agnostic", False),
53
+ ),
54
)
66
- mock_vllm_config.parallel_config.rank = kwargs.get("rank", 0)
67
- mock_vllm_config.use_v2_model_runner = kwargs.get("use_v2_model_runner", False)
68
-
69
- mock_kv_cache_config = MagicMock()
70
- mock_kv_cache_config.kv_cache_groups = kwargs.get("kv_cache_groups", [])
71
-
72
- mock_offloading_spec = MagicMock(spec=OffloadingSpec)
73
- mock_offloading_spec.vllm_config = mock_vllm_config
74
- mock_offloading_spec.kv_cache_config = mock_kv_cache_config
75
- mock_offloading_spec.block_size_factor = kwargs.get("block_size_factor", 1)
76
-
55
+ spec = MagicMock(spec=OffloadingSpec)
56
+ spec.config = config
57
return FileMapper.from_offloading_spec(
58
root_dir=kwargs.get("root_dir", "/tmp/cache"),
79
- offloading_spec=mock_offloading_spec,
80
- gpu_blocks_per_file=mock_offloading_spec.block_size_factor,
59
+ offloading_spec=spec,
60
+ blocks_per_file=config.cache.blocks_per_chunk,
61
parallel_agnostic=kwargs.get("parallel_agnostic", False),
62
)
63
@@ -92,7 +72,7 @@ def test_get_file_name_full_structure():
72
Path must match: <base_path>_r<rank>/<hhh>/<hh>_g<group_idx>/<hash_hex>.bin
73
74
Concretely:
95
- - The segment immediately after base_path must end with `_r0`
75
+ - The segment immediately after base_path must end with `_r3`
76
- The next segment is the first 3 hex chars of the block hash
77
- The next segment is <2 hex chars>_g<group_idx>
78
- The final segment is <full hash hex>.bin
@@ -105,7 +85,7 @@ def test_get_file_name_full_structure():
85
path = fm.get_file_name(key)
86
87
expected_path = (
108
- "/tmp/cache/test-model_588656ebcc66_r3/000/10_g2/0001020304050607.bin"
88
+ "/tmp/cache/test-model_42b94bdc9933_r3/000/10_g2/0001020304050607.bin"
89
)
90
assert path == expected_path
91
@@ -114,19 +94,30 @@ def test_get_run_config_fields():
94
fm = make_mapper_from_offloading_spec(
95
model_name="my-model",
96
dtype="bfloat16",
117
- tp_size=2,
97
+ tp_size=4,
98
+ pp_size=3,
99
+ pcp_size=2,
100
+ dcp_size=2,
101
+ groups=((64, "layer0"),),
102
+ tokens_per_hash=64,
103
+ blocks_per_chunk=3,
104
)
105
cfg = fm.get_run_config()
106
assert cfg == {
107
"model_name": "my-model",
122
- "hash_block_size": 16,
123
- "gpu_blocks_per_file": 1,
124
- "tp_size": 2,
125
- "pp_size": 1,
126
- "pcp_size": 1,
127
- "dcp_size": 1,
108
+ "tokens_per_hash": 64,
109
+ "blocks_per_file": 3,
110
+ "tp_size": 4,
111
+ "pp_size": 3,
112
+ "pcp_size": 2,
113
+ "dcp_size": 2,
114
"dtype": "bfloat16",
129
- "kv_cache_groups": [],
115
+ "kv_cache_groups": [
116
+ {
117
+ "tokens_per_block": 64,
118
+ "layer_names": ["layer0"],
119
+ }
120
+ ],
121
"inference_engine": "vllm",
122
}
123
@@ -137,90 +128,61 @@ def test_get_config_file_path():
128
assert config_path == f"{fm.base_path}/config.json"
129
130
140
-# ---------------------------------------------------------------------------
141
-# parallel_agnostic: honored only for a single non-MLA full-attention group
142
-# ---------------------------------------------------------------------------
143
-
144
-
145
-def _full_attention_group() -> KVCacheGroupSpec:
146
- return KVCacheGroupSpec(
147
- layer_names=["layer0"],
148
- kv_cache_spec=FullAttentionSpec(
149
- block_size=16, num_kv_heads=4, head_size=128, dtype=torch.float32
150
- ),
131
+def test_hybrid_file_identity_uses_resolved_tokens_per_hash():
132
+ # For heterogeneous groups the namespace records the resolved hash
133
+ # granularity (GCD of the group block sizes), which is the actual
134
+ # granularity of the offload block hashes.
135
+ fm = make_mapper_from_offloading_spec(
136
+ groups=((12, "full_layer"), (16, "mla_layer")),
137
+ tokens_per_hash=4,
138
)
139
+ assert fm.fields["tokens_per_hash"] == 4
140
+ assert fm.fields["kv_cache_groups"] == [
141
+ {"tokens_per_block": 12, "layer_names": ["full_layer"]},
142
+ {"tokens_per_block": 16, "layer_names": ["mla_layer"]},
143
+ ]
144
145
154
-def _sliding_window_group() -> KVCacheGroupSpec:
155
- return KVCacheGroupSpec(
156
- layer_names=["layer0"],
157
- kv_cache_spec=SlidingWindowSpec(
158
- block_size=16,
159
- num_kv_heads=4,
160
- head_size=128,
161
- dtype=torch.float32,
162
- sliding_window=128,
163
- ),
164
- )
146
+# ---------------------------------------------------------------------------
147
+# parallel_agnostic: opt-in honored only when the config marks the layout
148
+# parallelism-agnostic (predicate computation is covered in test_factory.py)
149
+# ---------------------------------------------------------------------------
150
151
167
-def test_parallel_agnostic_enabled_for_single_full_attention():
168
- # tp/rank are collapsed out of the namespace so the cache is shared
169
- # across tensor-parallel sizes.
152
+def test_parallel_agnostic_collapses_namespace_when_config_allows():
153
fm = make_mapper_from_offloading_spec(
171
- tp_size=2,
154
+ tp_size=4,
155
+ pp_size=3,
156
+ pcp_size=2,
157
+ dcp_size=2,
158
rank=1,
173
- kv_cache_groups=[_full_attention_group()],
159
+ is_parallelism_agnostic=True,
160
parallel_agnostic=True,
161
)
162
assert fm.fields["tp_size"] == 1
163
+ assert fm.fields["pp_size"] == 1
164
+ assert fm.fields["pcp_size"] == 1
165
+ assert fm.fields["dcp_size"] == 1
166
assert fm.rank == 0
167
168
180
-def test_parallel_agnostic_disabled_for_multiple_groups():
181
- # More than one KV-cache group (hybrid model) => keep per-layout namespacing.
169
+def test_parallel_agnostic_ignored_when_config_disallows():
170
fm = make_mapper_from_offloading_spec(
171
tp_size=2,
184
- kv_cache_groups=[_full_attention_group(), _full_attention_group()],
185
- parallel_agnostic=True,
186
- )
187
- assert fm.fields["tp_size"] == 2
188
-
189
-
190
-def test_parallel_agnostic_disabled_for_non_full_attention():
191
- # Single group but not full attention (sliding window) => keep namespacing.
192
- fm = make_mapper_from_offloading_spec(
193
- tp_size=2,
194
- kv_cache_groups=[_sliding_window_group()],
172
+ rank=1,
173
+ is_parallelism_agnostic=False,
174
parallel_agnostic=True,
175
)
176
assert fm.fields["tp_size"] == 2
198
-
199
-
200
-def test_parallel_agnostic_excludes_mla():
201
- # MLA latent KV is replicated per rank, so its offloaded blocks are not
202
- # parallelism-invariant: the opt-in must not collapse tp/rank.
203
- group = KVCacheGroupSpec(
204
- layer_names=["layer0"],
205
- kv_cache_spec=MLAAttentionSpec(
206
- block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32
207
- ),
208
- )
209
- fm = make_mapper_from_offloading_spec(
210
- tp_size=2, rank=1, kv_cache_groups=[group], parallel_agnostic=True
211
- )
212
- assert fm.fields["tp_size"] == 2
177
assert fm.rank == 1
178
179
216
-def test_parallel_agnostic_disabled_on_v2_model_runner():
217
- # V2's KV layout is not known to be parallelism-invariant: don't collapse.
180
+def test_namespace_kept_without_parallel_agnostic_opt_in():
181
fm = make_mapper_from_offloading_spec(
182
tp_size=2,
183
rank=1,
221
- kv_cache_groups=[_full_attention_group()],
222
- use_v2_model_runner=True,
223
- parallel_agnostic=True,
184
+ is_parallelism_agnostic=True,
185
+ parallel_agnostic=False,
186
)
187
assert fm.fields["tp_size"] == 2
188
assert fm.rank == 1
tests/v1/kv_offload/tiering/p2p/test_data_transport.py
+1
-1
@@ -47,7 +47,7 @@ class TestDataTransportBase:
47
def test_config_fingerprint_deterministic(self):
48
"""Same config fields → same fingerprint."""
49
view = self._make_view()
50
- fields = {"model": "llama", "dtype": "float16", "block_size_factor": 1}
50
+ fields = {"model": "llama", "dtype": "float16", "blocks_per_chunk": 1}
51
with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None):
52
t1 = NixlTransport("test:1", view, config_fields=fields)
53
t2 = NixlTransport("test:2", view, config_fields=fields)
tests/v1/kv_offload/tiering/p2p/test_manager.py
+3
-3
@@ -1419,9 +1419,9 @@ class TestBindHostPortDefaults:
1419
or SimpleNamespace(),
1420
)
1421
spec = SimpleNamespace(
1422
- block_size_factor=1,
1423
- vllm_config=SimpleNamespace(
1424
- parallel_config=SimpleNamespace(data_parallel_index=dp_index)
1422
+ blocks_per_chunk=1,
1423
+ config=SimpleNamespace(
1424
+ parallel=SimpleNamespace(data_parallel_index=dp_index)
1425
),
1426
)
1427
mgr = P2PSecondaryTierManager(spec, memoryview(b""), **kwargs)
tests/v1/kv_offload/tiering/test_fs_tier.py
+34
-22
@@ -22,11 +22,18 @@ from vllm.distributed.kv_events import MEDIUM_FS
22
from vllm.v1.kv_offload.base import (
23
LookupResult,
24
OffloadingEvent,
25
+ OffloadingKVEventsConfig,
26
OffloadKey,
27
ReqContext,
28
ScheduleEndContext,
29
make_offload_key,
30
)
31
+from vllm.v1.kv_offload.config import (
32
+ OffloadingCacheConfig,
33
+ OffloadingConfig,
34
+ OffloadingModelConfig,
35
+ OffloadingParallelConfig,
36
+)
37
from vllm.v1.kv_offload.tiering.base import JobMetadata
38
from vllm.v1.kv_offload.tiering.fs.manager import (
39
FileSystemTierManager,
@@ -41,35 +48,40 @@ _BLOCK_ELEMENTS = 128 * mmap.PAGESIZE # 2MB per block for pagesize 4096.
48
_DTYPE: torch.dtype = torch.float32
49
_CTX = ReqContext(req_id="test")
50
44
-_MOCK_VLLM_CONFIG = MagicMock()
45
-_MOCK_VLLM_CONFIG.model_config.model = "test-model"
46
-_MOCK_VLLM_CONFIG.cache_config.block_size = 16
47
-_MOCK_VLLM_CONFIG.cache_config.cache_dtype = "torch.float32"
48
-_MOCK_VLLM_CONFIG.parallel_config.tensor_parallel_size = 1
49
-_MOCK_VLLM_CONFIG.parallel_config.pipeline_parallel_size = 1
50
-_MOCK_VLLM_CONFIG.parallel_config.prefill_context_parallel_size = 1
51
-_MOCK_VLLM_CONFIG.parallel_config.decode_context_parallel_size = 1
52
-_MOCK_VLLM_CONFIG.parallel_config.rank = 0
53
-
54
-_MOCK_KV_CACHE_CONFIG = MagicMock()
55
-_MOCK_KV_CACHE_CONFIG.kv_cache_groups = []
56
-
57
-_MOCK_OFFLOADING_SPEC = MagicMock()
58
-_MOCK_OFFLOADING_SPEC.vllm_config = _MOCK_VLLM_CONFIG
59
-_MOCK_OFFLOADING_SPEC.kv_cache_config = _MOCK_KV_CACHE_CONFIG
60
-_MOCK_OFFLOADING_SPEC.block_size_factor = 1
61
-
51
52
def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock:
53
"""Mock spec with an explicit global KV events flag."""
54
spec = MagicMock()
66
- spec.vllm_config = _MOCK_VLLM_CONFIG
67
- spec.kv_cache_config = _MOCK_KV_CACHE_CONFIG
68
- spec.block_size_factor = 1
69
- spec.kv_events_config.enable_kv_cache_events = enable_kv_cache_events
55
+ spec.config = OffloadingConfig(
56
+ groups=(),
57
+ worker_kv_bytes_per_block=0,
58
+ enable_kv_cache_events=enable_kv_cache_events,
59
+ extra_config={},
60
+ engine_id="test-engine",
61
+ model=OffloadingModelConfig(name="test-model", dtype="float32"),
62
+ cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1),
63
+ parallel=OffloadingParallelConfig(
64
+ rank=0,
65
+ world_size=1,
66
+ tp_size=1,
67
+ pp_size=1,
68
+ pcp_size=1,
69
+ dcp_size=1,
70
+ data_parallel_index=0,
71
+ is_parallelism_agnostic=False,
72
+ ),
73
+ )
74
+ spec.blocks_per_chunk = 1
75
+ spec.kv_events_config = OffloadingKVEventsConfig(
76
+ enable_kv_cache_events=enable_kv_cache_events,
77
+ self_describing_kv_events=False,
78
+ )
79
return spec
80
81
82
+_MOCK_OFFLOADING_SPEC = _make_offloading_spec(enable_kv_cache_events=False)
83
+
84
+
85
def key(n: int) -> OffloadKey:
86
return make_offload_key(n.to_bytes(8, "big"), 0)
87
tests/v1/kv_offload/tiering/test_obj_tier.py
+30
-15
@@ -19,11 +19,18 @@ import torch
19
20
from vllm.v1.kv_offload.base import (
21
LookupResult,
22
+ OffloadingKVEventsConfig,
23
OffloadKey,
24
ReqContext,
25
ScheduleEndContext,
26
make_offload_key,
27
)
28
+from vllm.v1.kv_offload.config import (
29
+ OffloadingCacheConfig,
30
+ OffloadingConfig,
31
+ OffloadingModelConfig,
32
+ OffloadingParallelConfig,
33
+)
34
from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
35
from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig
36
from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager
@@ -33,24 +40,30 @@ from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManag
40
# ---------------------------------------------------------------------------
41
42
36
-def _make_vllm_config():
37
- return SimpleNamespace(
38
- model_config=SimpleNamespace(model="test/model"),
39
- cache_config=SimpleNamespace(block_size=16, cache_dtype="float16"),
40
- parallel_config=SimpleNamespace(
41
- tensor_parallel_size=1,
42
- pipeline_parallel_size=1,
43
- prefill_context_parallel_size=1,
44
- decode_context_parallel_size=1,
43
+def _make_offloading_config(enable_kv_cache_events: bool) -> OffloadingConfig:
44
+ return OffloadingConfig(
45
+ groups=(),
46
+ worker_kv_bytes_per_block=0,
47
+ enable_kv_cache_events=enable_kv_cache_events,
48
+ extra_config={},
49
+ engine_id="test-engine",
50
+ model=OffloadingModelConfig(name="test/model", dtype="float16"),
51
+ cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1),
52
+ parallel=OffloadingParallelConfig(
53
rank=0,
54
+ world_size=1,
55
+ tp_size=1,
56
+ pp_size=1,
57
+ pcp_size=1,
58
+ dcp_size=1,
59
+ data_parallel_index=0,
60
+ is_parallelism_agnostic=False,
61
),
47
- use_v2_model_runner=False,
62
)
63
64
65
_OFFLOADING_SPEC = SimpleNamespace(
52
- vllm_config=_make_vllm_config(),
53
- kv_cache_config=SimpleNamespace(kv_cache_groups=[]),
66
+ config=_make_offloading_config(enable_kv_cache_events=False),
67
)
68
69
_STORE_CONFIG = {
@@ -182,9 +195,11 @@ class MockNixlAgent:
195
def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace:
196
"""Offloading spec stub with an explicit global KV events flag."""
197
return SimpleNamespace(
185
- vllm_config=_make_vllm_config(),
186
- kv_cache_config=SimpleNamespace(kv_cache_groups=[]),
187
- kv_events_config=SimpleNamespace(enable_kv_cache_events=enable_kv_cache_events),
198
+ config=_make_offloading_config(enable_kv_cache_events),
199
+ kv_events_config=OffloadingKVEventsConfig(
200
+ enable_kv_cache_events=enable_kv_cache_events,
201
+ self_describing_kv_events=False,
202
+ ),
203
)
204
205
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
new
+131
@@ -0,0 +1,131 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Translate vLLM KV cache metadata for native offloading backends."""
4
+
5
+from typing import TYPE_CHECKING
6
+
7
+from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
8
+from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
9
+from vllm.v1.kv_offload.config import (
10
+ OffloadingCacheConfig,
11
+ OffloadingConfig,
12
+ OffloadingGroupConfig,
13
+ OffloadingModelConfig,
14
+ OffloadingParallelConfig,
15
+)
16
+
17
+if TYPE_CHECKING:
18
+ from vllm.config import VllmConfig
19
+ from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheTensor
20
+
21
+
22
+def is_kv_cache_tensor_packed(kv_cache_tensor: "KVCacheTensor") -> bool:
23
+ """Return whether a KV cache tensor uses a packed block stride."""
24
+ return bool(kv_cache_tensor.block_stride)
25
+
26
+
27
+def build_offloading_config(
28
+ vllm_config: "VllmConfig",
29
+ kv_cache_config: "KVCacheConfig",
30
+) -> OffloadingConfig:
31
+ """Translate vLLM configuration into the native offloading boundary."""
32
+ kv_transfer_config = vllm_config.kv_transfer_config
33
+ assert kv_transfer_config is not None
34
+ extra_config = kv_transfer_config.kv_connector_extra_config
35
+ assert kv_transfer_config.engine_id is not None
36
+ engine_id = kv_transfer_config.engine_id
37
+
38
+ parallel_config = vllm_config.parallel_config
39
+ context_parallel_factor = (
40
+ parallel_config.decode_context_parallel_size
41
+ * parallel_config.prefill_context_parallel_size
42
+ )
43
+ groups = tuple(
44
+ OffloadingGroupConfig(
45
+ tokens_per_block=(group.kv_cache_spec.block_size * context_parallel_factor),
46
+ layer_names=tuple(group.layer_names),
47
+ )
48
+ for group in kv_cache_config.kv_cache_groups
49
+ )
50
+
51
+ _, tokens_per_hash = resolve_kv_cache_block_sizes(kv_cache_config, vllm_config)
52
+ for group in groups:
53
+ assert group.tokens_per_block % tokens_per_hash == 0, (
54
+ f"tokens_per_block={group.tokens_per_block} not divisible by "
55
+ f"tokens_per_hash={tokens_per_hash}. "
56
+ f"Hybrid models (e.g. Mamba+Attention) need "
57
+ f"--enable-prefix-caching to align block sizes."
58
+ )
59
+
60
+ blocks_per_chunk = 1
61
+ tokens_per_chunk = extra_config.get("block_size")
62
+ if tokens_per_chunk is not None:
63
+ tokens_per_chunk_int = int(tokens_per_chunk)
64
+ unique_tokens_per_block = {group.tokens_per_block for group in groups}
65
+ assert len(unique_tokens_per_block) == 1, (
66
+ "If 'block_size' is specified in kv_connector_extra_config, "
67
+ "there must be at least one KV cache group, "
68
+ "and all groups must have the same block size."
69
+ )
70
+ tokens_per_block = unique_tokens_per_block.pop()
71
+ assert tokens_per_chunk_int % tokens_per_block == 0
72
+ blocks_per_chunk = tokens_per_chunk_int // tokens_per_block
73
+
74
+ worker_kv_bytes_per_block = 0
75
+ if kv_cache_config.num_blocks > 0:
76
+ packed_tensors = tuple(
77
+ is_kv_cache_tensor_packed(tensor)
78
+ for tensor in kv_cache_config.kv_cache_tensors
79
+ )
80
+ is_packed = any(packed_tensors)
81
+ assert not is_packed or all(packed_tensors)
82
+ total_gpu_kv_bytes = (
83
+ kv_cache_config.kv_cache_tensors[0].size
84
+ if is_packed
85
+ else sum(tensor.size for tensor in kv_cache_config.kv_cache_tensors)
86
+ )
87
+ worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks
88
+
89
+ # Only a single non-MLA full-attention group is parallelism-invariant:
90
+ # MLA latent KV is replicated per rank (never head-sharded), and the V2
91
+ # model runner's KV layout is not known to be parallelism-invariant.
92
+ single_group = (
93
+ kv_cache_config.kv_cache_groups[0].kv_cache_spec
94
+ if len(kv_cache_config.kv_cache_groups) == 1
95
+ else None
96
+ )
97
+ is_parallelism_agnostic = (
98
+ not vllm_config.use_v2_model_runner
99
+ and single_group is not None
100
+ and isinstance(single_group, FullAttentionSpec)
101
+ and not isinstance(single_group, MLAAttentionSpec)
102
+ )
103
+
104
+ kv_events_config = vllm_config.kv_events_config
105
+ return OffloadingConfig(
106
+ groups=groups,
107
+ worker_kv_bytes_per_block=worker_kv_bytes_per_block,
108
+ enable_kv_cache_events=(
109
+ kv_events_config is not None and kv_events_config.enable_kv_cache_events
110
+ ),
111
+ extra_config=extra_config,
112
+ engine_id=engine_id,
113
+ model=OffloadingModelConfig(
114
+ name=vllm_config.model_config.model,
115
+ dtype=str(vllm_config.cache_config.cache_dtype).replace("torch.", ""),
116
+ ),
117
+ cache=OffloadingCacheConfig(
118
+ tokens_per_hash=tokens_per_hash,
119
+ blocks_per_chunk=blocks_per_chunk,
120
+ ),
121
+ parallel=OffloadingParallelConfig(
122
+ rank=parallel_config.rank,
123
+ world_size=parallel_config.world_size,
124
+ tp_size=parallel_config.tensor_parallel_size,
125
+ pp_size=parallel_config.pipeline_parallel_size,
126
+ pcp_size=parallel_config.prefill_context_parallel_size,
127
+ dcp_size=parallel_config.decode_context_parallel_size,
128
+ data_parallel_index=parallel_config.data_parallel_index,
129
+ is_parallelism_agnostic=is_parallelism_agnostic,
130
+ ),
131
+ )
vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
+12
-12
@@ -101,7 +101,7 @@ class OffloadingEventsTracker:
101
self,
102
req: Request,
103
group_config: "GroupOffloadConfig",
104
- offload_block_idx: int,
104
+ chunk_idx: int,
105
offload_key: OffloadKey,
106
) -> None:
107
"""Snapshot the KV cache event payload for one offloaded chunk.
@@ -111,9 +111,9 @@ class OffloadingEventsTracker:
111
"""
112
if not self.self_describing_enabled:
113
return
114
- if group_config.sliding_window_size_in_blocks is not None:
114
+ if group_config.sliding_window_size_in_chunks is not None:
115
return
116
- meta = self._build_event_metadata(req, group_config, offload_block_idx)
116
+ meta = self._build_event_metadata(req, group_config, chunk_idx)
117
self._pending_event_metadata[offload_key] = meta
118
119
def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]:
@@ -142,19 +142,19 @@ class OffloadingEventsTracker:
142
self,
143
req: Request,
144
group_config: "GroupOffloadConfig",
145
- offload_block_idx: int,
145
+ chunk_idx: int,
146
) -> _OffloadEventMetadata:
147
"""Build the payload snapshot for one offloaded chunk: its
148
constituent per-block hashes, the whole chunk's tokens, and the
149
per-block ``block_size``."""
150
- hbf = group_config.hash_block_size_factor
150
+ hbf = group_config.hashes_per_chunk
151
assert hbf > 0
152
- assert offload_block_idx >= 0
152
+ assert chunk_idx >= 0
153
# per-block token count (= the GPU/hash block size)
154
- sub_block_size = group_config.offloaded_block_size // hbf
154
+ tokens_per_hash = group_config.tokens_per_chunk // hbf
155
# chunk c covers hash-blocks [c*hbf, (c+1)*hbf); its tail block's hash
156
# is the chunk's OffloadKey.
157
- first_hash_idx = offload_block_idx * hbf
157
+ first_hash_idx = chunk_idx * hbf
158
last_hash_idx = first_hash_idx + hbf
159
assert first_hash_idx >= 0
160
assert last_hash_idx <= len(req.block_hashes)
@@ -164,7 +164,7 @@ class OffloadingEventsTracker:
164
chunk_hashes.append(block_hash)
165
assert len(chunk_hashes) == hbf
166
167
- if group_config.sliding_window_size_in_blocks is not None:
167
+ if group_config.sliding_window_size_in_chunks is not None:
168
# record_store filters these out before calling this helper.
169
raise AssertionError("self-describing events only support full attention")
170
@@ -175,8 +175,8 @@ class OffloadingEventsTracker:
175
parent_block_hash = req.block_hashes[first_hash_idx - 1]
176
assert parent_block_hash is not None
177
178
- tok_start = offload_block_idx * group_config.offloaded_block_size
179
- tok_end = tok_start + group_config.offloaded_block_size
178
+ tok_start = chunk_idx * group_config.tokens_per_chunk
179
+ tok_end = tok_start + group_config.tokens_per_chunk
180
assert tok_end <= len(req.all_token_ids)
181
token_ids = tuple(req.all_token_ids[tok_start:tok_end])
182
@@ -190,7 +190,7 @@ class OffloadingEventsTracker:
190
block_hashes=tuple(chunk_hashes),
191
parent_block_hash=parent_block_hash,
192
token_ids=token_ids,
193
- block_size=sub_block_size,
193
+ block_size=tokens_per_hash,
194
lora_id=lora_id,
195
lora_name=lora_name,
196
extra_keys=None,
vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py
+1
-1
@@ -321,10 +321,10 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
321
self.histogram_transfer_size: dict[tuple[int, str], PromMetricT] = {}
322
self.counter_kv_bytes: dict[tuple[int, str], PromMetricT] = {}
323
self.counter_kv_transfer_time: dict[tuple[int, str], PromMetricT] = {}
324
- spec_cls = OffloadingSpecFactory.get_spec_cls(vllm_config)
324
kv_transfer_config = vllm_config.kv_transfer_config
325
assert kv_transfer_config is not None
326
extra_config = kv_transfer_config.kv_connector_extra_config
327
+ spec_cls = OffloadingSpecFactory.get_spec_cls(extra_config)
328
self._offloading_metric_metadata: dict[str, OffloadingMetricMetadata] = {
329
**spec_cls.build_metric_definitions(extra_config),
330
**get_connector_metric_definitions(),
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+217
-207
@@ -6,6 +6,7 @@ from dataclasses import dataclass, field
6
from itertools import islice
7
from typing import Any, NamedTuple
8
9
+from vllm.config import VllmConfig
10
from vllm.distributed.kv_events import KVCacheEvent
11
from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
12
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
@@ -31,6 +32,7 @@ from vllm.v1.core.kv_cache_manager import KVCacheBlocks
32
from vllm.v1.core.sched.output import SchedulerOutput
33
from vllm.v1.kv_cache_interface import (
34
FullAttentionSpec,
35
+ KVCacheConfig,
36
KVCacheSpec,
37
MambaSpec,
38
SlidingWindowSpec,
@@ -74,32 +76,32 @@ class TransferJobStatus:
76
77
class GroupOffloadConfig(NamedTuple):
78
group_idx: int
77
- gpu_block_size: int
78
- offloaded_block_size: int
79
- hash_block_size_factor: int
79
+ tokens_per_block: int
80
+ tokens_per_chunk: int
81
+ hashes_per_chunk: int
82
# KV cache spec metadata propagated onto emitted BlockStored events so
83
# KV-aware consumers can classify and filter the group.
84
kv_event_group_spec: OffloadingEventGroupSpec
85
# None below means full attention
84
- sliding_window_size_in_blocks: int | None
85
- # Number of this group's offloaded blocks per full-attention alignment
86
- # segment. Used to skip storing SWA blocks that can never serve a load
86
+ sliding_window_size_in_chunks: int | None
87
+ # Number of this group's offloaded chunks per full-attention alignment
88
+ # segment. Used to skip storing SWA chunks that can never serve a load
89
# hit (e.g. DeepSeek V4 where SWA groups have much smaller block sizes
90
# than the MLA full-attention group).
91
# None for full-attention groups or when the optimization doesn't apply.
90
- alignment_block_count: int | None = None
91
- # True for EAGLE/MTP draft-model attention groups. The trailing block
92
+ alignment_chunk_count: int | None = None
93
+ # True for EAGLE/MTP draft-model attention groups. The trailing chunk
94
# of these groups is volatile and lacks a stable hash, so it must
95
# be excluded from store and load scheduling.
96
is_eagle_group: bool = False
97
98
97
-def get_sliding_window_size_in_blocks(
98
- kv_cache_spec: KVCacheSpec, offloaded_block_size: int
99
+def get_sliding_window_size_in_chunks(
100
+ kv_cache_spec: KVCacheSpec, tokens_per_chunk: int
101
) -> int | None:
102
if isinstance(kv_cache_spec, SlidingWindowSpec):
103
assert kv_cache_spec.sliding_window > 0
102
- return cdiv(kv_cache_spec.sliding_window, offloaded_block_size)
104
+ return cdiv(kv_cache_spec.sliding_window, tokens_per_chunk)
105
106
if isinstance(kv_cache_spec, MambaSpec):
107
# Mamba depends on a single state
@@ -109,116 +111,121 @@ def get_sliding_window_size_in_blocks(
111
return None
112
113
112
-def resolve_mamba_align_size(spec: "OffloadingSpec") -> int | None:
114
+def resolve_mamba_align_size(
115
+ spec: "OffloadingSpec", kv_cache_config: KVCacheConfig
116
+) -> int | None:
117
"""Scan all KV cache groups in *spec* and return the single mamba alignment
118
size, or None if no group requires mamba alignment.
119
120
For MambaSpec groups in "align" cache mode the hit window must be rounded
117
- down to a multiple of the offloaded block size. Asserts that all such
121
+ down to a multiple of the offloaded chunk size. Asserts that all such
122
groups agree on the same value.
123
"""
124
mamba_align_size: int | None = None
121
- for idx, gpu_block_size in enumerate(spec.gpu_block_size):
122
- kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec
125
+ for idx, tokens_per_block in enumerate(spec.tokens_per_block):
126
+ kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec
127
if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align":
124
- offload_block_size = gpu_block_size * spec.block_size_factor
125
- assert mamba_align_size is None or mamba_align_size == offload_block_size
126
- mamba_align_size = offload_block_size
128
+ tokens_per_chunk = tokens_per_block * spec.blocks_per_chunk
129
+ assert mamba_align_size is None or mamba_align_size == tokens_per_chunk
130
+ mamba_align_size = tokens_per_chunk
131
return mamba_align_size
132
133
134
class SchedulerOffloadConfig(NamedTuple):
135
kv_group_configs: tuple[GroupOffloadConfig, ...]
132
- block_size_factor: int
136
+ blocks_per_chunk: int
137
num_workers: int
138
offload_prompt_only: bool
139
140
@classmethod
137
- def from_spec(cls, spec: OffloadingSpec) -> "SchedulerOffloadConfig":
141
+ def from_spec(
142
+ cls,
143
+ spec: OffloadingSpec,
144
+ vllm_config: VllmConfig,
145
+ kv_cache_config: KVCacheConfig,
146
+ ) -> "SchedulerOffloadConfig":
147
# Determine the alignment token count from the full-attention group(s).
139
- # This is the offloaded_block_size of the full-attention group; load
148
+ # This is the tokens_per_chunk of the full-attention group; load
149
# hits are always aligned to this boundary, so SWA blocks earlier in
150
# each segment can never serve a load hit. Relevant for hybrid
151
# architectures like DeepSeek V4 (MLA + SWA groups).
143
- full_attn_offloaded_block_sizes: set[int] = set()
144
- for idx, gpu_block_size in enumerate(spec.gpu_block_size):
145
- kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec
146
- sw = get_sliding_window_size_in_blocks(
147
- kv_spec, gpu_block_size * spec.block_size_factor
152
+ full_attn_tokens_per_chunk: set[int] = set()
153
+ for idx, tokens_per_block in enumerate(spec.tokens_per_block):
154
+ kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec
155
+ sw = get_sliding_window_size_in_chunks(
156
+ kv_spec, tokens_per_block * spec.blocks_per_chunk
157
)
158
if sw is None:
150
- full_attn_offloaded_block_sizes.add(
151
- gpu_block_size * spec.block_size_factor
152
- )
159
+ full_attn_tokens_per_chunk.add(tokens_per_block * spec.blocks_per_chunk)
160
161
# Only apply the optimization if there's a single consistent
162
# full-attention alignment size.
163
alignment_tokens: int | None = None
157
- if len(full_attn_offloaded_block_sizes) == 1:
158
- alignment_tokens = full_attn_offloaded_block_sizes.pop()
164
+ if len(full_attn_tokens_per_chunk) == 1:
165
+ alignment_tokens = full_attn_tokens_per_chunk.pop()
166
160
- def _alignment_block_count(
161
- offloaded_block_size: int,
162
- sliding_window_size_in_blocks: int | None,
167
+ def _alignment_chunk_count(
168
+ tokens_per_chunk: int,
169
+ sliding_window_size_in_chunks: int | None,
170
) -> int | None:
164
- if alignment_tokens is None or sliding_window_size_in_blocks is None:
171
+ if alignment_tokens is None or sliding_window_size_in_chunks is None:
172
return None
166
- if alignment_tokens <= offloaded_block_size:
173
+ if alignment_tokens <= tokens_per_chunk:
174
return None
168
- per_segment = alignment_tokens // offloaded_block_size
169
- if sliding_window_size_in_blocks >= per_segment:
175
+ per_segment = alignment_tokens // tokens_per_chunk
176
+ if sliding_window_size_in_chunks >= per_segment:
177
return None
178
return per_segment
179
180
eagle_groups = {
181
idx
175
- for idx, g in enumerate(spec.kv_cache_config.kv_cache_groups)
182
+ for idx, g in enumerate(kv_cache_config.kv_cache_groups)
183
if g.is_eagle_group
184
}
185
186
use_eagle = (
180
- spec.vllm_config.speculative_config is not None
181
- and spec.vllm_config.speculative_config.use_eagle()
187
+ vllm_config.speculative_config is not None
188
+ and vllm_config.speculative_config.use_eagle()
189
)
190
if use_eagle and not eagle_groups:
184
- eagle_groups = set(range(len(spec.kv_cache_config.kv_cache_groups)))
191
+ eagle_groups = set(range(len(kv_cache_config.kv_cache_groups)))
192
193
if eagle_groups:
194
logger.info(
195
"KV offloading: EAGLE/MTP draft attention groups %s "
189
- "detected. The trailing block of these groups will be "
196
+ "detected. The trailing chunk of these groups will be "
197
"excluded from offloading due to volatility.",
198
sorted(eagle_groups),
199
)
200
201
return cls(
195
- num_workers=spec.vllm_config.parallel_config.world_size,
202
+ num_workers=vllm_config.parallel_config.world_size,
203
kv_group_configs=tuple(
204
GroupOffloadConfig(
205
group_idx=idx,
199
- gpu_block_size=gpu_block_size,
200
- offloaded_block_size=gpu_block_size * spec.block_size_factor,
201
- hash_block_size_factor=(
202
- (gpu_block_size * spec.block_size_factor)
203
- // spec.hash_block_size
206
+ tokens_per_block=tokens_per_block,
207
+ tokens_per_chunk=tokens_per_block * spec.blocks_per_chunk,
208
+ hashes_per_chunk=(
209
+ (tokens_per_block * spec.blocks_per_chunk)
210
+ // spec.tokens_per_hash
211
),
205
- sliding_window_size_in_blocks=(
206
- sw := get_sliding_window_size_in_blocks(
207
- spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec,
208
- gpu_block_size * spec.block_size_factor,
212
+ sliding_window_size_in_chunks=(
213
+ sw := get_sliding_window_size_in_chunks(
214
+ kv_cache_config.kv_cache_groups[idx].kv_cache_spec,
215
+ tokens_per_block * spec.blocks_per_chunk,
216
)
217
),
211
- alignment_block_count=_alignment_block_count(
212
- gpu_block_size * spec.block_size_factor, sw
218
+ alignment_chunk_count=_alignment_chunk_count(
219
+ tokens_per_block * spec.blocks_per_chunk, sw
220
),
221
kv_event_group_spec=get_offloading_event_group_spec(
215
- spec.kv_cache_config.kv_cache_groups[idx]
222
+ kv_cache_config.kv_cache_groups[idx]
223
),
224
is_eagle_group=idx in eagle_groups,
225
)
219
- for idx, gpu_block_size in enumerate(spec.gpu_block_size)
226
+ for idx, tokens_per_block in enumerate(spec.tokens_per_block)
227
),
221
- block_size_factor=spec.block_size_factor,
228
+ blocks_per_chunk=spec.blocks_per_chunk,
229
offload_prompt_only=spec.offload_prompt_only,
230
)
231
@@ -227,11 +234,11 @@ class SchedulerOffloadConfig(NamedTuple):
234
class RequestGroupState:
235
offload_keys: list[OffloadKey] = field(default_factory=list)
236
block_ids: list[int] = field(default_factory=list)
230
- # index of next block (of size offloaded_block_size) to offload
231
- next_stored_block_idx: int = 0
232
- # number of offloaded blocks hit (including GPU prefix cache)
237
+ # Index of the next chunk to offload.
238
+ next_stored_chunk_idx: int = 0
239
+ # Number of offloaded chunks hit (including GPU prefix cache)
240
# when the request first started
234
- num_hit_blocks: int = 0
241
+ num_hit_chunks: int = 0
242
243
244
@dataclass(slots=True)
@@ -278,11 +285,11 @@ class RequestOffloadState:
285
):
286
for req_block_hash in islice(
287
self.req.block_hashes,
281
- group_config.hash_block_size_factor * len(group_state.offload_keys)
282
- + group_config.hash_block_size_factor
288
+ group_config.hashes_per_chunk * len(group_state.offload_keys)
289
+ + group_config.hashes_per_chunk
290
- 1,
291
None,
285
- group_config.hash_block_size_factor,
292
+ group_config.hashes_per_chunk,
293
):
294
group_state.offload_keys.append(
295
make_offload_key(req_block_hash, group_config.group_idx)
@@ -298,46 +305,46 @@ class RequestOffloadState:
305
for group_state, new_blocks in zip(self.group_states, new_block_id_groups):
306
group_state.block_ids.extend(new_blocks)
307
301
- def storable_blocks(
308
+ def storable_chunks(
309
self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int
310
) -> int:
304
- """Number of leading offloaded blocks eligible for store.
311
+ """Number of leading offloaded chunks eligible for store.
312
306
- For eagle/MTP groups the volatile trailing block of the offloadable
313
+ For eagle/MTP groups the volatile trailing chunk of the offloadable
314
range is excluded while decoding: the draft-layer KV of the last
315
accepted position may be rewritten after spec-token rejection. During
309
- prefill the trailing block is stable (the draft input for a chunk's
316
+ prefill the trailing chunk is stable (the draft input for a chunk's
317
last position is the next prompt token), so it is stored immediately.
318
The exclusion must be applied consistently everywhere
312
- ``next_stored_block_idx`` is derived: otherwise the trailing block of
319
+ ``next_stored_chunk_idx`` is derived: otherwise the trailing chunk of
320
each step is skipped on collection but jumped over by
314
- ``next_stored_block_idx``, so it is never re-considered and a
321
+ ``next_stored_chunk_idx``, so it is never re-considered and a
322
permanent hole breaks prefix-reuse lookup.
323
"""
317
- num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
324
+ num_chunks = num_offloadable_tokens // group_config.tokens_per_chunk
325
is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens
326
if group_config.is_eagle_group and is_decoding:
320
- num_blocks = max(0, num_blocks - 1)
321
- return num_blocks
327
+ num_chunks = max(0, num_chunks - 1)
328
+ return num_chunks
329
330
def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
324
- # max(): at the prefill->decode transition of a block-aligned prompt,
325
- # storable_blocks drops by one (the eagle exclusion kicks in), and the
326
- # index must not move backwards past already-stored blocks.
331
+ # max(): at the prefill->decode transition of a chunk-aligned prompt,
332
+ # storable_chunks drops by one (the eagle exclusion kicks in), and the
333
+ # index must not move backwards past already-stored chunks.
334
for group_config, group_state in zip(
335
self.config.kv_group_configs, self.group_states
336
):
330
- group_state.next_stored_block_idx = max(
331
- group_state.next_stored_block_idx,
332
- self.storable_blocks(group_config, num_offloadable_tokens),
337
+ group_state.next_stored_chunk_idx = max(
338
+ group_state.next_stored_chunk_idx,
339
+ self.storable_chunks(group_config, num_offloadable_tokens),
340
)
341
335
- def update_num_hit_blocks(self, num_cached_tokens: int) -> None:
342
+ def update_num_hit_chunks(self, num_cached_tokens: int) -> None:
343
for group_config, group_state in zip(
344
self.config.kv_group_configs, self.group_states
345
):
339
- group_state.num_hit_blocks = (
340
- num_cached_tokens // group_config.offloaded_block_size
346
+ group_state.num_hit_chunks = (
347
+ num_cached_tokens // group_config.tokens_per_chunk
348
)
349
350
@@ -354,22 +361,26 @@ class OffloadingConnectorScheduler:
361
def __init__(
362
self,
363
spec: OffloadingSpec,
364
+ vllm_config: VllmConfig,
365
+ kv_cache_config: KVCacheConfig,
366
):
358
- self.config = SchedulerOffloadConfig.from_spec(spec)
367
+ self.config = SchedulerOffloadConfig.from_spec(
368
+ spec, vllm_config, kv_cache_config
369
+ )
370
self.manager: OffloadingManager = spec.get_manager()
371
self._connector_stats = OffloadingConnectorStats()
372
373
full_attention_groups: list[int] = []
374
sliding_window_groups: list[int] = []
375
for group_config in self.config.kv_group_configs:
365
- if group_config.sliding_window_size_in_blocks is None:
376
+ if group_config.sliding_window_size_in_chunks is None:
377
full_attention_groups.append(group_config.group_idx)
378
else:
379
sliding_window_groups.append(group_config.group_idx)
380
381
# sort sliding window groups by window size in decreasing order
382
def _sliding_window_sort_key(i: int) -> int:
372
- val = self.config.kv_group_configs[i].sliding_window_size_in_blocks
383
+ val = self.config.kv_group_configs[i].sliding_window_size_in_chunks
384
assert val is not None
385
return val
386
@@ -378,7 +389,9 @@ class OffloadingConnectorScheduler:
389
# used by _lookup
390
self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups)
391
self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups
381
- self._mamba_align_size: int | None = resolve_mamba_align_size(spec)
392
+ self._mamba_align_size: int | None = resolve_mamba_align_size(
393
+ spec, kv_cache_config
394
+ )
395
396
self._req_status: dict[ReqId, RequestOffloadState] = {}
397
self._current_batch_load_jobs: dict[int, TransferJob] = {}
@@ -386,9 +399,9 @@ class OffloadingConnectorScheduler:
399
# GPU block IDs allocated in the current engine step
400
self._current_batch_allocated_block_ids: set[int] = set()
401
# if GPU prefix caching is enabled,
389
- # track loaded blocks to avoid redundant loads
390
- self._blocks_being_loaded: set[OffloadKey] | None = (
391
- set() if spec.vllm_config.cache_config.enable_prefix_caching else None
402
+ # Track loaded chunks to avoid redundant loads.
403
+ self._chunks_being_loaded: set[OffloadKey] | None = (
404
+ set() if vllm_config.cache_config.enable_prefix_caching else None
405
)
406
407
# Job ID counter shared by loads and stores.
@@ -434,7 +447,7 @@ class OffloadingConnectorScheduler:
447
def _maximal_prefix_lookup(
448
self, keys: Iterable[OffloadKey], req_context: ReqContext
449
) -> int | None:
437
- """Return the number of consecutive offloaded blocks from the start,
450
+ """Return the number of consecutive offloaded chunks from the start,
451
or None if the backend deferred a lookup."""
452
hit_count = 0
453
defer_lookup = False
@@ -490,18 +503,18 @@ class OffloadingConnectorScheduler:
503
for group_config, group_state in zip(
504
self.config.kv_group_configs, req_status.group_states
505
):
493
- if group_config.sliding_window_size_in_blocks is None:
506
+ if group_config.sliding_window_size_in_chunks is None:
507
self.manager.touch(group_state.offload_keys, req_status.req_context)
508
else:
496
- # we aim to keep just blocks that are necessary to hit
497
- # the original request (+ decoded blocks)
498
- blocks_to_skip = max(
509
+ # Keep only chunks needed to hit the original request, plus
510
+ # decoded chunks.
511
+ chunks_to_skip = max(
512
0,
500
- group_state.num_hit_blocks
501
- - group_config.sliding_window_size_in_blocks,
513
+ group_state.num_hit_chunks
514
+ - group_config.sliding_window_size_in_chunks,
515
)
516
self.manager.touch(
504
- group_state.offload_keys[blocks_to_skip:],
517
+ group_state.offload_keys[chunks_to_skip:],
518
req_status.req_context,
519
)
520
@@ -531,7 +544,7 @@ class OffloadingConnectorScheduler:
544
defer_lookup = False
545
lookup_groups = self._lookup_groups
546
534
- # Tracks which eagle groups have already popped their volatile trailing block
547
+ # Tracks which eagle groups have already popped their volatile trailing chunk
548
# in the current convergence iteration. Reset when a non-eagle group
549
# tightens the hit boundary, requiring a fresh pop.
550
eagle_verified: set[int] = set()
@@ -544,79 +557,76 @@ class OffloadingConnectorScheduler:
557
group_idx
558
]
559
group_state: RequestGroupState = req_status.group_states[group_idx]
547
- offloaded_block_size = group_config.offloaded_block_size
560
+ tokens_per_chunk = group_config.tokens_per_chunk
561
offload_keys = group_state.offload_keys
562
563
assert (
551
- len(offload_keys)
552
- >= req_status.req.num_tokens // offloaded_block_size
564
+ len(offload_keys) >= req_status.req.num_tokens // tokens_per_chunk
565
)
566
567
is_eagle_unverified = (
568
group_config.is_eagle_group and group_idx not in eagle_verified
569
)
570
559
- # Constrain to block-aligned boundary for this group
571
+ # Constrain to a chunk-aligned boundary for this group.
572
max_hit_size_tokens = min(
561
- max_hit_size_tokens, len(offload_keys) * offloaded_block_size
573
+ max_hit_size_tokens, len(offload_keys) * tokens_per_chunk
574
)
563
- if max_hit_size_tokens - num_computed_tokens < offloaded_block_size:
564
- # we can only load less than a block, better skip
575
+ if max_hit_size_tokens - num_computed_tokens < tokens_per_chunk:
576
+ # We can only load less than a chunk, so skip.
577
return 0
578
567
- sliding_window_size_in_blocks = (
568
- group_config.sliding_window_size_in_blocks
579
+ sliding_window_size_in_chunks = (
580
+ group_config.sliding_window_size_in_chunks
581
)
582
571
- # For eagle groups, query one extra block that will be popped.
583
+ # For eagle groups, query one extra chunk that will be popped.
584
# We only need to increase the query size for sliding window groups.
585
query_max = max_hit_size_tokens
574
- if is_eagle_unverified and sliding_window_size_in_blocks is not None:
586
+ if is_eagle_unverified and sliding_window_size_in_chunks is not None:
587
query_max = min(
576
- max_hit_size_tokens + offloaded_block_size,
577
- len(offload_keys) * offloaded_block_size,
588
+ max_hit_size_tokens + tokens_per_chunk,
589
+ len(offload_keys) * tokens_per_chunk,
590
)
591
580
- num_blocks = min(
581
- cdiv(query_max, offloaded_block_size), len(offload_keys)
582
- )
583
- start_block_idx = num_computed_tokens // offloaded_block_size
584
- offload_keys = offload_keys[start_block_idx:num_blocks]
592
+ num_chunks = min(cdiv(query_max, tokens_per_chunk), len(offload_keys))
593
+ start_chunk_idx = num_computed_tokens // tokens_per_chunk
594
+ offload_keys = offload_keys[start_chunk_idx:num_chunks]
595
596
# end index (in the sliced offload_keys) up to which we
597
# have backend-confirmed hits
588
- num_hit_blocks: int | None
589
- if sliding_window_size_in_blocks is None:
590
- num_hit_blocks = self._maximal_prefix_lookup(
598
+ num_hit_chunks: int | None
599
+ if sliding_window_size_in_chunks is None:
600
+ num_hit_chunks = self._maximal_prefix_lookup(
601
offload_keys, req_status.req_context
602
)
603
else:
594
- required_window = sliding_window_size_in_blocks
604
+ required_window = sliding_window_size_in_chunks
605
if is_eagle_unverified:
606
required_window += 1
597
- num_hit_blocks = self._sliding_window_lookup(
607
+ num_hit_chunks = self._sliding_window_lookup(
608
offload_keys,
609
required_window,
610
req_status.req_context,
611
)
602
- if num_hit_blocks == 0:
612
+ if num_hit_chunks == 0:
613
return 0
614
605
- if num_hit_blocks is None:
615
+ if num_hit_chunks is None:
616
defer_lookup = True
617
else:
618
if is_eagle_unverified:
609
- num_hit_blocks -= 1
619
+ num_hit_chunks -= 1
620
eagle_verified.add(group_idx)
621
622
max_hit_size_tokens = min(
623
max_hit_size_tokens,
614
- offloaded_block_size * (start_block_idx + num_hit_blocks),
624
+ tokens_per_chunk * (start_chunk_idx + num_hit_chunks),
625
)
626
627
new_num_hit_tokens = max_hit_size_tokens - num_computed_tokens
618
- if new_num_hit_tokens < offloaded_block_size:
619
- # we can only load less than a block, better skip
628
+ if new_num_hit_tokens < tokens_per_chunk:
629
+ # We can only load less than a chunk, so skip.
630
return 0
631
632
if new_num_hit_tokens < num_hit_tokens:
@@ -632,7 +642,7 @@ class OffloadingConnectorScheduler:
642
# sliding window works with the new_num_hit_tokens
643
lookup_groups = self._sliding_window_groups
644
635
- looked_up_sliding_window |= sliding_window_size_in_blocks is not None
645
+ looked_up_sliding_window |= sliding_window_size_in_chunks is not None
646
num_hit_tokens = new_num_hit_tokens
647
648
if defer_lookup:
@@ -642,28 +652,28 @@ class OffloadingConnectorScheduler:
652
)
653
return None
654
645
- # possibly delay request if any of the hit blocks is already being loaded
646
- if self._blocks_being_loaded:
655
+ # Possibly delay the request if any hit chunk is already being loaded.
656
+ if self._chunks_being_loaded:
657
for group_config, group_state in zip(
658
self.config.kv_group_configs, req_status.group_states
659
):
650
- offloaded_block_size = group_config.offloaded_block_size
651
- sliding_window_size_in_blocks = (
652
- group_config.sliding_window_size_in_blocks
660
+ tokens_per_chunk = group_config.tokens_per_chunk
661
+ sliding_window_size_in_chunks = (
662
+ group_config.sliding_window_size_in_chunks
663
)
664
offload_keys = group_state.offload_keys
655
- num_blocks = cdiv(
656
- num_computed_tokens + num_hit_tokens, offloaded_block_size
665
+ num_chunks = cdiv(
666
+ num_computed_tokens + num_hit_tokens, tokens_per_chunk
667
)
658
- start_block_idx = num_computed_tokens // offloaded_block_size
659
- offload_keys = offload_keys[start_block_idx:num_blocks]
660
- if sliding_window_size_in_blocks is not None:
661
- offload_keys = offload_keys[-sliding_window_size_in_blocks:]
662
- if any(key in self._blocks_being_loaded for key in offload_keys):
663
- # hit blocks are being loaded, delay request
668
+ start_chunk_idx = num_computed_tokens // tokens_per_chunk
669
+ offload_keys = offload_keys[start_chunk_idx:num_chunks]
670
+ if sliding_window_size_in_chunks is not None:
671
+ offload_keys = offload_keys[-sliding_window_size_in_chunks:]
672
+ if any(key in self._chunks_being_loaded for key in offload_keys):
673
+ # Hit chunks are being loaded, so delay the request.
674
logger.debug(
675
"Delaying request %s since some of its"
666
- " blocks are already being loaded",
676
+ " chunks are already being loaded",
677
req_status.req.request_id,
678
)
679
return None
@@ -740,7 +750,7 @@ class OffloadingConnectorScheduler:
750
req_status.deferred_lookup_start_time = lookup_start
751
else:
752
self._maybe_observe_lookup_async_delay(req_status)
743
- req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0))
753
+ req_status.update_num_hit_chunks(num_computed_tokens + (num_hit_tokens or 0))
754
755
self._touch(req_status)
756
@@ -771,10 +781,10 @@ class OffloadingConnectorScheduler:
781
block.block_id for block in group_blocks if block.block_id != 0
782
)
783
774
- gpu_block_size = group_config.gpu_block_size
775
- offloaded_block_size = group_config.offloaded_block_size
784
+ tokens_per_block = group_config.tokens_per_block
785
+ tokens_per_chunk = group_config.tokens_per_chunk
786
offload_keys = group_state.offload_keys
777
- num_gpu_blocks = cdiv(num_cached_tokens, gpu_block_size)
787
+ num_gpu_blocks = cdiv(num_cached_tokens, tokens_per_block)
788
789
assert len(group_blocks) >= num_gpu_blocks
790
num_locally_computed_gpu_blocks = num_gpu_blocks
@@ -786,24 +796,24 @@ class OffloadingConnectorScheduler:
796
797
assert (
798
num_locally_computed_tokens
789
- <= num_locally_computed_gpu_blocks * gpu_block_size
799
+ <= num_locally_computed_gpu_blocks * tokens_per_block
800
)
801
num_pending_gpu_blocks = num_gpu_blocks - num_locally_computed_gpu_blocks
802
793
- if group_config.sliding_window_size_in_blocks is not None:
803
+ if group_config.sliding_window_size_in_chunks is not None:
804
assert (
805
num_pending_gpu_blocks
796
- <= group_config.sliding_window_size_in_blocks
797
- * self.config.block_size_factor
806
+ <= group_config.sliding_window_size_in_chunks
807
+ * self.config.blocks_per_chunk
808
)
809
800
- num_blocks = cdiv(num_cached_tokens, offloaded_block_size)
801
- assert len(offload_keys) >= num_blocks
810
+ num_chunks = cdiv(num_cached_tokens, tokens_per_chunk)
811
+ assert len(offload_keys) >= num_chunks
812
if num_pending_gpu_blocks:
803
- start_block_idx = (
804
- num_locally_computed_gpu_blocks // self.config.block_size_factor
813
+ start_chunk_idx = (
814
+ num_locally_computed_gpu_blocks // self.config.blocks_per_chunk
815
)
806
- keys_to_load.extend(offload_keys[start_block_idx:num_blocks])
816
+ keys_to_load.extend(offload_keys[start_chunk_idx:num_chunks])
817
818
dst_block_ids.extend(
819
block.block_id
@@ -814,11 +824,11 @@ class OffloadingConnectorScheduler:
824
group_sizes.append(num_pending_gpu_blocks)
825
block_indices.append(num_locally_computed_gpu_blocks)
826
817
- # Skip prefix-hit blocks for block-level policy; for
818
- # request-level, next_stored_block_idx stays at 0 so all
819
- # blocks (including hits) are offloaded.
827
+ # Skip prefix-hit chunks for block-level policy; for
828
+ # request-level, next_stored_chunk_idx stays at 0 so all
829
+ # chunks (including hits) are offloaded.
830
if req_status.offloading_context.policy == OffloadPolicy.BLOCK_LEVEL:
821
- group_state.next_stored_block_idx = num_blocks
831
+ group_state.next_stored_chunk_idx = num_chunks
832
833
src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context)
834
dst_spec = GPULoadStoreSpec(
@@ -841,8 +851,8 @@ class OffloadingConnectorScheduler:
851
is_store=False,
852
)
853
844
- if self._blocks_being_loaded is not None:
845
- self._blocks_being_loaded.update(keys_to_load)
854
+ if self._chunks_being_loaded is not None:
855
+ self._chunks_being_loaded.update(keys_to_load)
856
857
def _update_req_states(self, scheduler_output: SchedulerOutput) -> None:
858
"""
@@ -877,16 +887,16 @@ class OffloadingConnectorScheduler:
887
# Zero out stale block_ids in sliding window groups' pending-store
888
# positions. Only sliding window groups can have stale entries (blocks
889
# freed by remove_skipped_blocks then reallocated). Only positions in
880
- # [next_stored_block_idx * bsf, end) need checking where end is the
890
+ # [next_stored_chunk_idx * bsf, end) need checking where end is the
891
# pre-extend length: earlier positions were already offloaded, later
892
# ones are fresh allocations from this step.
893
if self._sliding_window_groups and self._current_batch_allocated_block_ids:
884
- block_size_factor = self.config.block_size_factor
894
+ blocks_per_chunk = self.config.blocks_per_chunk
895
for req_id, req_status in self._req_status.items():
896
ends = new_block_ids_end.get(req_id)
897
for i, grp_idx in enumerate(self._sliding_window_groups):
898
group_state = req_status.group_states[grp_idx]
889
- start = group_state.next_stored_block_idx * block_size_factor
899
+ start = group_state.next_stored_chunk_idx * blocks_per_chunk
900
end = ends[i] if ends is not None else len(group_state.block_ids)
901
for j in range(start, end):
902
if (
@@ -899,7 +909,7 @@ class OffloadingConnectorScheduler:
909
self,
910
scheduler_output: SchedulerOutput,
911
) -> dict[int, TransferJob]:
902
- block_size_factor = self.config.block_size_factor
912
+ blocks_per_chunk = self.config.blocks_per_chunk
913
store_jobs: dict[int, TransferJob] = {}
914
for req_id in scheduler_output.num_scheduled_tokens:
915
req_status = self._req_status.get(req_id)
@@ -915,59 +925,59 @@ class OffloadingConnectorScheduler:
925
if max_offload_tokens is not None:
926
num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens)
927
918
- # Skip decode-phase blocks: clamp to the prompt length so only
919
- # prefill (prompt) blocks become eligible for store. next_stored_idx
920
- # never advances past this boundary, so decode blocks are never
928
+ # Skip decode-phase chunks: clamp to the prompt length so only
929
+ # prefill chunks become eligible for store. next_stored_chunk_idx
930
+ # never advances past this boundary, so decode chunks are never
931
# queued in this or any later step.
932
if self.config.offload_prompt_only:
933
num_offloadable_tokens = min(
934
num_offloadable_tokens, req.num_prompt_tokens
935
)
936
927
- # Filter out blocks skipped due to sliding window attention / SSM
937
+ # Filter out chunks skipped due to sliding window attention / SSM
938
# or unreachable by the load path's alignment constraints.
939
new_offload_keys: list[OffloadKey] = []
940
for group_config, group_state in zip(
941
self.config.kv_group_configs, req_status.group_states
942
):
933
- num_blocks = req_status.storable_blocks(
943
+ num_chunks = req_status.storable_chunks(
944
group_config, num_offloadable_tokens
945
)
946
937
- start_block_idx = group_state.next_stored_block_idx
938
- if num_blocks <= start_block_idx:
947
+ start_chunk_idx = group_state.next_stored_chunk_idx
948
+ if num_chunks <= start_chunk_idx:
949
continue
940
- offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
941
- # For each block to offload, take the last corresponding GPU block.
942
- # e.g. if block size factor is 3 and GPU block IDs are
943
- # 1 5 6 7 2 4 9 3 8 then we'll take blocks 6 4 8.
950
+ offload_keys = group_state.offload_keys[start_chunk_idx:num_chunks]
951
+ # For each chunk, take the last corresponding GPU block. For
952
+ # blocks_per_chunk=3 and GPU block IDs 1 5 6 7 2 4 9 3 8,
953
+ # this selects GPU blocks 6 4 8.
954
# A block_id of 0 means either a sliding window / SSM skip
955
# or a stale entry that was zeroed out — skip it either way.
956
offload_block_ids = group_state.block_ids[
947
- start_block_idx * block_size_factor
948
- + block_size_factor
949
- - 1 : num_blocks * block_size_factor : block_size_factor
957
+ start_chunk_idx * blocks_per_chunk
958
+ + blocks_per_chunk
959
+ - 1 : num_chunks * blocks_per_chunk : blocks_per_chunk
960
]
961
assert len(offload_keys) == len(offload_block_ids)
962
953
- alignment_block_count = group_config.alignment_block_count
954
- tail = group_config.sliding_window_size_in_blocks
963
+ alignment_chunk_count = group_config.alignment_chunk_count
964
+ tail = group_config.sliding_window_size_in_chunks
965
966
for key_idx, (offload_key, block_id) in enumerate(
967
zip(offload_keys, offload_block_ids)
968
):
969
if block_id == 0:
970
continue
961
- # Skip SWA blocks that can never serve a load hit:
971
+ # Skip SWA chunks that can never serve a load hit:
972
# within each full-attention alignment segment, only the
963
- # trailing `tail` blocks are reachable by
973
+ # trailing `tail` chunks are reachable by
974
# _sliding_window_lookup. For DeepSeek V4 with 100K
975
# tokens this reduces SWA stores by ~78%.
966
- if alignment_block_count is not None:
976
+ if alignment_chunk_count is not None:
977
assert tail is not None
968
- abs_block_idx = start_block_idx + key_idx
969
- pos_in_segment = abs_block_idx % alignment_block_count
970
- if pos_in_segment < alignment_block_count - tail:
978
+ abs_chunk_idx = start_chunk_idx + key_idx
979
+ pos_in_segment = abs_chunk_idx % alignment_chunk_count
980
+ if pos_in_segment < alignment_chunk_count - tail:
981
continue
982
new_offload_keys.append(offload_key)
983
@@ -982,7 +992,7 @@ class OffloadingConnectorScheduler:
992
self._connector_stats.increase_counter(
993
_ConnectorMetricName.ALLOCATION_FAILURE
994
)
985
- logger.warning("Request %s: cannot store blocks", req_id)
995
+ logger.warning("Request %s: cannot store chunks", req_id)
996
continue
997
998
if not store_output.keys_to_store:
@@ -1002,29 +1012,29 @@ class OffloadingConnectorScheduler:
1012
self.config.kv_group_configs, req_status.group_states
1013
):
1014
is_sliding_window = (
1005
- group_config.sliding_window_size_in_blocks is not None
1015
+ group_config.sliding_window_size_in_chunks is not None
1016
)
1007
- num_blocks = req_status.storable_blocks(
1017
+ num_chunks = req_status.storable_chunks(
1018
group_config, num_offloadable_tokens
1019
)
1010
- start_block_idx = group_state.next_stored_block_idx
1020
+ start_chunk_idx = group_state.next_stored_chunk_idx
1021
block_ids = group_state.block_ids
1022
num_group_blocks = 0
1023
start_gpu_block_idx: int | None = None
1024
for idx, offload_key in enumerate(
1015
- group_state.offload_keys[start_block_idx:num_blocks]
1025
+ group_state.offload_keys[start_chunk_idx:num_chunks]
1026
):
1027
if offload_key not in keys_to_store:
1028
continue
1029
1020
- offloaded_block_idx = start_block_idx + idx
1030
+ chunk_idx = start_chunk_idx + idx
1031
1032
self._events_tracker.record_store(
1023
- req, group_config, offloaded_block_idx, offload_key
1033
+ req, group_config, chunk_idx, offload_key
1034
)
1035
1026
- gpu_block_idx = offloaded_block_idx * block_size_factor
1027
- for i in range(block_size_factor):
1036
+ gpu_block_idx = chunk_idx * blocks_per_chunk
1037
+ for i in range(blocks_per_chunk):
1038
block_id = block_ids[gpu_block_idx + i]
1039
if block_id == 0:
1040
continue
@@ -1039,8 +1049,8 @@ class OffloadingConnectorScheduler:
1049
1050
group_sizes.append(num_group_blocks)
1051
block_indices.append(start_gpu_block_idx or 0)
1042
- group_state.next_stored_block_idx = max(
1043
- group_state.next_stored_block_idx, num_blocks
1052
+ group_state.next_stored_chunk_idx = max(
1053
+ group_state.next_stored_chunk_idx, num_chunks
1054
)
1055
1056
src_spec = GPULoadStoreSpec(
@@ -1076,7 +1086,7 @@ class OffloadingConnectorScheduler:
1086
)
1087
1088
logger.debug(
1079
- "Request %s offloading %s blocks upto %d tokens (job %d)",
1089
+ "Request %s offloading %s chunks upto %d tokens (job %d)",
1090
req_id,
1091
len(keys_to_store),
1092
num_offloadable_tokens,
@@ -1198,8 +1208,8 @@ class OffloadingConnectorScheduler:
1208
self.manager.complete_store(job_status.keys, req_status.req_context)
1209
else:
1210
self.manager.complete_load(job_status.keys, req_status.req_context)
1201
- if self._blocks_being_loaded:
1202
- self._blocks_being_loaded.difference_update(job_status.keys)
1211
+ if self._chunks_being_loaded:
1212
+ self._chunks_being_loaded.difference_update(job_status.keys)
1213
if self._block_id_to_pending_jobs:
1214
# Sliding window blocks are tracked from store creation
1215
# and must be cleaned up unconditionally.
@@ -1289,7 +1299,7 @@ class OffloadingConnectorScheduler:
1299
yield from self._events_tracker.take_events(self.manager.take_events())
1300
1301
def reset_cache(self) -> None:
1292
- """Reset the offloading manager cache, evicting all stored blocks."""
1302
+ """Reset the offloading manager cache, evicting all stored chunks."""
1303
1304
# reset_cache cannot be called in the middle of a schedule step
1305
assert not self._current_batch_load_jobs
@@ -1306,10 +1316,10 @@ class OffloadingConnectorScheduler:
1316
# Reset offloading manager cache
1317
self.manager.reset_cache()
1318
1309
- # Reset store progress so active requests re-offload from block 0
1319
+ # Reset store progress so active requests re-offload from chunk 0.
1320
for status in self._req_status.values():
1321
for group_state in status.group_states:
1312
- group_state.next_stored_block_idx = 0
1322
+ group_state.next_stored_chunk_idx = 0
1323
status.transfer_jobs.clear()
1324
1325
# Discard jobs and save job_counter to be able to discard worker responses
@@ -1323,8 +1333,8 @@ class OffloadingConnectorScheduler:
1333
1334
# Note: _current_batch_jobs_to_flush is intentionally NOT cleared.
1335
# The load flush IDs collected above must be delivered to workers.
1326
- if self._blocks_being_loaded is not None:
1327
- self._blocks_being_loaded.clear()
1336
+ if self._chunks_being_loaded is not None:
1337
+ self._chunks_being_loaded.clear()
1338
1339
def shutdown(self) -> None:
1340
self.manager.shutdown()
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+14
-5
@@ -10,10 +10,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
10
OffloadingWorkerMetadata,
11
ReqId,
12
)
13
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
14
+ is_kv_cache_tensor_packed,
15
+)
16
from vllm.logger import init_logger
17
from vllm.v1.attention.backend import AttentionBackend
18
from vllm.v1.kv_cache_interface import (
19
AttentionSpec,
20
+ KVCacheConfig,
21
MambaSpec,
22
UniformTypeKVCacheSpecs,
23
)
@@ -33,8 +37,13 @@ logger = init_logger(__name__)
37
class OffloadingConnectorWorker:
38
"""Implementation of Worker side methods"""
39
36
- def __init__(self, spec: OffloadingSpec):
40
+ def __init__(
41
+ self,
42
+ spec: OffloadingSpec,
43
+ kv_cache_config: KVCacheConfig,
44
+ ):
45
self.spec = spec
46
+ self.kv_cache_config = kv_cache_config
47
self.worker: OffloadingWorker | None = None
48
49
# job_id -> req_id for in-flight loads.
@@ -50,7 +59,7 @@ class OffloadingConnectorWorker:
59
def register_kv_caches(
60
self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
61
):
53
- kv_cache_config = self.spec.kv_cache_config
62
+ kv_cache_config = self.kv_cache_config
63
num_blocks = kv_cache_config.num_blocks
64
65
# Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use
@@ -58,7 +67,7 @@ class OffloadingConnectorWorker:
67
# General (non-packed) layouts size the tensor at page_size_bytes per
68
# manager block, so page_size_bytes is the correct offloading stride.
69
layer_is_packed: dict[str, bool] = {
61
- ln: bool(kv_tensor.block_stride)
70
+ ln: is_kv_cache_tensor_packed(kv_tensor)
71
for kv_tensor in kv_cache_config.kv_cache_tensors
72
for ln in kv_tensor.shared_by
73
}
@@ -142,7 +151,7 @@ class OffloadingConnectorWorker:
151
(
152
t
153
for t in kv_cache_config.kv_cache_tensors
145
- if t.block_stride and t.shared_by
154
+ if is_kv_cache_tensor_packed(t) and t.shared_by
155
),
156
None,
157
)
@@ -237,7 +246,7 @@ class OffloadingConnectorWorker:
246
num_blocks_physical_dim = physical_to_logical.index(num_blocks_logical_dim)
247
assert num_blocks_physical_dim == 0
248
240
- kv_cache_groups = self.spec.kv_cache_config.kv_cache_groups
249
+ kv_cache_groups = self.kv_cache_config.kv_cache_groups
250
assert len(kv_cache_groups) == 1
251
kv_cache_spec = kv_cache_groups[0].kv_cache_spec
252
num_layers = len(kv_cache_groups[0].layer_names)
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+9
-3
@@ -23,6 +23,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
23
OffloadingConnectorMetadata,
24
OffloadingWorkerMetadata,
25
)
26
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
27
+ build_offloading_config,
28
+)
29
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
30
OffloadingConnectorStats,
31
OffloadPromMetrics,
@@ -56,14 +59,17 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
59
):
60
super().__init__(vllm_config, role, kv_cache_config)
61
59
- spec = OffloadingSpecFactory.create_spec(vllm_config, kv_cache_config)
62
+ offloading_config = build_offloading_config(vllm_config, kv_cache_config)
63
+ spec = OffloadingSpecFactory.create_spec(offloading_config)
64
65
self.connector_scheduler: OffloadingConnectorScheduler | None = None
66
self.connector_worker: OffloadingConnectorWorker | None = None
67
if role == KVConnectorRole.SCHEDULER:
64
- self.connector_scheduler = OffloadingConnectorScheduler(spec)
68
+ self.connector_scheduler = OffloadingConnectorScheduler(
69
+ spec, vllm_config, kv_cache_config
70
+ )
71
elif role == KVConnectorRole.WORKER:
66
- self.connector_worker = OffloadingConnectorWorker(spec)
72
+ self.connector_worker = OffloadingConnectorWorker(spec, kv_cache_config)
73
74
def shutdown(self) -> None:
75
if self.connector_worker is not None:
vllm/v1/kv_offload/base.py
+9
-56
@@ -14,14 +14,13 @@ import numpy as np
14
import torch
15
16
from vllm.logger import init_logger
17
-from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
17
18
if TYPE_CHECKING:
20
- from vllm.config import VllmConfig
19
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
20
OffloadingConnectorStats,
21
)
24
- from vllm.v1.kv_cache_interface import KVCacheConfig
22
+
23
+from vllm.v1.kv_offload.config import OffloadingConfig
24
25
# `OffloadKey` identifies an offloaded block. It combines a block hash with
26
# its KV cache group index, encoded as raw bytes to avoid tuple GC overhead.
@@ -482,22 +481,15 @@ class OffloadingSpec(ABC):
481
"""Return Prometheus metric definitions emitted by this spec."""
482
return {}
483
485
- def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
484
+ def __init__(self, config: OffloadingConfig):
485
logger.warning(
486
"Initializing OffloadingSpec. This API is experimental and "
487
"subject to change in the future as we iterate the design."
488
)
490
- self.vllm_config = vllm_config
491
- self.kv_cache_config = kv_cache_config
492
-
493
- kv_transfer_config = vllm_config.kv_transfer_config
494
- assert kv_transfer_config is not None
495
- self.extra_config = kv_transfer_config.kv_connector_extra_config
496
- kv_events_config = vllm_config.kv_events_config
489
+ self.config = config
490
+ self.extra_config = config.extra_config
491
self.kv_events_config = OffloadingKVEventsConfig(
498
- enable_kv_cache_events=(
499
- kv_events_config is not None and kv_events_config.enable_kv_cache_events
500
- ),
492
+ enable_kv_cache_events=config.enable_kv_cache_events,
493
self_describing_kv_events=bool(
494
self.extra_config.get("self_describing_kv_events", False)
495
),
@@ -511,48 +503,9 @@ class OffloadingSpec(ABC):
503
self.extra_config.get("offload_prompt_only", True)
504
)
505
514
- parallel_config = vllm_config.parallel_config
515
- context_parallel_factor = (
516
- parallel_config.decode_context_parallel_size
517
- * parallel_config.prefill_context_parallel_size
518
- )
519
-
520
- # gpu block size per group
521
- self.gpu_block_size: tuple[int, ...] = tuple(
522
- kv_cache_group.kv_cache_spec.block_size * context_parallel_factor
523
- for kv_cache_group in kv_cache_config.kv_cache_groups
524
- )
525
-
526
- # hash_block_size must match what the scheduler uses for
527
- # Request.block_hashes (resolved via resolve_kv_cache_block_sizes).
528
- _, self.hash_block_size = resolve_kv_cache_block_sizes(
529
- kv_cache_config, vllm_config
530
- )
531
-
532
- for block_size in self.gpu_block_size:
533
- assert block_size % self.hash_block_size == 0, (
534
- f"gpu_block_size={block_size} not divisible by "
535
- f"hash_block_size={self.hash_block_size}. "
536
- f"Hybrid models (e.g. Mamba+Attention) need "
537
- f"--enable-prefix-caching to align block sizes."
538
- )
539
-
540
- # offloaded_block_size / gpu_block_size
541
- self.block_size_factor: int = 1
542
-
543
- offloaded_block_size = self.extra_config.get("block_size")
544
- if offloaded_block_size is not None:
545
- offloaded_block_size_int = int(offloaded_block_size)
546
- gpu_block_sizes = set(self.gpu_block_size)
547
- assert len(gpu_block_sizes) == 1, (
548
- "If 'block_size' is specified in kv_connector_extra_config, "
549
- "there must be at least one KV cache group, "
550
- "and all groups must have the same block size."
551
- )
552
- gpu_block_size = gpu_block_sizes.pop()
553
-
554
- assert offloaded_block_size_int % gpu_block_size == 0
555
- self.block_size_factor = offloaded_block_size_int // gpu_block_size
506
+ self.tokens_per_block = tuple(group.tokens_per_block for group in config.groups)
507
+ self.tokens_per_hash = config.cache.tokens_per_hash
508
+ self.blocks_per_chunk = config.cache.blocks_per_chunk
509
510
@abstractmethod
511
def get_manager(self) -> OffloadingManager:
vllm/v1/kv_offload/config.py
new
+70
@@ -0,0 +1,70 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Normalized configuration consumed by native offloading backends."""
4
+
5
+from collections.abc import Mapping
6
+from dataclasses import dataclass
7
+from typing import Any
8
+
9
+
10
+@dataclass(frozen=True)
11
+class OffloadingGroupConfig:
12
+ # Total token span covered by one block across all workers
13
+ # (accounts for context parallelism).
14
+ tokens_per_block: int
15
+ # Layer names belonging to this group.
16
+ layer_names: tuple[str, ...]
17
+
18
+
19
+@dataclass(frozen=True)
20
+class OffloadingModelConfig:
21
+ # Model identifier (e.g. HuggingFace model path).
22
+ name: str
23
+ # KV cache data type (e.g. "float16").
24
+ dtype: str
25
+
26
+
27
+@dataclass(frozen=True)
28
+class OffloadingCacheConfig:
29
+ # Tokens per block hash.
30
+ tokens_per_hash: int
31
+ # Blocks coalesced into one offload chunk.
32
+ blocks_per_chunk: int
33
+
34
+
35
+@dataclass(frozen=True)
36
+class OffloadingParallelConfig:
37
+ # Worker index in [0, world_size). 0 on the scheduler side.
38
+ rank: int
39
+ # Total number of workers.
40
+ world_size: int
41
+ # Tensor parallel size.
42
+ tp_size: int
43
+ # Pipeline parallel size.
44
+ pp_size: int
45
+ # Prefill context parallel size.
46
+ pcp_size: int
47
+ # Decode context parallel size.
48
+ dcp_size: int
49
+ # Data parallel replica index of this engine.
50
+ data_parallel_index: int
51
+ # True when concatenating a block's data across all workers yields
52
+ # the same result regardless of the parallelism configuration.
53
+ is_parallelism_agnostic: bool
54
+
55
+
56
+@dataclass(frozen=True)
57
+class OffloadingConfig:
58
+ groups: tuple[OffloadingGroupConfig, ...]
59
+ # KV bytes stored by one worker per block.
60
+ worker_kv_bytes_per_block: int
61
+ # Whether the scheduler emits KV cache events. When true,
62
+ # the offloading backend should emit events as well.
63
+ enable_kv_cache_events: bool
64
+ # Offloading-specific configuration from kv_connector_extra_config.
65
+ extra_config: Mapping[str, Any]
66
+ # Unique identifier for this engine, distinct per DP rank.
67
+ engine_id: str
68
+ model: OffloadingModelConfig
69
+ cache: OffloadingCacheConfig
70
+ parallel: OffloadingParallelConfig
vllm/v1/kv_offload/cpu/gpu_worker.py
+28
-32
@@ -72,7 +72,7 @@ class Transfer:
72
73
def compute_sub_block_ptrs(
74
block_ids: np.ndarray,
75
- block_size_factor: int,
75
+ blocks_per_chunk: int,
76
output: np.ndarray,
77
tensor: torch.Tensor,
78
skip_count: int = 0,
@@ -80,38 +80,38 @@ def compute_sub_block_ptrs(
80
"""
81
Compute byte pointers for sub-blocks of the given block IDs.
82
83
- Each block in block_ids contains block_size_factor sub-blocks.
83
+ Each block in block_ids contains blocks_per_chunk sub-blocks.
84
The pointer for sub-block j of block b is:
85
- base_ptr + b * row_stride + j * sub_block_size
85
+ base_ptr + b * row_stride + j * block_page_size
86
87
- where sub_block_size = tensor.shape[1] // block_size_factor (gpu page size).
87
+ where block_page_size = tensor.shape[1] // blocks_per_chunk (gpu page size).
88
89
- This handles tensors where row_stride != block_size_factor * sub_block_size
89
+ This handles tensors where row_stride != blocks_per_chunk * block_page_size
90
(e.g. non-contiguous CPU tensors).
91
92
Args:
93
block_ids: array of block IDs at the tensor's native granularity.
94
- block_size_factor: number of sub-blocks per block.
94
+ blocks_per_chunk: number of sub-blocks per block.
95
output: pre-allocated pointer array to write pointers into.
96
tensor: the source or destination tensor.
97
skip_count: sub-blocks to skip in the first block.
98
"""
99
- assert skip_count < block_size_factor
99
+ assert skip_count < blocks_per_chunk
100
101
num_sub_blocks = len(output)
102
base_ptr = tensor.data_ptr()
103
row_stride = tensor.stride(0)
104
105
- if block_size_factor == 1:
105
+ if blocks_per_chunk == 1:
106
# Fast path: 1:1 mapping, no sub-block expansion needed.
107
output[:] = base_ptr + block_ids.astype(np.uint64)[:num_sub_blocks] * row_stride
108
return
109
110
- # Vectorized expansion for block_size_factor > 1.
111
- assert tensor.shape[1] % block_size_factor == 0
112
- sub_block_size = tensor.shape[1] // block_size_factor
113
- sub_offsets = np.arange(block_size_factor, dtype=np.uint64) * sub_block_size
114
- # (num_blocks, 1) + (1, block_size_factor) -> (num_blocks, block_size_factor)
110
+ # Vectorized expansion for blocks_per_chunk > 1.
111
+ assert tensor.shape[1] % blocks_per_chunk == 0
112
+ block_page_size = tensor.shape[1] // blocks_per_chunk
113
+ sub_offsets = np.arange(blocks_per_chunk, dtype=np.uint64) * block_page_size
114
+ # (num_blocks, 1) + (1, blocks_per_chunk) -> (num_blocks, blocks_per_chunk)
115
all_ptrs = (
116
base_ptr + block_ids.astype(np.uint64)[:, np.newaxis] * row_stride
117
) + sub_offsets[np.newaxis, :]
@@ -175,7 +175,7 @@ class SingleDirectionOffloadingHandler:
175
self,
176
gpu_tensors: list[torch.Tensor],
177
cpu_tensors: list[torch.Tensor],
178
- block_size_factor: int,
178
+ blocks_per_chunk: int,
179
kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]],
180
gpu_to_cpu: bool,
181
mmap_region: SharedOffloadRegion | None = None,
@@ -205,7 +205,7 @@ class SingleDirectionOffloadingHandler:
205
assert cpu_tensor.device.type == "cpu"
206
_, gpu_page_size = gpu_tensor.shape
207
_, cpu_page_size = cpu_tensor.shape
208
- assert cpu_page_size == gpu_page_size * block_size_factor
208
+ assert cpu_page_size == gpu_page_size * blocks_per_chunk
209
210
self.src_tensors: list[torch.Tensor] = (
211
gpu_tensors if gpu_to_cpu else cpu_tensors
@@ -220,9 +220,9 @@ class SingleDirectionOffloadingHandler:
220
)
221
222
# GPU blocks may be smaller
223
- # cpu_page_size = gpu_page_size * block_size_factor.
224
- self.src_block_size_factor = 1 if self.gpu_to_cpu else block_size_factor
225
- self.dst_block_size_factor = block_size_factor if self.gpu_to_cpu else 1
223
+ # cpu_page_size = gpu_page_size * blocks_per_chunk.
224
+ self.src_blocks_per_chunk = 1 if self.gpu_to_cpu else blocks_per_chunk
225
+ self.dst_blocks_per_chunk = blocks_per_chunk if self.gpu_to_cpu else 1
226
227
# mmap_region to clean up on shutdown (gpu_to_cpu handler owns it)
228
self._mmap_region = mmap_region
@@ -313,20 +313,16 @@ class SingleDirectionOffloadingHandler:
313
if group_size == 0:
314
continue
315
316
- src_logical_blocks_to_skip = block_idx % self.src_block_size_factor
317
- dst_logical_blocks_to_skip = block_idx % self.dst_block_size_factor
316
+ src_logical_blocks_to_skip = block_idx % self.src_blocks_per_chunk
317
+ dst_logical_blocks_to_skip = block_idx % self.dst_blocks_per_chunk
318
src_logical_blocks_count = group_size + src_logical_blocks_to_skip
319
dst_logical_blocks_count = group_size + dst_logical_blocks_to_skip
320
321
- dst_blocks_count = cdiv(
322
- dst_logical_blocks_count, self.dst_block_size_factor
323
- )
321
+ dst_blocks_count = cdiv(dst_logical_blocks_count, self.dst_blocks_per_chunk)
322
dst_end_offset = dst_offset + dst_blocks_count
323
assert dst_end_offset <= num_dst_blocks
324
327
- src_blocks_count = cdiv(
328
- src_logical_blocks_count, self.src_block_size_factor
329
- )
325
+ src_blocks_count = cdiv(src_logical_blocks_count, self.src_blocks_per_chunk)
326
src_end_offset = src_offset + src_blocks_count
327
assert src_end_offset <= num_src_blocks
328
@@ -339,14 +335,14 @@ class SingleDirectionOffloadingHandler:
335
336
compute_sub_block_ptrs(
337
group_src,
342
- self.src_block_size_factor,
338
+ self.src_blocks_per_chunk,
339
all_src[op_idx:end_idx],
340
self.src_tensors[t_idx],
341
skip_count=src_logical_blocks_to_skip,
342
)
343
compute_sub_block_ptrs(
344
group_dst,
349
- self.dst_block_size_factor,
345
+ self.dst_blocks_per_chunk,
346
all_dst[op_idx:end_idx],
347
self.dst_tensors[t_idx],
348
skip_count=dst_logical_blocks_to_skip,
@@ -476,7 +472,7 @@ class CPUOffloadingWorker(OffloadingWorker):
472
def __init__(
473
self,
474
kv_caches: CanonicalKVCaches,
479
- block_size_factor: int,
475
+ blocks_per_chunk: int,
476
num_cpu_blocks: int,
477
mmap_region: SharedOffloadRegion | None = None,
478
):
@@ -492,7 +488,7 @@ class CPUOffloadingWorker(OffloadingWorker):
488
gpu_tensor = kv_cache_tensor.tensor.view(torch.int8).view(
489
(-1, gpu_page_size_bytes)
490
)
495
- cpu_page_size_bytes = gpu_page_size_bytes * block_size_factor
491
+ cpu_page_size_bytes = gpu_page_size_bytes * blocks_per_chunk
492
493
if mmap_region is not None:
494
cpu_tensor = mmap_region.create_next_view(cpu_page_size_bytes)
@@ -518,7 +514,7 @@ class CPUOffloadingWorker(OffloadingWorker):
514
self._store_handler = SingleDirectionOffloadingHandler(
515
gpu_tensors=gpu_tensors,
516
cpu_tensors=cpu_tensors,
521
- block_size_factor=block_size_factor,
517
+ blocks_per_chunk=blocks_per_chunk,
518
kv_cache_groups_data_refs=kv_caches.group_data_refs,
519
gpu_to_cpu=True,
520
mmap_region=mmap_region,
@@ -527,7 +523,7 @@ class CPUOffloadingWorker(OffloadingWorker):
523
self._load_handler = SingleDirectionOffloadingHandler(
524
gpu_tensors=gpu_tensors,
525
cpu_tensors=cpu_tensors,
530
- block_size_factor=block_size_factor,
526
+ blocks_per_chunk=blocks_per_chunk,
527
kv_cache_groups_data_refs=kv_caches.group_data_refs,
528
gpu_to_cpu=False,
529
)
vllm/v1/kv_offload/cpu/spec.py
+16
-31
@@ -4,10 +4,8 @@ from typing import Any
4
5
from typing_extensions import override
6
7
-from vllm.config import VllmConfig
7
from vllm.platforms import current_platform
8
from vllm.utils.math_utils import round_up
10
-from vllm.v1.kv_cache_interface import KVCacheConfig
9
from vllm.v1.kv_offload.base import (
10
CanonicalKVCaches,
11
OffloadingCounterMetadata,
@@ -18,6 +16,7 @@ from vllm.v1.kv_offload.base import (
16
OffloadingSpec,
17
OffloadingWorker,
18
)
19
+from vllm.v1.kv_offload.config import OffloadingConfig
20
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
21
from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker
22
from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
@@ -73,8 +72,8 @@ class CPUOffloadingSpec(OffloadingSpec):
72
)
73
return definitions
74
76
- def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
77
- super().__init__(vllm_config, kv_cache_config)
75
+ def __init__(self, config: OffloadingConfig):
76
+ super().__init__(config)
77
78
cpu_bytes_to_use = self.extra_config.get("cpu_bytes_to_use")
79
if not cpu_bytes_to_use:
@@ -82,42 +81,28 @@ class CPUOffloadingSpec(OffloadingSpec):
81
"cpu_bytes_to_use must be specified in kv_connector_extra_config"
82
)
83
85
- world_size = vllm_config.parallel_config.world_size
84
+ world_size = config.parallel.world_size
85
self.num_blocks = 0
87
- self.kv_bytes_per_offloaded_block = 0
86
+ self.kv_bytes_per_chunk = 0
87
self.cpu_page_size_per_worker = 0
89
- assert kv_cache_config is not None
90
- if kv_cache_config.num_blocks > 0 and world_size > 0:
91
- is_packed = any(t.block_stride for t in kv_cache_config.kv_cache_tensors)
92
- assert not is_packed or all(
93
- t.block_stride for t in kv_cache_config.kv_cache_tensors
94
- )
95
- total_gpu_kv_bytes = (
96
- kv_cache_config.kv_cache_tensors[0].size
97
- if is_packed
98
- else sum(t.size for t in kv_cache_config.kv_cache_tensors)
99
- )
100
- kv_bytes_per_block = (
101
- total_gpu_kv_bytes // kv_cache_config.num_blocks
102
- ) * world_size
103
- kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor
88
+ if config.worker_kv_bytes_per_block > 0 and world_size > 0:
89
+ kv_bytes_per_block = config.worker_kv_bytes_per_block * world_size
90
+ kv_bytes_per_chunk = kv_bytes_per_block * self.blocks_per_chunk
91
92
# calculate cpu_page_size_per_worker
106
- self.cpu_page_size_per_worker = kv_bytes_per_offloaded_block // world_size
93
+ self.cpu_page_size_per_worker = kv_bytes_per_chunk // world_size
94
95
# calculate num_blocks
109
- aligned_kv_bytes_per_offloaded_block = round_up(
110
- kv_bytes_per_offloaded_block, self.BLOCK_SIZE_ALIGNMENT
111
- )
112
- self.num_blocks = (
113
- int(cpu_bytes_to_use) // aligned_kv_bytes_per_offloaded_block
96
+ aligned_kv_bytes_per_chunk = round_up(
97
+ kv_bytes_per_chunk, self.BLOCK_SIZE_ALIGNMENT
98
)
99
+ self.num_blocks = int(cpu_bytes_to_use) // aligned_kv_bytes_per_chunk
100
116
- # Expose aligned_kv_bytes_per_offloaded_block as
117
- # kv_bytes_per_offloaded_block. Note that this might contain
101
+ # Expose aligned_kv_bytes_per_chunk as
102
+ # kv_bytes_per_chunk. Note that this might contain
103
# some padding. i.e. each offloaded block is of the form,
104
# |--- W0-B0---|---- W1-B0---| ... |---- Wn-B0---| *** maybe-pad *** |
120
- self.kv_bytes_per_offloaded_block = aligned_kv_bytes_per_offloaded_block
105
+ self.kv_bytes_per_chunk = aligned_kv_bytes_per_chunk
106
107
# scheduler-side
108
self._manager: OffloadingManager | None = None
@@ -150,7 +135,7 @@ class CPUOffloadingSpec(OffloadingSpec):
135
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
136
return CPUOffloadingWorker(
137
kv_caches=kv_caches,
153
- block_size_factor=self.block_size_factor,
138
+ blocks_per_chunk=self.blocks_per_chunk,
139
num_cpu_blocks=self.num_blocks,
140
)
141
vllm/v1/kv_offload/factory.py
+8
-22
@@ -1,15 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import importlib
4
-from collections.abc import Callable
5
-from typing import TYPE_CHECKING
4
+from collections.abc import Callable, Mapping
5
+from typing import Any
6
7
from vllm.logger import init_logger
8
from vllm.v1.kv_offload.base import OffloadingSpec
9
-
10
-if TYPE_CHECKING:
11
- from vllm.config import VllmConfig
12
- from vllm.v1.kv_cache_interface import KVCacheConfig
9
+from vllm.v1.kv_offload.config import OffloadingConfig
10
11
logger = init_logger(__name__)
12
@@ -30,10 +27,7 @@ class OffloadingSpecFactory:
27
cls._registry[name] = loader
28
29
@classmethod
33
- def get_spec_cls(cls, config: "VllmConfig") -> type[OffloadingSpec]:
34
- kv_transfer_config = config.kv_transfer_config
35
- assert kv_transfer_config is not None
36
- extra_config = kv_transfer_config.kv_connector_extra_config
30
+ def get_spec_cls(cls, extra_config: Mapping[str, Any]) -> type[OffloadingSpec]:
31
spec_name = extra_config.get("spec_name", "CPUOffloadingSpec")
32
if spec_name in cls._registry:
33
spec_cls = cls._registry[spec_name]()
@@ -47,19 +41,11 @@ class OffloadingSpecFactory:
41
return spec_cls
42
43
@classmethod
50
- def create_spec(
51
- cls,
52
- config: "VllmConfig",
53
- kv_cache_config: "KVCacheConfig",
54
- ) -> OffloadingSpec:
55
- kv_transfer_config = config.kv_transfer_config
56
- assert kv_transfer_config is not None
57
- spec_name = kv_transfer_config.kv_connector_extra_config.get(
58
- "spec_name", "CPUOffloadingSpec"
59
- )
60
- spec_cls = cls.get_spec_cls(config)
44
+ def create_spec(cls, config: OffloadingConfig) -> OffloadingSpec:
45
+ spec_name = config.extra_config.get("spec_name", "CPUOffloadingSpec")
46
+ spec_cls = cls.get_spec_cls(config.extra_config)
47
logger.info("Creating offloading spec with name: %s", spec_name)
62
- return spec_cls(config, kv_cache_config)
48
+ return spec_cls(config)
49
50
51
# Register various specs here.
vllm/v1/kv_offload/file_mapper.py
+19
-35
@@ -4,7 +4,6 @@
4
import hashlib
5
import json
6
7
-from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
7
from vllm.v1.kv_offload.base import (
8
OffloadingSpec,
9
OffloadKey,
@@ -25,8 +24,8 @@ class FileMapper:
24
self,
25
root_dir: str,
26
model_name: str,
28
- hash_block_size: int,
29
- gpu_blocks_per_file: int,
27
+ tokens_per_hash: int,
28
+ blocks_per_file: int,
29
tp_size: int,
30
pp_size: int,
31
pcp_size: int,
@@ -49,8 +48,8 @@ class FileMapper:
48
self.rank: int = rank
49
self.fields: dict = {
50
"model_name": model_name,
52
- "hash_block_size": hash_block_size,
53
- "gpu_blocks_per_file": gpu_blocks_per_file,
51
+ "tokens_per_hash": tokens_per_hash,
52
+ "blocks_per_file": blocks_per_file,
53
"tp_size": tp_size,
54
"pp_size": pp_size,
55
"pcp_size": pcp_size,
@@ -66,47 +65,32 @@ class FileMapper:
65
cls,
66
root_dir: str,
67
offloading_spec: OffloadingSpec,
69
- gpu_blocks_per_file: int = 1,
68
+ blocks_per_file: int = 1,
69
parallel_agnostic: bool = False,
70
) -> "FileMapper":
71
"""Build a FileMapper from an OffloadingSpec."""
73
- vllm_config = offloading_spec.vllm_config
74
- kv_cache_config = offloading_spec.kv_cache_config
75
-
76
- parallel_config = vllm_config.parallel_config
77
- dtype = str(vllm_config.cache_config.cache_dtype).replace("torch.", "")
72
+ config = offloading_spec.config
73
kv_cache_groups = [
74
{
80
- "block_size": group.kv_cache_spec.block_size,
75
+ "tokens_per_block": group.tokens_per_block,
76
"layer_names": list(group.layer_names),
77
}
83
- for group in kv_cache_config.kv_cache_groups
78
+ for group in config.groups
79
]
85
- # Only a single full-attention group is parallelism-invariant. MLA is
86
- # excluded: its latent KV is replicated per rank, never head-sharded.
87
- # The V2 model runner is excluded: its KV layout is not known to be
88
- # parallelism-invariant.
89
- groups = kv_cache_config.kv_cache_groups
90
- spec = groups[0].kv_cache_spec if len(groups) == 1 else None
91
- parallel_agnostic = (
92
- parallel_agnostic
93
- and not vllm_config.use_v2_model_runner
94
- and isinstance(spec, FullAttentionSpec)
95
- and not isinstance(spec, MLAAttentionSpec)
96
- )
80
+ parallel = config.parallel
81
return cls(
82
root_dir=root_dir,
99
- model_name=vllm_config.model_config.model,
100
- hash_block_size=vllm_config.cache_config.block_size,
101
- gpu_blocks_per_file=gpu_blocks_per_file,
102
- tp_size=parallel_config.tensor_parallel_size,
103
- pp_size=parallel_config.pipeline_parallel_size,
104
- pcp_size=parallel_config.prefill_context_parallel_size,
105
- dcp_size=parallel_config.decode_context_parallel_size,
106
- rank=parallel_config.rank,
107
- dtype=dtype,
83
+ model_name=config.model.name,
84
+ tokens_per_hash=config.cache.tokens_per_hash,
85
+ blocks_per_file=blocks_per_file,
86
+ tp_size=parallel.tp_size,
87
+ pp_size=parallel.pp_size,
88
+ pcp_size=parallel.pcp_size,
89
+ dcp_size=parallel.dcp_size,
90
+ rank=parallel.rank,
91
+ dtype=config.model.dtype,
92
kv_cache_groups=kv_cache_groups,
109
- parallel_agnostic=parallel_agnostic,
93
+ parallel_agnostic=(parallel_agnostic and parallel.is_parallelism_agnostic),
94
)
95
96
def get_file_name(self, key: OffloadKey) -> str:
vllm/v1/kv_offload/tiering/fs/manager.py
+3
-3
@@ -113,8 +113,8 @@ class FileSystemTierManager(SecondaryTierManager):
113
):
114
"""
115
Args:
116
- offloading_spec: contains the vllm_config, kv_cache_config
117
- and block_size_factor.
116
+ offloading_spec: Contains normalized offloading configuration and
117
+ blocks_per_chunk.
118
primary_kv_view: Memoryview of the primary tier's CPU KV cache.
119
tier_type: Tier type identifier, set by SecondaryTierFactory.
120
root_dir: Root directory for block files.
@@ -150,7 +150,7 @@ class FileSystemTierManager(SecondaryTierManager):
150
self.file_mapper = FileMapper.from_offloading_spec(
151
root_dir=root_dir,
152
offloading_spec=offloading_spec,
153
- gpu_blocks_per_file=offloading_spec.block_size_factor,
153
+ blocks_per_file=offloading_spec.blocks_per_chunk,
154
parallel_agnostic=True,
155
)
156
vllm/v1/kv_offload/tiering/p2p/manager.py
+4
-4
@@ -127,8 +127,8 @@ class P2PSecondaryTierManager(SecondaryTierManager):
127
configuration reference.
128
129
Args:
130
- offloading_spec: Owning ``OffloadingSpec`` (provides
131
- ``vllm_config`` and the offloaded block layout).
130
+ offloading_spec: Owning ``OffloadingSpec`` (provides normalized
131
+ model, parallel, and cache layout configuration).
132
primary_kv_view: Memoryview over the CPU primary tier; the
133
NIXL agent registers this region for RDMA transfers.
134
tier_type: Tier identifier (defaults to ``"p2p"``).
@@ -164,7 +164,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
164
# One control socket per DP replica: offset the base by the global
165
# data-parallel index so replicas on a host don't collide (mirrors
166
# NIXL). For DP=1 the index is 0, leaving the base port unchanged.
167
- dp_index = offloading_spec.vllm_config.parallel_config.data_parallel_index
167
+ dp_index = offloading_spec.config.parallel.data_parallel_index
168
port = int(port) + dp_index
169
# Two decoupled identities:
170
# _local_id (``host:port``): the ZMQ control identity that peers
@@ -181,7 +181,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
181
config_fields = FileMapper.from_offloading_spec(
182
root_dir="",
183
offloading_spec=offloading_spec,
184
- gpu_blocks_per_file=offloading_spec.block_size_factor,
184
+ blocks_per_file=offloading_spec.blocks_per_chunk,
185
parallel_agnostic=True,
186
).get_run_config()
187
self._data: DataTransport = NixlTransport(
vllm/v1/kv_offload/tiering/spec.py
+9
-12
@@ -36,14 +36,13 @@ from typing import Any
36
import torch
37
from typing_extensions import override
38
39
-from vllm.config import VllmConfig
39
from vllm.logger import init_logger
41
-from vllm.v1.kv_cache_interface import KVCacheConfig
40
from vllm.v1.kv_offload.base import (
41
CanonicalKVCaches,
42
OffloadingManager,
43
OffloadingMetricMetadata,
44
)
45
+from vllm.v1.kv_offload.config import OffloadingConfig
46
from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker
47
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
48
from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
@@ -87,8 +86,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
86
metrics.update(tier_cls.build_metric_definitions(tier_config))
87
return metrics
88
90
- def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
91
- super().__init__(vllm_config, kv_cache_config)
89
+ def __init__(self, config: OffloadingConfig):
90
+ super().__init__(config)
91
# Redeclare for mypy: parent sets this but `--follow-imports skip` hides it
92
self._manager: OffloadingManager | None = None
93
if self.kv_events_config.self_describing_kv_events:
@@ -110,10 +109,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
109
110
# engine_id is unique per DP replica (suffixed with _dp{rank} in both
111
# the Ray and multiprocessing paths), so it names a per-replica offload
113
- # region. Non-None is guaranteed by OffloadingSpec.__init__.
114
- assert vllm_config.kv_transfer_config is not None
115
- assert vllm_config.kv_transfer_config.engine_id is not None
116
- self._engine_id: str = vllm_config.kv_transfer_config.engine_id
112
+ # region.
113
+ self._engine_id = config.engine_id
114
115
@override
116
def get_manager(self) -> OffloadingManager:
@@ -134,7 +131,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
131
engine_id=self._engine_id,
132
num_blocks=self.num_blocks,
133
rank=None,
137
- kv_bytes_per_block=self.kv_bytes_per_offloaded_block,
134
+ kv_bytes_per_block=self.kv_bytes_per_chunk,
135
cpu_page_size=self.cpu_page_size_per_worker,
136
)
137
self._scheduler_mmap = scheduler_mmap
@@ -196,18 +193,18 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
193
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
194
# Fold the global physical device index into the replica-local
195
# [0, world_size) slot range.
199
- world_size = self.vllm_config.parallel_config.world_size
196
+ world_size = self.config.parallel.world_size
197
rank = torch.accelerator.current_device_index() % world_size
198
worker_mmap = SharedOffloadRegion(
199
engine_id=self._engine_id,
200
num_blocks=self.num_blocks,
201
rank=rank,
205
- kv_bytes_per_block=self.kv_bytes_per_offloaded_block,
202
+ kv_bytes_per_block=self.kv_bytes_per_chunk,
203
cpu_page_size=self.cpu_page_size_per_worker,
204
)
205
return CPUOffloadingWorker(
206
kv_caches=kv_caches,
210
- block_size_factor=self.block_size_factor,
207
+ blocks_per_chunk=self.blocks_per_chunk,
208
num_cpu_blocks=self.num_blocks,
209
mmap_region=worker_mmap,
210
)