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>

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())

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