Add blocks_per_chunk configuration for KV offloading to support heterogeneous KV cache groups (#48878)
Signed-off-by: Debasish-87 <22btics06@suiit.ac.in> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Debasish Mohanty committed
Jul 17, 2026 at 11:30 UTC
472d330c214f88c1a2ba4ade75d21f25782e8bf5
2 files changed
+55
-1
tests/v1/kv_offload/test_factory.py
+37
@@ -612,3 +612,40 @@ def test_build_metric_definitions_returns_counter_at_threshold():
612
config.kv_transfer_config.kv_connector_extra_config
613
)
614
assert CPUOffloadingMetrics.STORES_SKIPPED in metrics
615
+
616
+
617
+def test_offloading_spec_accepts_blocks_per_chunk_for_heterogeneous_groups():
618
+ config = _make_layout_vllm_config(
619
+ cpu_bytes_to_use=65536,
620
+ extra_config={"blocks_per_chunk": 2},
621
+ )
622
+
623
+ spec = _create_spec(config, _make_hybrid_kv_cache_config())
624
+
625
+ assert spec.tokens_per_block == (12, 16)
626
+ assert spec.blocks_per_chunk == 2
627
+
628
+
629
+def test_block_size_and_blocks_per_chunk_are_mutually_exclusive():
630
+ config = _make_layout_vllm_config(
631
+ cpu_bytes_to_use=65536,
632
+ extra_config={
633
+ "block_size": 64,
634
+ "blocks_per_chunk": 2,
635
+ },
636
+ )
637
+
638
+ with pytest.raises(ValueError, match="Specify only one"):
639
+ _create_spec(config, _make_kv_cache_config())
640
+
641
+
642
+def test_blocks_per_chunk_must_be_positive():
643
+ config = _make_layout_vllm_config(
644
+ cpu_bytes_to_use=65536,
645
+ extra_config={
646
+ "blocks_per_chunk": 0,
647
+ },
648
+ )
649
+
650
+ with pytest.raises(ValueError, match="greater than 0"):
651
+ _create_spec(config, _make_kv_cache_config())
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+18
-1
@@ -58,15 +58,32 @@ def build_offloading_config(
58
)
59
60
blocks_per_chunk = 1
61
+ blocks_per_chunk_config = extra_config.get("blocks_per_chunk")
62
tokens_per_chunk = extra_config.get("block_size")
62
- if tokens_per_chunk is not None:
63
+
64
+ if blocks_per_chunk_config is not None and tokens_per_chunk is not None:
65
+ raise ValueError(
66
+ "Specify only one of 'block_size' or 'blocks_per_chunk' "
67
+ "in kv_connector_extra_config."
68
+ )
69
+
70
+ if blocks_per_chunk_config is not None:
71
+ blocks_per_chunk = int(blocks_per_chunk_config)
72
+
73
+ if blocks_per_chunk <= 0:
74
+ raise ValueError("'blocks_per_chunk' must be greater than 0.")
75
+
76
+ elif tokens_per_chunk is not None:
77
tokens_per_chunk_int = int(tokens_per_chunk)
78
+
79
unique_tokens_per_block = {group.tokens_per_block for group in groups}
80
+
81
assert len(unique_tokens_per_block) == 1, (
82
"If 'block_size' is specified in kv_connector_extra_config, "
83
"there must be at least one KV cache group, "
84
"and all groups must have the same block size."
85
)
86
+
87
tokens_per_block = unique_tokens_per_block.pop()
88
assert tokens_per_chunk_int % tokens_per_block == 0
89
blocks_per_chunk = tokens_per_chunk_int // tokens_per_block