[kv_offload] Skip scratch groups (#57145)
Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Itay Etelis committed
Sep 17, 2026 at 13:22 UTC
b6e7c1f1f0430b5d4784aea391c42067581b7f76
4 files changed
+81
-5
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+33
@@ -19,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
19
from vllm.platforms import current_platform
20
from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config
21
from vllm.v1.kv_cache_interface import (
22
+ CircularBufferSpec,
23
FullAttentionSpec,
24
HiddenStateCacheSpec,
25
KVCacheConfig,
@@ -340,6 +341,38 @@ def test_worker_kv_bytes_preserves_tensor_layout(packed: bool):
341
assert offloading_config.cache.blocks_per_chunk == 2
342
343
344
+def test_offloading_skips_scratch_group():
345
+ kv_cache_config = _make_mamba_hybrid_kv_cache_config()
346
+ kv_cache_config.kv_cache_groups.insert(
347
+ 1,
348
+ KVCacheGroupSpec(
349
+ ["scratch"],
350
+ CircularBufferSpec(
351
+ block_size=4, num_kv_heads=1, head_size=128, dtype=torch.float32
352
+ ),
353
+ ),
354
+ )
355
+ page_size = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes
356
+ kv_cache_config.kv_cache_tensors = [
357
+ KVCacheTensor(
358
+ size=page_size * kv_cache_config.num_blocks,
359
+ layers=["full_layer", "scratch", "mamba_layer"],
360
+ layer_stride=0,
361
+ block_stride=page_size,
362
+ )
363
+ ]
364
+ config = _make_vllm_config()
365
+
366
+ offloading_config = build_offloading_config(config, kv_cache_config)
367
+ scheduler_config = SchedulerOffloadConfig.from_spec(
368
+ MockOffloadingSpec(offloading_config), config, kv_cache_config
369
+ )
370
+
371
+ assert [group.group_id for group in offloading_config.groups] == [0, 2]
372
+ assert [group.group_idx for group in scheduler_config.kv_group_configs] == [0, 2]
373
+ assert offloading_config.worker_kv_bytes_per_block == page_size
374
+
375
+
376
def test_hisparse_offloads_only_indexer_group():
377
source = KVCacheGroupSpec(
378
["source"],
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+42
@@ -13,6 +13,7 @@ from vllm.platforms import current_platform
13
from vllm.utils.torch_utils import get_dtype_size
14
from vllm.v1.attention.backends.registry import AttentionBackendEnum
15
from vllm.v1.kv_cache_interface import (
16
+ CircularBufferSpec,
17
FullAttentionSpec,
18
KVCacheConfig,
19
KVCacheGroupSpec,
@@ -520,6 +521,47 @@ def test_register_kv_caches(backend):
521
assert actual.mapping is not None
522
523
524
+def test_register_packed_kv_caches_skips_scratch_group():
525
+ attn_spec = FullAttentionSpec(
526
+ block_size=BLOCK_SIZE,
527
+ num_kv_heads=NUM_KV_HEADS,
528
+ head_size=HEAD_SIZE,
529
+ dtype=DTYPE,
530
+ )
531
+ page = attn_spec.page_size_bytes
532
+ layers = ["layer0", "layer1"]
533
+ kv_cache_config = KVCacheConfig(
534
+ num_blocks=NUM_BLOCKS,
535
+ kv_cache_tensors=[
536
+ KVCacheTensor(
537
+ size=2 * page * NUM_BLOCKS,
538
+ layers=layers,
539
+ layer_stride=page,
540
+ block_stride=2 * page,
541
+ )
542
+ ],
543
+ kv_cache_groups=[
544
+ KVCacheGroupSpec(layers, attn_spec),
545
+ KVCacheGroupSpec(
546
+ ["scratch"],
547
+ CircularBufferSpec(
548
+ block_size=4, num_kv_heads=1, head_size=1, dtype=DTYPE
549
+ ),
550
+ ),
551
+ ],
552
+ )
553
+ kv_caches = _allocate_kv_caches(
554
+ kv_cache_config, [], device=torch.device(f"{DEVICE_TYPE}:0")
555
+ )
556
+ worker, spec = _make_worker(kv_cache_config)
557
+ worker.register_kv_caches(kv_caches)
558
+
559
+ canonical = spec.get_worker.call_args[0][0]
560
+ assert len(canonical.tensors) == 1
561
+ assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, 2 * page)
562
+ assert canonical.group_data_refs == [[CanonicalKVCacheRef(0, 2 * page)]]
563
+
564
+
565
@pytest.mark.parametrize("backend", ATTN_BACKENDS)
566
def test_register_kv_caches_uniform_type(backend):
567
"""Test register_kv_caches with UniformTypeKVCacheSpecs.
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+3
-3
@@ -35,7 +35,7 @@ if TYPE_CHECKING:
35
36
def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ...]:
37
if kv_cache_config.hisparse_host_num_blocks is None:
38
- return tuple(range(len(kv_cache_config.kv_cache_groups)))
38
+ return kv_cache_config.prefix_cacheable_group_ids
39
return tuple(
40
group_id
41
for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
@@ -135,12 +135,12 @@ def build_offloading_config(
135
)
136
137
worker_kv_bytes_per_block = 0
138
- all_groups_selected = len(selected_groups) == len(kv_cache_config.kv_cache_groups)
138
if (
140
- all_groups_selected
139
+ kv_cache_config.hisparse_host_num_blocks is None
140
and kv_cache_config.num_blocks > 0
141
and kv_cache_config.kv_cache_tensors
142
):
143
+ # Scratch filtering must preserve the scheduler/worker allocation stride.
144
# Every KVCacheTensor describes placement within the same backing allocation,
145
# so its size is the total, not a per-tensor share.
146
total_gpu_kv_bytes = kv_cache_config.kv_cache_tensors[0].size
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+3
-2
@@ -141,8 +141,9 @@ class OffloadingConnectorWorker:
141
),
142
None,
143
)
144
- if packed_layer_name is not None and len(selected_groups) == len(
145
- kv_cache_config.kv_cache_groups
144
+ if (
145
+ packed_layer_name is not None
146
+ and kv_cache_config.hisparse_host_num_blocks is None
147
):
148
(tensor,) = tensors_per_block[packed_layer_name]
149
num_blocks = tensor.shape[0]