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