From 472d330c214f88c1a2ba4ade75d21f25782e8bf5 Mon Sep 17 00:00:00 2001 From: Debasish Mohanty <139056705+Debasish-87@users.noreply.github.com> Date: Fri, 17 Jul 2026 11:30:12 +0530 Subject: [PATCH] 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 --- tests/v1/kv_offload/test_factory.py | 37 +++++++++++++++++++ .../kv_connector/v1/offloading/config.py | 19 +++++++++- 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 76403a267bc..570924cfccc 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -612,3 +612,40 @@ def test_build_metric_definitions_returns_counter_at_threshold(): config.kv_transfer_config.kv_connector_extra_config ) assert CPUOffloadingMetrics.STORES_SKIPPED in metrics + + +def test_offloading_spec_accepts_blocks_per_chunk_for_heterogeneous_groups(): + config = _make_layout_vllm_config( + cpu_bytes_to_use=65536, + extra_config={"blocks_per_chunk": 2}, + ) + + spec = _create_spec(config, _make_hybrid_kv_cache_config()) + + assert spec.tokens_per_block == (12, 16) + assert spec.blocks_per_chunk == 2 + + +def test_block_size_and_blocks_per_chunk_are_mutually_exclusive(): + config = _make_layout_vllm_config( + cpu_bytes_to_use=65536, + extra_config={ + "block_size": 64, + "blocks_per_chunk": 2, + }, + ) + + with pytest.raises(ValueError, match="Specify only one"): + _create_spec(config, _make_kv_cache_config()) + + +def test_blocks_per_chunk_must_be_positive(): + config = _make_layout_vllm_config( + cpu_bytes_to_use=65536, + extra_config={ + "blocks_per_chunk": 0, + }, + ) + + with pytest.raises(ValueError, match="greater than 0"): + _create_spec(config, _make_kv_cache_config()) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index bf86d02ec46..5cf4dffd1af 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -58,15 +58,32 @@ def build_offloading_config( ) blocks_per_chunk = 1 + blocks_per_chunk_config = extra_config.get("blocks_per_chunk") tokens_per_chunk = extra_config.get("block_size") - if tokens_per_chunk is not None: + + if blocks_per_chunk_config is not None and tokens_per_chunk is not None: + raise ValueError( + "Specify only one of 'block_size' or 'blocks_per_chunk' " + "in kv_connector_extra_config." + ) + + if blocks_per_chunk_config is not None: + blocks_per_chunk = int(blocks_per_chunk_config) + + if blocks_per_chunk <= 0: + raise ValueError("'blocks_per_chunk' must be greater than 0.") + + elif tokens_per_chunk is not None: tokens_per_chunk_int = int(tokens_per_chunk) + unique_tokens_per_block = {group.tokens_per_block for group in groups} + assert len(unique_tokens_per_block) == 1, ( "If 'block_size' is specified in kv_connector_extra_config, " "there must be at least one KV cache group, " "and all groups must have the same block size." ) + tokens_per_block = unique_tokens_per_block.pop() assert tokens_per_chunk_int % tokens_per_block == 0 blocks_per_chunk = tokens_per_chunk_int // tokens_per_block