A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 10:00 UTC
[kv_offload+HMA][13/N]: Enable HMA support (#41445)
This is the final PR in a series to enables HMA support for the offloading connector. The connector advertises `SupportsHMA` and is validated with unit tests and e2e tests. Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
May 1, 2026 at 14:30 UTC
2fa1f8ec00cf85b15422cb4c0e8eb3632ee13ea8
6 files changed
+662
-251
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+359
-108
@@ -4,6 +4,7 @@ from collections.abc import Iterable
4
from unittest.mock import MagicMock
5
6
import pytest
7
+import torch
8
9
from tests.v1.kv_connector.unit.offloading_connector.utils import (
10
generate_store_output,
@@ -15,6 +16,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
16
OffloadingConnectorScheduler,
17
)
18
from vllm.v1.core.kv_cache_utils import BlockHash
19
+from vllm.v1.kv_cache_interface import (
20
+ FullAttentionSpec,
21
+ KVCacheGroupSpec,
22
+ SlidingWindowSpec,
23
+)
24
from vllm.v1.kv_offload.base import (
25
OffloadingEvent,
26
OffloadingManager,
@@ -26,32 +32,32 @@ from vllm.v1.request import RequestStatus
32
33
@pytest.mark.parametrize("async_scheduling", [True, False])
34
def test_offloading_connector(request_runner, async_scheduling: bool):
29
- offloaded_block_size = 12
30
- gpu_block_size = 4
35
+ block_size = 4
36
+ block_size_factor = 3
37
+ offloaded_block_size = block_size * block_size_factor
38
num_gpu_blocks = 100
32
- block_size_factor = offloaded_block_size // gpu_block_size
39
40
runner = request_runner(
35
- offloaded_block_size=offloaded_block_size,
36
- gpu_block_size=gpu_block_size,
41
+ block_size=block_size,
42
num_gpu_blocks=num_gpu_blocks,
43
async_scheduling=async_scheduling,
44
+ block_size_factor=block_size_factor,
45
)
46
47
# 3 blocks, store just the middle block (skip first and last)
48
# blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
49
runner.new_request(token_ids=[0] * offloaded_block_size * 3)
44
- runner.manager.prepare_store.side_effect = (
45
- lambda keys, req_context: generate_store_output(list(keys)[1:2])
50
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
51
+ generate_store_output(list(keys)[1:2])
52
)
53
runner.run(decoded_tokens=[0])
54
55
# add block missing 1 token -> no offload
56
runner.run(
57
decoded_tokens=[0] * (offloaded_block_size - 1),
52
- expected_stored_gpu_block_indexes=(3, 4, 5),
58
+ expected_stored=(3, 4, 5),
59
)
54
- runner.manager.prepare_store.assert_not_called()
60
+ runner.manager.touch.assert_not_called()
61
62
# +1 token -> single block, fail prepare_store
63
runner.manager.prepare_store.side_effect = lambda keys, req_context: None
@@ -60,19 +66,19 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
66
67
# 1 more block (+ token for async scheduling)
68
# now set block_hashes_to_store = []
63
- runner.manager.prepare_store.side_effect = (
64
- lambda keys, req_context: generate_store_output([])
69
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
70
+ generate_store_output([])
71
)
72
runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
73
74
# 1 more block (+ token for kicking off offloading)
75
# now check touch was called with all 6 blocks
70
- runner.manager.prepare_store.side_effect = (
71
- lambda keys, req_context: generate_store_output(keys)
76
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
77
+ generate_store_output(keys)
78
)
79
runner.run(
80
decoded_tokens=[0] * (offloaded_block_size + 1),
75
- expected_stored_gpu_block_indexes=(15, 16, 17),
81
+ expected_stored=(15, 16, 17),
82
)
83
runner.manager.touch.assert_called()
84
block_hashes1 = list(runner.manager.touch.call_args.args[0])
@@ -95,23 +101,23 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
101
# terminate request
102
runner.run(
103
decoded_tokens=[EOS_TOKEN_ID],
98
- expected_stored_gpu_block_indexes=tuple(range(6 * block_size_factor)),
104
+ expected_stored=tuple(range(6 * block_size_factor)),
105
)
106
107
# full_block_tokens - num_computed_tokens < offloaded_block_size
108
runner.new_request(
103
- token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
109
+ token_ids=[0] * block_size + [1] * (offloaded_block_size - block_size)
110
)
105
- runner.manager.prepare_store.side_effect = (
106
- lambda keys, req_context: generate_store_output([])
111
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
112
+ generate_store_output([])
113
)
114
runner.run(decoded_tokens=[EOS_TOKEN_ID])
115
runner.manager.lookup.assert_not_called()
116
117
# single block lookup with no hits
118
runner.new_request(token_ids=[1] * offloaded_block_size)
113
- runner.manager.prepare_store.side_effect = (
114
- lambda keys, req_context: generate_store_output([])
119
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
120
+ generate_store_output([])
121
)
122
runner.run(decoded_tokens=[EOS_TOKEN_ID])
123
runner.manager.lookup.assert_called_once()
@@ -119,25 +125,21 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
125
# single block lookup with a hit
126
runner.scheduler.reset_prefix_cache()
127
runner.new_request(token_ids=[0] * offloaded_block_size)
122
- runner.manager.prepare_store.side_effect = (
123
- lambda keys, req_context: generate_store_output([])
128
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
129
+ generate_store_output([])
130
)
131
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
126
- runner.run(
127
- decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
128
- )
132
+ runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(0, 1, 2))
133
134
# single block lookup with a hit in a middle block
135
runner.new_request(
136
token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
137
)
134
- runner.manager.prepare_store.side_effect = (
135
- lambda keys, req_context: generate_store_output([])
138
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
139
+ generate_store_output([])
140
)
141
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
138
- runner.run(
139
- decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
140
- )
142
+ runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5))
143
144
# test take_events
145
def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
@@ -167,15 +169,16 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
169
170
@pytest.mark.parametrize("async_scheduling", [True, False])
171
def test_request_preemption(request_runner, async_scheduling: bool):
170
- offloaded_block_size = 12
171
- gpu_block_size = 4
172
+ block_size = 4
173
+ block_size_factor = 3
174
+ offloaded_block_size = block_size * block_size_factor
175
num_gpu_blocks = 100
176
177
runner = request_runner(
175
- offloaded_block_size=offloaded_block_size,
176
- gpu_block_size=gpu_block_size,
178
+ block_size=block_size,
179
num_gpu_blocks=num_gpu_blocks,
180
async_scheduling=async_scheduling,
181
+ block_size_factor=block_size_factor,
182
)
183
184
free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
@@ -184,8 +187,8 @@ def test_request_preemption(request_runner, async_scheduling: bool):
187
# 2 blocks, store all, without flushing
188
# blocks = [0, 1, 2], [3, 4, 5]
189
runner.new_request(token_ids=[0] * offloaded_block_size * 2)
187
- runner.manager.prepare_store.side_effect = (
188
- lambda keys, req_context: generate_store_output(keys)
190
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
191
+ generate_store_output(keys)
192
)
193
runner.run(
194
decoded_tokens=[0],
@@ -193,11 +196,11 @@ def test_request_preemption(request_runner, async_scheduling: bool):
196
)
197
198
# decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
196
- runner.manager.prepare_store.side_effect = (
197
- lambda keys, req_context: generate_store_output(keys)
199
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
200
+ generate_store_output(keys)
201
)
202
runner.run(
200
- decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
203
+ decoded_tokens=[0] * (2 * offloaded_block_size - block_size),
204
complete_transfers=False,
205
)
206
@@ -208,8 +211,8 @@ def test_request_preemption(request_runner, async_scheduling: bool):
211
runner.run(
212
decoded_tokens=[],
213
complete_transfers=False,
211
- expected_flushed_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
212
- expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
214
+ expected_flushed=(0, 1, 2, 3, 4, 5, 6, 7, 8),
215
+ expected_stored=(0, 1, 2, 3, 4, 5, 6, 7, 8),
216
)
217
218
# restore KV cache space and reset GPU prefix cache
@@ -219,17 +222,17 @@ def test_request_preemption(request_runner, async_scheduling: bool):
222
# request should now return from preemption
223
# re-load [0, ..., 8] from the CPU and store [9, 10, 11]
224
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 3
222
- runner.manager.prepare_store.side_effect = (
223
- lambda keys, req_context: generate_store_output(keys)
225
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
226
+ generate_store_output(keys)
227
)
228
runner.run(
226
- decoded_tokens=[0] * gpu_block_size,
227
- expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
229
+ decoded_tokens=[0] * block_size,
230
+ expected_loaded=(0, 1, 2, 3, 4, 5, 6, 7, 8),
231
)
232
233
runner.run(
234
decoded_tokens=[EOS_TOKEN_ID],
232
- expected_stored_gpu_block_indexes=(9, 10, 11),
235
+ expected_stored=(9, 10, 11),
236
)
237
238
# All stores completed before request_finished -> fence index empty.
@@ -238,25 +241,26 @@ def test_request_preemption(request_runner, async_scheduling: bool):
241
242
@pytest.mark.parametrize("async_scheduling", [True, False])
243
def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
241
- offloaded_block_size = 12
242
- gpu_block_size = 4
244
+ block_size = 4
245
+ block_size_factor = 3
246
+ offloaded_block_size = block_size * block_size_factor
247
num_gpu_blocks = 100
248
249
runner = request_runner(
246
- offloaded_block_size=offloaded_block_size,
247
- gpu_block_size=gpu_block_size,
250
+ block_size=block_size,
251
num_gpu_blocks=num_gpu_blocks,
252
async_scheduling=async_scheduling,
253
+ block_size_factor=block_size_factor,
254
)
255
256
# store 1 blocks
257
runner.new_request(token_ids=[0] * offloaded_block_size)
254
- runner.manager.prepare_store.side_effect = (
255
- lambda keys, req_context: generate_store_output(keys)
258
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
259
+ generate_store_output(keys)
260
)
261
runner.run(
262
decoded_tokens=[EOS_TOKEN_ID],
259
- expected_stored_gpu_block_indexes=(0, 1, 2),
263
+ expected_stored=(0, 1, 2),
264
)
265
266
# start a request to load the first block, but don't complete
@@ -284,12 +288,12 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
288
assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
289
290
# complete transfers
287
- runner.manager.prepare_store.side_effect = (
288
- lambda keys, req_context: generate_store_output([])
291
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
292
+ generate_store_output([])
293
)
294
runner.run(
295
decoded_tokens=[EOS_TOKEN_ID],
292
- expected_loaded_gpu_block_indexes=(0, 1, 2),
296
+ expected_loaded=(0, 1, 2),
297
)
298
299
# second request will use the GPU prefix cache
@@ -301,25 +305,26 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
305
306
@pytest.mark.parametrize("async_scheduling", [True, False])
307
def test_abort_loading_requests(request_runner, async_scheduling: bool):
304
- offloaded_block_size = 12
305
- gpu_block_size = 4
308
+ block_size = 4
309
+ block_size_factor = 3
310
+ offloaded_block_size = block_size * block_size_factor
311
num_gpu_blocks = 100
312
313
runner = request_runner(
309
- offloaded_block_size=offloaded_block_size,
310
- gpu_block_size=gpu_block_size,
314
+ block_size=block_size,
315
num_gpu_blocks=num_gpu_blocks,
316
async_scheduling=async_scheduling,
317
+ block_size_factor=block_size_factor,
318
)
319
320
# store 1 blocks
321
runner.new_request(token_ids=[0] * offloaded_block_size)
317
- runner.manager.prepare_store.side_effect = (
318
- lambda keys, req_context: generate_store_output(keys)
322
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
323
+ generate_store_output(keys)
324
)
325
runner.run(
326
decoded_tokens=[EOS_TOKEN_ID],
322
- expected_stored_gpu_block_indexes=(0, 1, 2),
327
+ expected_stored=(0, 1, 2),
328
)
329
330
# start a request to load the first block, but don't complete
@@ -345,13 +350,264 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
350
# complete loading request
351
runner.run(
352
decoded_tokens=[],
348
- expected_loaded_gpu_block_indexes=(0, 1, 2),
353
+ expected_loaded=(0, 1, 2),
354
)
355
356
# assert request is deleted
357
assert req_id not in runner.scheduler.requests
358
359
360
+@pytest.mark.parametrize("async_scheduling", [True, False])
361
+def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bool):
362
+ block_size = 4
363
+ num_gpu_blocks = 100
364
+ # sliding_window=8 -> 2 offloaded blocks (block_size_factor=1)
365
+ sliding_window = 8
366
+
367
+ kv_cache_groups = [
368
+ KVCacheGroupSpec(
369
+ ["layer0"],
370
+ FullAttentionSpec(
371
+ block_size=block_size,
372
+ num_kv_heads=1,
373
+ head_size=1,
374
+ dtype=torch.float32,
375
+ ),
376
+ ),
377
+ KVCacheGroupSpec(
378
+ ["layer1"],
379
+ SlidingWindowSpec(
380
+ block_size=block_size,
381
+ num_kv_heads=1,
382
+ head_size=1,
383
+ dtype=torch.float32,
384
+ sliding_window=sliding_window,
385
+ ),
386
+ ),
387
+ ]
388
+
389
+ runner = request_runner(
390
+ block_size=block_size,
391
+ num_gpu_blocks=num_gpu_blocks,
392
+ async_scheduling=async_scheduling,
393
+ kv_cache_groups=kv_cache_groups,
394
+ )
395
+
396
+ # Verify group configs: group 0 = full attention, group 1 = sliding window
397
+ kv_group_configs = runner.connector_scheduler.config.kv_group_configs
398
+ assert len(kv_group_configs) == 2
399
+ assert kv_group_configs[0].sliding_window_size_in_blocks is None
400
+ assert kv_group_configs[1].sliding_window_size_in_blocks == 2
401
+
402
+ # Blocks [0, 1, 2] miss
403
+ runner.new_request(token_ids=[0] * block_size * 3)
404
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
405
+ generate_store_output(keys)
406
+ )
407
+ runner.run(decoded_tokens=[0])
408
+ # _touch called from get_num_new_matched_tokens (2 groups) and
409
+ # _get_reqs_to_store (2 groups) → 4 touch calls total.
410
+ touch_calls = runner.manager.touch.call_args_list
411
+ assert len(touch_calls) == 4
412
+ assert len(touch_calls[0].args[0]) == 3
413
+ assert len(touch_calls[1].args[0]) == 3
414
+ assert len(touch_calls[2].args[0]) == 3
415
+ assert len(touch_calls[3].args[0]) == 3
416
+
417
+ # store 3 more block
418
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
419
+ generate_store_output(keys)
420
+ )
421
+ runner.run(
422
+ decoded_tokens=[0] * (block_size * 3 + 2),
423
+ expected_stored=(0, 1, 2, 3, 4, 5),
424
+ )
425
+
426
+ # touch called from _get_reqs_to_store * 3 blocks, once for each group
427
+ touch_calls = runner.manager.touch.call_args_list
428
+ assert len(touch_calls) == 6
429
+
430
+ runner.run(decoded_tokens=[EOS_TOKEN_ID])
431
+
432
+ runner.scheduler.reset_prefix_cache()
433
+
434
+ # full 3 blocks hit [0, 1, 2]
435
+ runner.new_request(token_ids=[0] * (block_size * 3 + 1))
436
+ runner.manager.lookup.return_value = True
437
+ runner.run(
438
+ decoded_tokens=[EOS_TOKEN_ID],
439
+ # Group 0 (full attn): prefix lookup hits 3 → loads blocks 0,1,2
440
+ # Group 1 (sliding window, window=2): only the last 2 blocks
441
+ # are within the window → loads blocks 1,2
442
+ expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)),
443
+ )
444
+
445
+ # one touch in get_num_new_matched_tokens x 2 groups
446
+ touch_calls = runner.manager.touch.call_args_list
447
+ assert len(touch_calls) == 2
448
+ # full attention group touched all 3 blocks
449
+ assert len(touch_calls[0].args[0]) == 3
450
+ # sliding window group touched just the last 2 blocks
451
+ assert len(touch_calls[1].args[0]) == 2
452
+
453
+ # 3 blocks are hit on GPU [0, 1, 2]
454
+ # 1 block loaded [3,]
455
+ runner.new_request(token_ids=[0] * (block_size * 4 + 1))
456
+ runner.manager.lookup.return_value = True
457
+ runner.run(
458
+ decoded_tokens=[EOS_TOKEN_ID],
459
+ # Group 0 (full attn): prefix lookup hits 3 → loads blocks 0,1,2
460
+ # Group 1 (sliding window, window=2): only the last 2 blocks
461
+ # are within the window → loads blocks 1,2
462
+ expected_loaded=((0, 3), (1, 3)),
463
+ )
464
+
465
+
466
+@pytest.mark.parametrize("async_scheduling", [True, False])
467
+def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool):
468
+ hash_block_size = 4
469
+ num_gpu_blocks = 100
470
+
471
+ # Group 0: block_size=12 (offloaded_block_size=12)
472
+ # Group 1: block_size=16 (offloaded_block_size=16)
473
+ kv_cache_groups = [
474
+ KVCacheGroupSpec(
475
+ ["layer0"],
476
+ FullAttentionSpec(
477
+ block_size=hash_block_size * 3,
478
+ num_kv_heads=1,
479
+ head_size=1,
480
+ dtype=torch.float32,
481
+ ),
482
+ ),
483
+ KVCacheGroupSpec(
484
+ ["layer1"],
485
+ FullAttentionSpec(
486
+ block_size=hash_block_size * 4,
487
+ num_kv_heads=1,
488
+ head_size=1,
489
+ dtype=torch.float32,
490
+ ),
491
+ ),
492
+ ]
493
+
494
+ runner = request_runner(
495
+ block_size=hash_block_size,
496
+ num_gpu_blocks=num_gpu_blocks,
497
+ async_scheduling=async_scheduling,
498
+ kv_cache_groups=kv_cache_groups,
499
+ )
500
+
501
+ # Verify group configs
502
+ kv_group_configs = runner.connector_scheduler.config.kv_group_configs
503
+ assert len(kv_group_configs) == 2
504
+ assert kv_group_configs[0].gpu_block_size == 12
505
+ assert kv_group_configs[0].offloaded_block_size == 12
506
+ assert kv_group_configs[1].gpu_block_size == 16
507
+ assert kv_group_configs[1].offloaded_block_size == 16
508
+
509
+ # Prompt: 25 tokens, unaligned to both block sizes.
510
+ # Group 0 blocks: [0, 1], ending_token_offset = 24
511
+ # Group 1 blocks: [0,], ending_token_offset = 16
512
+ runner.new_request(token_ids=[0] * 25)
513
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
514
+ generate_store_output(keys)
515
+ )
516
+ runner.run(decoded_tokens=[0])
517
+ # _touch called from get_num_new_matched_tokens (2 groups) and
518
+ # _get_reqs_to_store (2 groups) → 4 touch calls total.
519
+ # Group 0 has 2 offload keys, group 1 has 1.
520
+ touch_calls = runner.manager.touch.call_args_list
521
+ assert len(touch_calls) == 4
522
+ assert len(touch_calls[0].args[0]) == 2
523
+ assert len(touch_calls[1].args[0]) == 1
524
+ assert len(touch_calls[2].args[0]) == 2
525
+ assert len(touch_calls[3].args[0]) == 1
526
+
527
+ # Get to 31 tokens
528
+ # No further blocks offloaded
529
+ runner.run(decoded_tokens=[0] * 6, expected_stored=((0, 0), (0, 1), (1, 0)))
530
+
531
+ # Get to 32 tokens
532
+ # Group 0 blocks: [0, 1], ending_token_offset = 24
533
+ # Group 1 blocks: [0, 1], ending_token_offset = 32
534
+ runner.run(decoded_tokens=[0])
535
+ # _get_reqs_to_store touch: only group 1 has a new block to store
536
+ touch_calls = runner.manager.touch.call_args_list
537
+ assert len(touch_calls) == 2
538
+ assert len(touch_calls[0].args[0]) == 2
539
+ assert len(touch_calls[1].args[0]) == 2
540
+
541
+ # Get to 35 tokens
542
+ # No further blocks offloaded
543
+ runner.run(decoded_tokens=[0] * 3, expected_stored=((1, 1),))
544
+
545
+ # Get to 36 tokens
546
+ # Group 0 blocks: [0, 1, 2], ending_token_offset = 36
547
+ # Group 1 blocks: [0, 1], ending_token_offset = 32
548
+ runner.run(decoded_tokens=[0])
549
+ # _get_reqs_to_store touch: only group 0 has a new block to store
550
+ touch_calls = runner.manager.touch.call_args_list
551
+ assert len(touch_calls) == 2
552
+ assert len(touch_calls[0].args[0]) == 3
553
+ assert len(touch_calls[1].args[0]) == 2
554
+
555
+ # Get to 47 tokens
556
+ # No further blocks offloaded
557
+ runner.run(decoded_tokens=[0] * 11, expected_stored=((0, 2),))
558
+
559
+ # Get to 48 tokens
560
+ # Group 0 blocks: [0, 1, 2, 3], ending_token_offset = 4
561
+ # Group 1 blocks: [0, 1, 2], ending_token_offset = 48
562
+ runner.run(decoded_tokens=[0])
563
+ # _get_reqs_to_store touch: both groups have a new block, each with 1 key
564
+ touch_calls = runner.manager.touch.call_args_list
565
+ assert len(touch_calls) == 2
566
+ assert len(touch_calls[0].args[0]) == 4
567
+ assert len(touch_calls[1].args[0]) == 3
568
+
569
+ runner.run(decoded_tokens=[0], expected_stored=((0, 3), (1, 2)))
570
+
571
+ # Get to 96 tokens
572
+ runner.run(
573
+ decoded_tokens=[0] * 47 + [EOS_TOKEN_ID],
574
+ expected_stored=((0, 4), (0, 5), (0, 6), (0, 7), (1, 3), (1, 4), (1, 5)),
575
+ )
576
+
577
+ runner.scheduler.reset_prefix_cache()
578
+
579
+ # Request with 48 matching tokens
580
+ # will match 48 tokens (4 block) from the first group
581
+ # 48 tokens (3 block) from the second group
582
+ # Total 48 tokens can be loaded
583
+ runner.new_request(token_ids=[0] * 48)
584
+ runner.manager.lookup.return_value = True
585
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
586
+ generate_store_output([])
587
+ )
588
+ runner.run(
589
+ decoded_tokens=[0],
590
+ expected_loaded=((0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2)),
591
+ )
592
+ runner.run(decoded_tokens=[EOS_TOKEN_ID])
593
+
594
+ # Request with 48+37 matching tokens
595
+ # 48 tokens will be hit on GPU
596
+ # extra 32 tokens will be loaded
597
+ # extra tokens [0, 36] (blocks [4, 5, 6]) from the first group
598
+ # extra tokens [0, 32] (block [3, 4]) from the second group
599
+ runner.new_request(token_ids=[0] * (48 + 37))
600
+ runner.manager.lookup.return_value = True
601
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
602
+ generate_store_output([])
603
+ )
604
+ runner.run(
605
+ decoded_tokens=[0],
606
+ expected_loaded=((0, 4), (0, 5), (0, 6), (1, 3), (1, 4)),
607
+ )
608
+ runner.run(decoded_tokens=[EOS_TOKEN_ID])
609
+
610
+
611
# ---------------------------------------------------------------------------
612
# Unit tests for _maximal_prefix_lookup / _sliding_window_lookup
613
# ---------------------------------------------------------------------------
@@ -493,26 +749,24 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo
749
750
This supports P/D disaggregation where the prefill instance offloads the
751
complete KV cache so a remote decode node can consume it."""
496
- offloaded_block_size = 12
752
gpu_block_size = 4
753
+ block_size_factor = 3
754
+ offloaded_block_size = gpu_block_size * block_size_factor
755
num_gpu_blocks = 100
756
757
runner = request_runner(
501
- offloaded_block_size=offloaded_block_size,
502
- gpu_block_size=gpu_block_size,
758
+ block_size_factor=block_size_factor,
759
+ block_size=gpu_block_size,
760
num_gpu_blocks=num_gpu_blocks,
761
async_scheduling=async_scheduling,
762
)
763
764
# Store 1 offloaded block (3 GPU blocks) via a normal request.
765
runner.new_request(token_ids=[0] * offloaded_block_size)
509
- runner.manager.prepare_store.side_effect = (
510
- lambda keys, req_context: generate_store_output(keys)
511
- )
512
- runner.run(
513
- decoded_tokens=[EOS_TOKEN_ID],
514
- expected_stored_gpu_block_indexes=(0, 1, 2),
766
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
767
+ generate_store_output(keys)
768
)
769
+ runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0, 1, 2))
770
771
# Reset GPU prefix cache so the next request must load from CPU.
772
runner.scheduler.reset_prefix_cache()
@@ -524,22 +778,17 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo
778
kv_transfer_params={"do_remote_decode": True},
779
)
780
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
527
- runner.manager.prepare_store.side_effect = (
528
- lambda keys, req_context: generate_store_output(keys)
781
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
782
+ generate_store_output(keys)
783
)
784
785
# Load the first offloaded block from CPU.
532
- runner.run(
533
- decoded_tokens=[0],
534
- expected_loaded_gpu_block_indexes=(0, 1, 2),
535
- )
786
+ runner.run(decoded_tokens=[0], expected_loaded=(0, 1, 2))
787
788
# Store must include ALL 6 GPU blocks (both the loaded prefix and
789
# the newly computed block), not just the 3 new ones.
539
- runner.run(
540
- decoded_tokens=[EOS_TOKEN_ID],
541
- expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5),
542
- )
790
+ runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0, 1, 2, 3, 4, 5))
791
+
792
# All stores completed before request_finished -> fence index empty.
793
assert runner.connector_scheduler._block_id_to_pending_jobs == {}
794
@@ -553,8 +802,8 @@ def test_loads_do_not_populate_fence_index(request_runner):
802
"""Loads don't populate _block_id_to_pending_jobs (protected by
803
delay_free_blocks while in flight)."""
804
runner = request_runner(
556
- offloaded_block_size=12,
557
- gpu_block_size=4,
805
+ block_size_factor=3,
806
+ block_size=4,
807
num_gpu_blocks=100,
808
async_scheduling=False,
809
)
@@ -572,15 +821,15 @@ def test_fence_at_update_state_after_alloc(request_runner):
821
req1 just freed.
822
"""
823
runner = request_runner(
575
- offloaded_block_size=4,
576
- gpu_block_size=4,
824
+ block_size_factor=1,
825
+ block_size=4,
826
num_gpu_blocks=2,
827
async_scheduling=False,
828
)
829
830
runner.new_request(token_ids=[0] * 4)
582
- runner.manager.prepare_store.side_effect = (
583
- lambda keys, req_context: generate_store_output(keys)
831
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
832
+ generate_store_output(keys)
833
)
834
runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
835
assert runner.connector_scheduler._block_id_to_pending_jobs
@@ -588,14 +837,14 @@ def test_fence_at_update_state_after_alloc(request_runner):
837
runner.scheduler.reset_prefix_cache()
838
runner.new_request(token_ids=[0] * 4)
839
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
591
- runner.manager.prepare_store.side_effect = (
592
- lambda keys, req_context: generate_store_output([])
840
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
841
+ generate_store_output([])
842
)
843
runner.run(
844
decoded_tokens=[],
845
complete_transfers=False,
597
- expected_stored_gpu_block_indexes=(0,),
598
- expected_flushed_gpu_block_indexes=(0,),
846
+ expected_stored=(0,),
847
+ expected_flushed=(0,),
848
)
849
assert runner.connector_scheduler._block_id_to_pending_jobs == {}
850
@@ -605,15 +854,15 @@ def test_fence_at_build_store_jobs(request_runner):
854
reusing a finished request's pending-store block is flushed by
855
_build_store_jobs's fence."""
856
runner = request_runner(
608
- offloaded_block_size=4,
609
- gpu_block_size=4,
857
+ block_size_factor=1,
858
+ block_size=4,
859
num_gpu_blocks=2,
860
async_scheduling=False,
861
)
862
863
runner.new_request(token_ids=[0] * 4)
615
- runner.manager.prepare_store.side_effect = (
616
- lambda keys, req_context: generate_store_output(keys)
864
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
865
+ generate_store_output(keys)
866
)
867
runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
868
assert runner.connector_scheduler._block_id_to_pending_jobs
@@ -621,13 +870,13 @@ def test_fence_at_build_store_jobs(request_runner):
870
runner.scheduler.reset_prefix_cache()
871
runner.new_request(token_ids=[1] * 4)
872
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 0
624
- runner.manager.prepare_store.side_effect = (
625
- lambda keys, req_context: generate_store_output([])
873
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
874
+ generate_store_output([])
875
)
876
runner.run(
877
decoded_tokens=[EOS_TOKEN_ID],
629
- expected_stored_gpu_block_indexes=(0,),
630
- expected_flushed_gpu_block_indexes=(0,),
878
+ expected_stored=(0,),
879
+ expected_flushed=(0,),
880
)
881
assert runner.connector_scheduler._block_id_to_pending_jobs == {}
882
@@ -636,20 +885,22 @@ def test_fence_at_build_store_jobs(request_runner):
885
def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
886
"""complete_store fires per-job, not deferred to request finish.
887
Each call carries only that store's keys."""
639
- offloaded_block_size = 12
888
+ gpu_block_size = 4
889
+ block_size_factor = 3
890
+ offloaded_block_size = gpu_block_size * block_size_factor
891
runner = request_runner(
641
- offloaded_block_size=offloaded_block_size,
642
- gpu_block_size=4,
892
+ block_size_factor=block_size_factor,
893
+ block_size=gpu_block_size,
894
num_gpu_blocks=100,
895
async_scheduling=async_scheduling,
896
)
897
runner.new_request(token_ids=[0] * offloaded_block_size)
647
- runner.manager.prepare_store.side_effect = (
648
- lambda keys, req_context: generate_store_output(keys)
898
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
899
+ generate_store_output(keys)
900
)
901
902
# First store: fires when block 0 is fully populated.
652
- runner.run(decoded_tokens=[0, 0], expected_stored_gpu_block_indexes=(0, 1, 2))
903
+ runner.run(decoded_tokens=[0, 0], expected_stored=(0, 1, 2))
904
assert runner.manager.complete_store.call_count == 1
905
first_call_keys = set(runner.manager.complete_store.call_args.args[0])
906
assert len(first_call_keys) == 1
@@ -658,7 +909,7 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
909
# Second store: fires when block 1 is fully populated, with different keys.
910
runner.run(
911
decoded_tokens=[0] * (offloaded_block_size + 1),
661
- expected_stored_gpu_block_indexes=(3, 4, 5),
912
+ expected_stored=(3, 4, 5),
913
)
914
assert runner.manager.complete_store.call_count == 1
915
second_call_keys = set(runner.manager.complete_store.call_args.args[0])
tests/v1/kv_connector/unit/offloading_connector/utils.py
+184
-82
@@ -3,7 +3,7 @@
3
from collections.abc import Iterable, Iterator
4
from dataclasses import dataclass
5
from typing import Any
6
-from unittest.mock import MagicMock
6
+from unittest.mock import MagicMock, patch
7
8
import pytest
9
import torch
@@ -150,46 +150,61 @@ class MockOffloadingSpec(OffloadingSpec):
150
return specs
151
152
153
+@dataclass(frozen=True)
154
+class GPUBlock:
155
+ group_idx: int
156
+ request_block_offset: int
157
+
158
+
159
@dataclass
160
class TransferSummary:
155
- gpu_block_indices: list[int]
161
+ gpu_blocks: list[GPUBlock]
162
offload_addresses: list[Any]
163
164
165
class RequestRunner:
166
def __init__(
167
self,
162
- offloaded_block_size: int,
163
- gpu_block_size: int,
168
+ block_size: int,
169
num_gpu_blocks: int,
170
+ block_size_factor: int = 1,
171
async_scheduling: bool = True,
172
+ kv_cache_groups: list[KVCacheGroupSpec] | None = None,
173
):
167
- self.offloaded_block_size: int = offloaded_block_size
168
- self.gpu_block_size: int = gpu_block_size
174
+ assert block_size_factor == 1 or kv_cache_groups is None, (
175
+ "block_size_factor > 1 requires all groups to have the same "
176
+ "block size, so kv_cache_groups must be None (use default group)"
177
+ )
178
+
179
+ self.block_size_factor: int = block_size_factor
180
+ self.block_size: int = block_size
181
self.num_gpu_blocks: int = num_gpu_blocks
182
self.async_scheduling: bool = async_scheduling
183
184
self.req_id: int = -1
185
186
vllm_config = create_vllm_config(
175
- block_size=gpu_block_size, max_num_batched_tokens=1000
187
+ block_size=block_size,
188
+ max_num_batched_tokens=1000,
189
+ disable_hybrid_kv_cache_manager=False,
190
)
191
vllm_config.scheduler_config.async_scheduling = async_scheduling
192
+
193
+ extra_config: dict[str, Any] = {
194
+ "spec_name": "MockOffloadingSpec",
195
+ "spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils", # noqa: E501
196
+ }
197
+ if block_size_factor > 1:
198
+ extra_config["block_size"] = block_size * block_size_factor
199
+
200
vllm_config.kv_transfer_config = KVTransferConfig(
201
kv_connector="OffloadingConnector",
202
kv_role="kv_both",
181
- kv_connector_extra_config={
182
- "spec_name": "MockOffloadingSpec",
183
- "spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils", # noqa: E501
184
- "block_size": offloaded_block_size,
185
- },
203
+ kv_connector_extra_config=extra_config,
204
)
205
188
- block_size = vllm_config.cache_config.block_size
189
- kv_cache_config = KVCacheConfig(
190
- num_blocks=num_gpu_blocks,
191
- kv_cache_tensors=[],
192
- kv_cache_groups=[
206
+ if kv_cache_groups is None:
207
+ kv_cache_groups = [
208
KVCacheGroupSpec(
209
["layer"],
210
FullAttentionSpec(
@@ -199,7 +214,12 @@ class RequestRunner:
214
dtype=torch.float32,
215
),
216
)
202
- ],
217
+ ]
218
+
219
+ kv_cache_config = KVCacheConfig(
220
+ num_blocks=num_gpu_blocks,
221
+ kv_cache_tensors=[],
222
+ kv_cache_groups=kv_cache_groups,
223
)
224
vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
225
self.num_kv_groups = len(kv_cache_config.kv_cache_groups)
@@ -219,11 +239,38 @@ class RequestRunner:
239
240
# register worker kv_caches to enable OffloadingWorker creations
241
# set_current_vllm_config is needed for get_kv_cache_layout() to work
222
- with set_current_vllm_config(vllm_config):
223
- self.worker_connector.register_cross_layers_kv_cache(
224
- kv_cache=torch.empty(0),
225
- attn_backend=FlashAttentionBackend,
226
- )
242
+ # Mock get_layers_from_vllm_config so that mock layer names
243
+ # resolve to layers whose get_attn_backend() returns
244
+ # FlashAttentionBackend.
245
+ def _mock_get_layers(_vllm_config, _layer_type, layer_names):
246
+ mock_layer = MagicMock()
247
+ mock_layer.get_attn_backend.return_value = FlashAttentionBackend
248
+ return {name: mock_layer for name in layer_names}
249
+
250
+ kv_caches: dict[str, torch.Tensor] = {}
251
+ for group in kv_cache_groups:
252
+ spec = group.kv_cache_spec
253
+ for layer_name in group.layer_names:
254
+ # Shape follows FlashAttention layout:
255
+ # (2, num_blocks, block_size, num_kv_heads, head_size)
256
+ kv_caches[layer_name] = torch.empty(
257
+ 2,
258
+ num_gpu_blocks,
259
+ spec.block_size,
260
+ spec.num_kv_heads,
261
+ spec.head_size,
262
+ dtype=spec.dtype,
263
+ )
264
+
265
+ with (
266
+ set_current_vllm_config(vllm_config),
267
+ patch(
268
+ "vllm.distributed.kv_transfer.kv_connector.v1"
269
+ ".offloading.worker.get_layers_from_vllm_config",
270
+ side_effect=_mock_get_layers,
271
+ ),
272
+ ):
273
+ self.worker_connector.register_kv_caches(kv_caches)
274
275
# extract connector of scheduler
276
scheduler_connector = self.scheduler.connector
@@ -238,10 +285,17 @@ class RequestRunner:
285
assert isinstance(manager, MagicMock)
286
self.manager: MagicMock = manager
287
241
- assert len(self.connector_scheduler.config.kv_group_configs) == 1
242
- kv_group_config = self.connector_scheduler.config.kv_group_configs[0]
243
- assert kv_group_config.gpu_block_size == gpu_block_size
244
- assert kv_group_config.offloaded_block_size == offloaded_block_size
288
+ num_kv_groups = len(kv_cache_config.kv_cache_groups)
289
+ assert len(self.connector_scheduler.config.kv_group_configs) == num_kv_groups
290
+ for group_config, kv_cache_group in zip(
291
+ self.connector_scheduler.config.kv_group_configs,
292
+ kv_cache_config.kv_cache_groups,
293
+ ):
294
+ gpu_block_size = kv_cache_group.kv_cache_spec.block_size
295
+ assert group_config.gpu_block_size == gpu_block_size
296
+ assert (
297
+ group_config.offloaded_block_size == gpu_block_size * block_size_factor
298
+ )
299
300
# extract OffloadingSpec of worker_connector
301
connector_worker = self.worker_connector.connector_worker
@@ -250,18 +304,18 @@ class RequestRunner:
304
assert isinstance(offloading_spec, MockOffloadingSpec)
305
self.offloading_spec: MockOffloadingSpec = offloading_spec
306
253
- # mapping (offloading address) -> gpu_block_index
254
- self.offloaded: dict[Any, int] = {}
307
+ # mapping (offloading address) -> GPUBlock
308
+ self.offloaded: dict[Any, GPUBlock] = {}
309
310
self.completed_loads: list[TransferSummary] = []
311
self.completed_stores: list[TransferSummary] = []
258
- self.flushed_gpu_block_indexes: set[int] = set()
312
+ self.flushed_gpu_blocks: set[GPUBlock] = set()
313
260
- # maps {block_id: block_offset}
261
- self.gpu_block_index: dict[int, int] = {}
314
+ # block_id -> GPUBlock
315
+ self.gpu_blocks: dict[int, GPUBlock] = {}
316
317
init_none_hash(sha256)
264
- self._block_hasher = get_request_block_hasher(gpu_block_size, sha256)
318
+ self._block_hasher = get_request_block_hasher(block_size, sha256)
319
320
self._dummy_ctx: ForwardContext = ForwardContext(
321
no_compile_layers={},
@@ -297,11 +351,9 @@ class RequestRunner:
351
assert isinstance(src_spec, GPULoadStoreSpec)
352
353
for block_id in src_spec.block_ids:
300
- self.flushed_gpu_block_indexes.add(
301
- self.gpu_block_index[block_id.item()]
302
- )
354
+ self.flushed_gpu_blocks.add(self.gpu_blocks[block_id.item()])
355
304
- block_size_factor = self.offloaded_block_size // self.gpu_block_size
356
+ block_size_factor = self.block_size_factor
357
358
for transfer_spec in self.offloading_spec.get_completed_transfers():
359
src_spec, dst_spec = transfer_spec
@@ -317,10 +369,11 @@ class RequestRunner:
369
370
assert isinstance(offload_spec, MockLoadStoreSpec)
371
assert isinstance(gpu_spec, GPULoadStoreSpec)
372
+ assert len(gpu_spec.group_sizes) == self.num_kv_groups
373
321
- gpu_block_indices: list[int] = []
374
+ gpu_blocks: list[GPUBlock] = []
375
for block_id in gpu_spec.block_ids:
323
- gpu_block_indices.append(self.gpu_block_index[block_id.item()])
376
+ gpu_blocks.append(self.gpu_blocks[block_id.item()])
377
378
# list of (offload_key, sub_block_offset)
379
offload_addresses: list[Any] = []
@@ -328,30 +381,49 @@ class RequestRunner:
381
for sub_block_idx in range(block_size_factor):
382
offload_addresses.append((offload_key, sub_block_idx))
383
331
- if store:
332
- assert len(gpu_block_indices) == len(offload_addresses)
384
+ assert gpu_spec.block_indices is not None
385
+ assert len(gpu_spec.block_indices) == self.num_kv_groups
386
334
- self.completed_stores.append(
335
- TransferSummary(gpu_block_indices, offload_addresses)
336
- )
337
- else:
338
- remainder_sub_block_count = len(offload_addresses) - len(
339
- gpu_block_indices
387
+ gpu_block_offset = 0
388
+ offload_address_offset = 0
389
+ for group_size, logical_offset in zip(
390
+ gpu_spec.group_sizes, gpu_spec.block_indices
391
+ ):
392
+ gpu_block_end_offset = gpu_block_offset + group_size
393
+ assert gpu_block_end_offset <= len(gpu_blocks)
394
+
395
+ offload_addresses_to_skip = logical_offset % block_size_factor
396
+ offload_addresses_end_offset = (
397
+ offload_address_offset + offload_addresses_to_skip + group_size
398
)
341
- assert remainder_sub_block_count >= 0
342
- assert remainder_sub_block_count < block_size_factor
343
- offload_addresses = offload_addresses[remainder_sub_block_count:]
399
+ assert offload_addresses_end_offset <= len(offload_addresses)
400
345
- self.completed_loads.append(
346
- TransferSummary(gpu_block_indices, offload_addresses)
401
+ offload_addresses = (
402
+ offload_addresses[:offload_address_offset]
403
+ + offload_addresses[
404
+ offload_address_offset + offload_addresses_to_skip :
405
+ ]
406
)
407
349
- def _update_gpu_block_idx(self):
350
- for blocks in self.scheduler.kv_cache_manager.coordinator.single_type_managers[
351
- 0
352
- ].req_to_blocks.values():
353
- for block_idx, block in enumerate(blocks):
354
- self.gpu_block_index[block.block_id] = block_idx
408
+ gpu_block_offset += group_size
409
+ offload_address_offset += group_size
410
+
411
+ assert gpu_block_offset == len(gpu_blocks)
412
+ assert offload_address_offset == len(offload_addresses)
413
+
414
+ transfer_summary = TransferSummary(gpu_blocks, offload_addresses)
415
+ if store:
416
+ self.completed_stores.append(transfer_summary)
417
+ else:
418
+ self.completed_loads.append(transfer_summary)
419
+
420
+ def _update_gpu_blocks(self):
421
+ for group_idx, manager in enumerate(
422
+ self.scheduler.kv_cache_manager.coordinator.single_type_managers
423
+ ):
424
+ for blocks in manager.req_to_blocks.values():
425
+ for block_idx, block in enumerate(blocks):
426
+ self.gpu_blocks[block.block_id] = GPUBlock(group_idx, block_idx)
427
428
def _run(self, decoded_tokens: list[int], complete_transfers: bool):
429
"""
@@ -376,7 +448,7 @@ class RequestRunner:
448
break
449
450
scheduler_output = self.scheduler.schedule()
379
- self._update_gpu_block_idx()
451
+ self._update_gpu_blocks()
452
453
kv_connector_metadata = scheduler_output.kv_connector_metadata
454
assert kv_connector_metadata is not None
@@ -447,13 +519,30 @@ class RequestRunner:
519
if EOS_TOKEN_ID in decoded_tokens:
520
assert not self.scheduler.running
521
522
+ def _to_gpu_blocks(
523
+ self, blocks: tuple[int | tuple[int, int], ...]
524
+ ) -> list[GPUBlock]:
525
+ gpu_blocks: list[GPUBlock] = []
526
+ for block in blocks:
527
+ if isinstance(block, int):
528
+ for group_idx in range(self.num_kv_groups):
529
+ gpu_blocks.append(
530
+ GPUBlock(group_idx=group_idx, request_block_offset=block)
531
+ )
532
+ else:
533
+ group_idx, offset = block
534
+ gpu_blocks.append(
535
+ GPUBlock(group_idx=group_idx, request_block_offset=offset)
536
+ )
537
+ return gpu_blocks
538
+
539
def run(
540
self,
541
decoded_tokens: list[int],
542
complete_transfers: bool = True,
454
- expected_stored_gpu_block_indexes: tuple[int, ...] = (),
455
- expected_loaded_gpu_block_indexes: tuple[int, ...] = (),
456
- expected_flushed_gpu_block_indexes: tuple[int, ...] = (),
543
+ expected_stored: tuple[int | tuple[int, int], ...] = (),
544
+ expected_loaded: tuple[int | tuple[int, int], ...] = (),
545
+ expected_flushed: tuple[int | tuple[int, int], ...] = (),
546
):
547
"""
548
Runs multiple engine (scheduler + worker) steps.
@@ -462,41 +551,49 @@ class RequestRunner:
551
Args:
552
decoded_tokens: the tokens to yield at each step.
553
complete_transfers: complete transfers immediately
465
- expected_stored_gpu_block_indexes: GPU block indexes
554
+ expected_stored: GPU blocks
555
that are expected to be written during the run.
467
- expected_loaded_gpu_block_indexes: GPU block indexes
556
+ expected_loaded: GPU blocks
557
that are expected to be loaded during the run.
469
- expected_flushed_gpu_block_indexes: GPU block indexes
558
+ expected_flushed: GPU blocks
559
that are expected to be flushed during the run.
560
+
561
+ A GPU block is either a (group_idx: int, request_block_offset: int)
562
+ or just request_block_offset: int.
563
+ The latter case is a convenience for representing all groups.
564
"""
565
566
+ expected_stored_gpu_blocks = self._to_gpu_blocks(expected_stored)
567
+ expected_loaded_gpu_blocks = self._to_gpu_blocks(expected_loaded)
568
+ expected_flushed_gpu_blocks = self._to_gpu_blocks(expected_flushed)
569
+
570
self.manager.reset_mock()
571
self._run(decoded_tokens, complete_transfers)
572
476
- loaded_gpu_block_indexes: set[int] = set()
573
+ loaded_gpu_blocks: set[GPUBlock] = set()
574
for transfer in self.completed_loads:
478
- for gpu_block_idx, offloaded_address in zip(
479
- transfer.gpu_block_indices, transfer.offload_addresses
575
+ for gpu_block, offloaded_address in zip(
576
+ transfer.gpu_blocks, transfer.offload_addresses
577
):
481
- loaded_gpu_block_indexes.add(gpu_block_idx)
482
- assert gpu_block_idx == self.offloaded[offloaded_address]
578
+ loaded_gpu_blocks.add(gpu_block)
579
+ assert gpu_block == self.offloaded[offloaded_address]
580
484
- assert set(expected_loaded_gpu_block_indexes) == loaded_gpu_block_indexes
581
+ assert set(expected_loaded_gpu_blocks) == loaded_gpu_blocks
582
self.completed_loads.clear()
583
487
- stored_gpu_block_indexes: set[int] = set()
584
+ stored_gpu_blocks: set[GPUBlock] = set()
585
for transfer in self.completed_stores:
489
- for gpu_block_idx, offloaded_address in zip(
490
- transfer.gpu_block_indices, transfer.offload_addresses
586
+ for gpu_block, offloaded_address in zip(
587
+ transfer.gpu_blocks, transfer.offload_addresses
588
):
492
- stored_gpu_block_indexes.add(gpu_block_idx)
493
- self.offloaded[offloaded_address] = gpu_block_idx
589
+ stored_gpu_blocks.add(gpu_block)
590
+ self.offloaded[offloaded_address] = gpu_block
591
495
- assert set(expected_stored_gpu_block_indexes) == stored_gpu_block_indexes
592
+ assert set(expected_stored_gpu_blocks) == stored_gpu_blocks
593
self.completed_stores.clear()
594
498
- assert set(expected_flushed_gpu_block_indexes) == self.flushed_gpu_block_indexes
499
- self.flushed_gpu_block_indexes.clear()
595
+ assert set(expected_flushed_gpu_blocks) == self.flushed_gpu_blocks
596
+ self.flushed_gpu_blocks.clear()
597
598
599
@pytest.fixture
@@ -504,13 +601,18 @@ def request_runner():
601
runners = []
602
603
def runner_factory(
507
- offloaded_block_size, gpu_block_size, num_gpu_blocks, async_scheduling
604
+ block_size,
605
+ num_gpu_blocks,
606
+ async_scheduling,
607
+ block_size_factor=1,
608
+ kv_cache_groups=None,
609
):
610
runner = RequestRunner(
510
- offloaded_block_size=offloaded_block_size,
511
- gpu_block_size=gpu_block_size,
611
+ block_size=block_size,
612
num_gpu_blocks=num_gpu_blocks,
613
+ block_size_factor=block_size_factor,
614
async_scheduling=async_scheduling,
615
+ kv_cache_groups=kv_cache_groups,
616
)
617
runners.append(runner)
618
return runner
tests/v1/kv_connector/unit/test_offloading_connector.py
+106
-58
@@ -14,13 +14,36 @@ from vllm.config import KVEventsConfig, KVTransferConfig
14
from vllm.distributed.kv_events import BlockStored, KVEventBatch
15
from vllm.platforms import current_platform
16
17
-CPU_BLOCK_SIZES = [48]
18
-ATTN_BACKENDS = []
19
-
17
+_ATTN_BACKENDS: list[str] = []
18
if current_platform.is_cuda():
21
- ATTN_BACKENDS = ["FLASH_ATTN", "FLASHINFER", "TRITON_ATTN"]
19
+ _ATTN_BACKENDS = ["FLASH_ATTN", "FLASHINFER", "TRITON_ATTN"]
20
elif current_platform.is_rocm():
23
- ATTN_BACKENDS = ["TRITON_ATTN"]
21
+ _ATTN_BACKENDS = ["TRITON_ATTN"]
22
+
23
+# (model, attn_backend | None, block_size | None, uses_hma)
24
+#
25
+# - Llama: tested with each attention backend and a custom block_size.
26
+# - Gemma-3: HMA (sliding window + full attention), default backend.
27
+# - Mamba-130m: HMA (attention-free, acts like sliding_window=1),
28
+# default backend. Prefix caching must be force-enabled.
29
+# - Falcon-H1-0.5B-Instruct: HMA (parallel SSM/attention in every layer).
30
+# After page-size unification the mamba and attention groups have
31
+# different block sizes.
32
+MODEL_PARAMS: list[tuple[str, str | None, int | None, bool]] = [
33
+ ("meta-llama/Llama-3.2-1B-Instruct", backend, 48, False)
34
+ for backend in _ATTN_BACKENDS
35
+]
36
+# HMA / Mamba models are only tested on CUDA (not ROCm).
37
+if current_platform.is_cuda():
38
+ MODEL_PARAMS += [
39
+ ("google/gemma-3-1b-it", None, 48, True),
40
+ ("state-spaces/mamba-130m-hf", None, 48, True),
41
+ # Falcon-H1: parallel hybrid (every layer has both attention and SSM).
42
+ # The mamba and attention groups end up with different GPU block sizes
43
+ # after page-size unification, so we leave cpu_block_size=None
44
+ # (block_size_factor stays 1).
45
+ ("tiiuae/Falcon-H1-0.5B-Instruct", None, None, True),
46
+ ]
47
48
# Maximum time (seconds) to wait for the async CPU offload transfer
49
# to complete before giving up.
@@ -110,7 +133,7 @@ def _wait_for_prefix_cache_reset(llm: LLM) -> None:
133
)
134
135
113
-def _latency_test(llm: LLM, subscriber: MockSubscriber):
136
+def _latency_test(llm: LLM, subscriber: MockSubscriber | None):
137
sampling_params = SamplingParams(max_tokens=1)
138
139
num_times_cpu_better_than_cold = 0
@@ -118,7 +141,10 @@ def _latency_test(llm: LLM, subscriber: MockSubscriber):
141
total_cold_time = 0.0
142
total_gpu_hit_time = 0.0
143
total_cpu_hit_time = 0.0
121
- prompt_token_ids = [0] * 10001
144
+ max_model_len = llm.llm_engine.vllm_config.model_config.max_model_len
145
+ # Use a long prompt that fits within the model's context window.
146
+ prompt_len = min(10001, max_model_len - 1)
147
+ prompt_token_ids = [0] * prompt_len
148
for i in tqdm(range(num_tests), desc="Running tests"):
149
prompt_token_ids[0] = i
150
prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)]
@@ -141,10 +167,11 @@ def _latency_test(llm: LLM, subscriber: MockSubscriber):
167
168
# Verify CPU stored events arrived (offload is done before we
169
# attempt to load from CPU).
144
- assert subscriber.get_new_cpu_stored_events(), (
145
- f"No CPU stored events received on iteration {i}; "
146
- "async offload may not have completed in time"
147
- )
170
+ if subscriber is not None:
171
+ assert subscriber.get_new_cpu_stored_events(), (
172
+ f"No CPU stored events received on iteration {i}; "
173
+ "async offload may not have completed in time"
174
+ )
175
176
# run generation again - this should trigger loading from CPU
177
start_time = time.time()
@@ -163,84 +190,105 @@ def _latency_test(llm: LLM, subscriber: MockSubscriber):
190
assert num_times_cpu_better_than_cold >= 0.8 * num_tests
191
192
166
-def _accuracy_test(llm: LLM, subscriber: MockSubscriber):
193
+def _accuracy_test(llm: LLM, subscriber: MockSubscriber | None):
194
sampling_params = SamplingParams(max_tokens=1)
168
- cpu_block_size = (
169
- llm.llm_engine.vllm_config.kv_transfer_config.kv_connector_extra_config[
170
- "block_size"
171
- ]
195
+ extra_config = (
196
+ llm.llm_engine.vllm_config.kv_transfer_config.kv_connector_extra_config
197
)
173
-
174
- subscriber.get_new_cpu_stored_events()
175
-
176
- # prepend prompt to be cpu block aligned
198
+ cpu_block_size = extra_config.get("block_size")
199
+ if cpu_block_size is None:
200
+ # No custom offloaded block_size: offloaded blocks match GPU blocks.
201
+ # Use the hash block_size (cache_config.block_size) for alignment.
202
+ cpu_block_size = llm.llm_engine.vllm_config.cache_config.block_size
203
+
204
+ if subscriber is not None:
205
+ subscriber.get_new_cpu_stored_events()
206
+
207
+ # Pad prompt so its token count is a multiple of cpu_block_size.
208
+ # Use the tokenizer directly to avoid expensive llm.generate() calls.
209
+ tokenizer = llm.get_tokenizer()
210
prompt = "Let's count to 10. One, two, three, four,"
178
- while (
179
- len(llm.generate(prompt, use_tqdm=False)[0].prompt_token_ids) % cpu_block_size
180
- != 0
181
- ):
211
+ while len(tokenizer.encode(prompt)) % cpu_block_size != 0:
212
prompt = ". " + prompt
213
184
- assert subscriber.get_new_cpu_stored_events()
214
+ # Seed the CPU cache with the prompt.
215
+ llm.generate(prompt, sampling_params, use_tqdm=False)
216
186
- test_count = 100
187
- success_count = 0
188
- for i in range(test_count):
189
- if (
190
- llm.generate(prompt, sampling_params, use_tqdm=False)[0].outputs[0].text
191
- == " five"
192
- ):
193
- success_count += 1
217
+ if subscriber is not None:
218
+ assert subscriber.get_new_cpu_stored_events()
219
220
+ test_count = 20
221
+ results = llm.generate([prompt] * test_count, sampling_params, use_tqdm=False)
222
+ success_count = sum(1 for r in results if r.outputs[0].text == " five")
223
assert success_count >= 0.5 * test_count
224
225
198
-@pytest.mark.parametrize("cpu_block_size", CPU_BLOCK_SIZES)
199
-@pytest.mark.parametrize("attn_backend", ATTN_BACKENDS)
200
-def test_cpu_offloading(cpu_block_size: int, attn_backend: str) -> None:
226
+@pytest.mark.parametrize("model, attn_backend, cpu_block_size, uses_hma", MODEL_PARAMS)
227
+def test_cpu_offloading(
228
+ model: str,
229
+ attn_backend: str | None,
230
+ cpu_block_size: int | None,
231
+ uses_hma: bool,
232
+) -> None:
233
"""
234
Tests OffloadingConnector with CPUOffloadingSpec.
235
"""
204
-
236
# configure OffloadingConnector (spec_name=CPUOffloadingSpec by default)
237
+ extra_config: dict = {"cpu_bytes_to_use": 500 << 20}
238
+ if cpu_block_size is not None:
239
+ extra_config["block_size"] = cpu_block_size
240
kv_transfer_config = KVTransferConfig(
241
kv_connector="OffloadingConnector",
242
kv_role="kv_both",
209
- kv_connector_extra_config={
210
- "cpu_bytes_to_use": 500 << 20,
211
- "block_size": cpu_block_size,
212
- },
243
+ kv_connector_extra_config=extra_config,
244
)
245
215
- port: int
216
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
217
- s.bind(("0.0.0.0", 0))
218
- port = s.getsockname()[1]
219
-
220
- events_endpoint = f"tcp://*:{port}"
221
- kv_events_config = KVEventsConfig(
222
- enable_kv_cache_events=True,
223
- publisher="zmq",
224
- endpoint=events_endpoint,
225
- topic="test",
226
- )
246
+ # KV events are incompatible with HMA (setting kv_events_config
247
+ # would force HMA off), so only enable them for non-HMA models.
248
+ subscriber: MockSubscriber | None = None
249
+ kv_events_config: KVEventsConfig | None = None
250
+ if not uses_hma:
251
+ port: int
252
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
253
+ s.bind(("0.0.0.0", 0))
254
+ port = s.getsockname()[1]
255
+
256
+ events_endpoint = f"tcp://*:{port}"
257
+ kv_events_config = KVEventsConfig(
258
+ enable_kv_cache_events=True,
259
+ publisher="zmq",
260
+ endpoint=events_endpoint,
261
+ topic="test",
262
+ )
263
+
264
+ # Attention-free / hybrid models disable prefix caching by default
265
+ # (ModelConfig.is_prefix_caching_supported returns False). Without it,
266
+ # mamba_block_size falls back to max_model_len, making GPU blocks too
267
+ # large for any reasonable offloaded block_size. Force-enable it.
268
+ force_prefix_caching = uses_hma
269
270
llm = LLM(
229
- model="meta-llama/Llama-3.2-1B-Instruct",
271
+ model=model,
272
+ max_model_len=4096,
273
gpu_memory_utilization=0.5,
274
kv_events_config=kv_events_config,
275
kv_transfer_config=kv_transfer_config,
233
- attention_config={"backend": attn_backend},
276
+ **({"attention_config": {"backend": attn_backend}} if attn_backend else {}),
277
+ # HMA models need explicit opt-in when kv_transfer_config is set
278
+ **({"disable_hybrid_kv_cache_manager": False} if uses_hma else {}),
279
+ **({"enable_prefix_caching": True} if force_prefix_caching else {}),
280
# ROCm: batch size 1 to reduce variability
281
**({"max_num_seqs": 1} if current_platform.is_rocm() else {}),
282
)
283
238
- events_endpoint = events_endpoint.replace("*", "127.0.0.1")
239
- subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
284
+ if kv_events_config is not None:
285
+ events_endpoint = events_endpoint.replace("*", "127.0.0.1")
286
+ subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
287
288
try:
289
_latency_test(llm, subscriber)
290
_accuracy_test(llm, subscriber)
291
finally:
245
- subscriber.close()
292
+ if subscriber is not None:
293
+ subscriber.close()
294
del llm
tests/v1/kv_connector/unit/utils.py
+2
@@ -103,6 +103,7 @@ def create_vllm_config(
103
kv_load_failure_policy: Literal["recompute", "fail"] = "fail",
104
kv_connector: str = "NixlConnector",
105
kv_role: str = "kv_both",
106
+ disable_hybrid_kv_cache_manager: bool | None = None,
107
) -> VllmConfig:
108
"""Initialize VllmConfig For Testing."""
109
model_config = ModelConfig(
@@ -118,6 +119,7 @@ def create_vllm_config(
119
max_model_len=max_model_len,
120
enable_chunked_prefill=enable_chunked_prefill,
121
is_encoder_decoder=model_config.is_encoder_decoder,
122
+ disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager,
123
)
124
# Cache config, optionally force APC
125
cache_config = CacheConfig(
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
-1
@@ -829,7 +829,6 @@ class OffloadingConnectorScheduler:
829
def request_finished(
830
self,
831
request: Request,
832
- block_ids: list[int],
832
) -> tuple[bool, dict[str, Any] | None]:
833
"""
834
Called when a request has finished, before its blocks are freed.
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+11
-2
@@ -10,6 +10,7 @@ from vllm.distributed.kv_events import KVCacheEvent
10
from vllm.distributed.kv_transfer.kv_connector.v1 import (
11
KVConnectorBase_V1,
12
KVConnectorRole,
13
+ SupportsHMA,
14
)
15
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
16
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
@@ -42,7 +43,7 @@ from vllm.v1.outputs import KVConnectorOutput
43
from vllm.v1.request import Request
44
45
45
-class OffloadingConnector(KVConnectorBase_V1):
46
+class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
47
@property
48
def prefer_cross_layer_blocks(self) -> bool:
49
return True
@@ -149,7 +150,15 @@ class OffloadingConnector(KVConnectorBase_V1):
150
block_ids: list[int],
151
) -> tuple[bool, dict[str, Any] | None]:
152
assert self.connector_scheduler is not None
152
- return self.connector_scheduler.request_finished(request, block_ids)
153
+ return self.connector_scheduler.request_finished(request)
154
+
155
+ def request_finished_all_groups(
156
+ self,
157
+ request: "Request",
158
+ block_ids: tuple[list[int], ...],
159
+ ) -> tuple[bool, dict[str, Any] | None]:
160
+ assert self.connector_scheduler is not None
161
+ return self.connector_scheduler.request_finished(request)
162
163
def take_events(self) -> Iterable[KVCacheEvent]:
164
assert self.connector_scheduler is not None