[Bugfix][KV Offload] Size CPU cache slots for one KV group

docs/features/kv_offloading_usage.md

@@ -26,6 +26,10 @@ flowchart LR
 
 The unit of operation is a **chunk** — a fixed-size piece of KV data covering a group of tokens. By default, a chunk maps to a single accelerator block. A configurable `blocks_per_chunk` parameter allows larger chunks, yielding larger I/Os to the host and secondary tiers.
 
+For models with multiple KV cache groups, each offload key stores data from one group. CPU slots are shared by all groups and sized for the largest selected group, including worker copies and alignment padding. Smaller groups can leave unused space within a slot. This applies to both `CPUOffloadingSpec` and the CPU primary tier of `TieringOffloadingSpec`.
+
+The compact multi-group layout uses a separate persistent-cache namespace. Files written with the previous layout remain on disk but are not reused.
+
 ## Single-Tier Setup (CPU Only)
 
 ```bash

tests/v1/kv_connector/unit/offloading_connector/test_config.py

@@ -914,3 +914,64 @@ def test_blocks_per_chunk_must_be_positive():
 
     with pytest.raises(ValueError, match="greater than 0"):
         build_offloading_config(config, _make_kv_cache_config())
+
+
+@pytest.mark.parametrize("head_sizes", [(128,) * 8, (64, 128, 256)])
+@pytest.mark.parametrize("blocks_per_chunk", [1, 3])
+def test_cpu_capacity_reserves_one_group_per_key(head_sizes, blocks_per_chunk):
+    """A group-specific key must not reserve all other groups' pages."""
+    from vllm.utils.math_utils import round_up
+    from vllm.v1.kv_offload.base import LookupResult, ReqContext, make_offload_key
+    from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
+    from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
+
+    groups = [
+        KVCacheGroupSpec(
+            [f"layer{i}"],
+            FullAttentionSpec(
+                block_size=16,
+                num_kv_heads=1,
+                head_size=head_size,
+                dtype=torch.float16,
+            ),
+        )
+        for i, head_size in enumerate(head_sizes)
+    ]
+    pages = [g.kv_cache_spec.page_size_bytes for g in groups]
+    row_bytes = round_up(
+        max(pages) * blocks_per_chunk, SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
+    )
+    capacity = 2 * len(groups)
+    config = _make_vllm_config(
+        extra_config={
+            "cpu_bytes_to_use": row_bytes * capacity,
+            "blocks_per_chunk": blocks_per_chunk,
+        }
+    )
+    cache = KVCacheConfig(
+        num_blocks=16,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=sum(pages) * 16,
+                layers=[name for g in groups for name in g.layer_names],
+                layer_stride=0,
+                block_stride=sum(pages),
+            )
+        ],
+        kv_cache_groups=groups,
+    )
+    spec = CPUOffloadingSpec(build_offloading_config(config, cache))
+    assert spec.num_chunks == capacity
+    assert spec.kv_bytes_per_chunk == row_bytes
+    manager = spec.get_manager()
+    ctx = ReqContext("capacity")
+    keys = [
+        make_offload_key(i.to_bytes(8, "big"), g)
+        for g in range(len(groups))
+        for i in range(2)
+    ]
+    store = manager.prepare_store(keys, ctx)
+    assert store is not None
+    assert store.evicted_keys == []
+    manager.complete_store(keys, ctx, success=True)
+    assert all(manager.lookup(key, ctx) is LookupResult.HIT for key in keys)

tests/v1/kv_connector/unit/offloading_connector/test_worker.py

@@ -107,6 +107,7 @@ def _make_worker(
 
     spec = MagicMock(spec=OffloadingSpec)
     spec.replicated_layout = replicated_layout
+    spec.compact_group_layout = False
     spec.config = MagicMock()
     spec.config.parallel.rank = rank
     spec.get_worker.return_value = MagicMock()
@@ -636,3 +637,44 @@ def test_register_kv_caches_uniform_type(backend):
     # opaque mapping rather than a certified, parallelism-agnostic one
     assert group_refs[0].mapping.parallelism_agnostic
     assert not group_refs[1].mapping.parallelism_agnostic
+
+
+@pytest.mark.parametrize("compact", [False, True])
+def test_packed_groups_keep_separate_transfer_regions(monkeypatch, compact):
+    """A group key must not copy unrelated layers from a packed GPU block."""
+    spec = FullAttentionSpec(
+        block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float16
+    )
+    page = spec.page_size_bytes
+    packed = torch.arange(4 * 2 * page, dtype=torch.int64).to(torch.int8)
+    packed = packed.view(4, 2 * page)
+    caches = {"layer0": packed[:, :page], "layer1": packed[:, page:]}
+    config = KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=packed.numel(),
+                layers=list(caches),
+                layer_stride=page,
+                block_stride=2 * page,
+            )
+        ],
+        kv_cache_groups=[KVCacheGroupSpec([name], spec) for name in caches],
+    )
+    monkeypatch.setattr(
+        "vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker."
+        "derive_canonical_mappings",
+        lambda *args: {},
+    )
+    worker, offloading_spec = _make_worker(config)
+    offloading_spec.compact_group_layout = compact
+    worker.register_kv_caches(caches)
+    registered = offloading_spec.get_worker.call_args.args[0]
+    for index, refs in enumerate(registered.group_data_refs):
+        assert len(refs) == 1
+        ref = refs[0]
+        tensor = registered.tensors[ref.tensor_idx].tensor
+        assert ref.page_size_bytes == (page if compact else 2 * page)
+        torch.testing.assert_close(
+            tensor, caches[f"layer{index}"] if compact else packed
+        )

tests/v1/kv_offload/cpu/test_gpu_worker.py

@@ -17,6 +17,8 @@
     CanonicalKVCacheRef,
     CanonicalKVCaches,
     CanonicalKVCacheTensor,
+    CanonicalPageMapping,
+    CopyRun,
     GPULoadStoreSpec,
     TransferResult,
 )
@@ -682,3 +684,110 @@ def test_load_waits_for_pending_compute_stream_writes(default_vllm_config) -> No
                 torch.testing.assert_close(gpu_tensor[block_id].cpu(), expected)
     finally:
         worker.shutdown()
+
+
+@pytest.mark.parametrize("layout", ["private", "shared", "canonical"])
+@torch.inference_mode()
+def test_compact_groups_roundtrip_after_slot_reuse(default_vllm_config, layout):
+    """Aliased, unequal groups share slots without losing sparse chunk tails."""
+    packed = torch.randint(1, 100, (18, 1536), dtype=torch.int8, device=DEVICES[0])
+    original = packed.clone()
+    # The first alias exposes less of the backing page than group 1 needs.
+    pages = [packed[:, :512], packed[:, 1024:]]
+    groups = [[(0, 512), (1, 512)], [(0, 1024)], [(1, 256), (0, 256)]]
+    canonical = layout == "canonical"
+    refs = []
+    for group in groups:
+        group_refs = []
+        for index, size in group:
+            # Simulate one TP shard's two fragments in a twice-as-large page.
+            mapping = (
+                CanonicalPageMapping(
+                    size * 2,
+                    size,
+                    (CopyRun(0, size // 2, size // 2, 2, size // 2, size),),
+                    1,
+                    0,
+                    True,
+                )
+                if canonical
+                else None
+            )
+            group_refs.append(CanonicalKVCacheRef(index, size, mapping))
+        refs.append(group_refs)
+    region = None
+    if layout != "private":
+        region = SharedOffloadRegion(
+            engine_id=str(uuid.uuid4()),
+            num_chunks=2,
+            rank=None if canonical else 1,
+            kv_bytes_per_chunk=8192,
+            cpu_page_size=3072,
+        )
+    worker = CPUOffloadingWorker(
+        kv_caches=CanonicalKVCaches(
+            [CanonicalKVCacheTensor(page, page.shape[1]) for page in pages], refs
+        ),
+        blocks_per_chunk=3,
+        num_cpu_chunks=2,
+        mmap_region=region,
+        canonical_layout=canonical,
+        compact_group_layout=True,
+    )
+
+    def transfer(job, gpu, cpu, store):
+        if store:
+            assert worker.submit_store(job, gpu, cpu)
+        else:
+            assert worker.submit_load(job, cpu, gpu)
+        deadline = time.monotonic() + 10
+        finished: list[TransferResult] = []
+        while not finished and time.monotonic() < deadline:
+            finished = worker.get_finished()
+        assert len(finished) == 1 and finished[0].success
+        assert finished[0].job_id == job
+
+    try:
+        # Fill both slots; group 2 has only the middle sub-block of its chunk.
+        transfer(
+            1,
+            GPULoadStoreSpec([0, 1, 2, 4], [3, 0, 1], [0, 0, 1]),
+            CPULoadStoreSpec([0, 1]),
+            True,
+        )
+        packed.zero_()
+        transfer(
+            2,
+            GPULoadStoreSpec([9, 10, 11, 13], [3, 0, 1], [0, 0, 1]),
+            CPULoadStoreSpec([0, 1]),
+            False,
+        )
+        expected = torch.zeros_like(packed)
+        expected[9:12, :512] = original[:3, :512]
+        expected[9:12, 1024:] = original[:3, 1024:]
+        expected[13, :256] = original[4, :256]
+        expected[13, 1024:1280] = original[4, 1024:1280]
+        torch.testing.assert_close(packed, expected, rtol=0, atol=0)
+
+        # Evict group 0's key and reuse its slot for the differently sized group 1.
+        packed[6:9] = original[6:9]
+        transfer(
+            3,
+            GPULoadStoreSpec([6, 7, 8], [0, 3, 0], [0, 0, 0]),
+            CPULoadStoreSpec([0]),
+            True,
+        )
+        packed.zero_()
+        transfer(
+            4,
+            GPULoadStoreSpec([15, 16, 17, 13], [0, 3, 1], [0, 0, 1]),
+            CPULoadStoreSpec([0, 1]),
+            False,
+        )
+        expected.zero_()
+        expected[15:18, :1024] = original[6:9, :1024]
+        expected[13, :256] = original[4, :256]
+        expected[13, 1024:1280] = original[4, 1024:1280]
+        torch.testing.assert_close(packed, expected, rtol=0, atol=0)
+    finally:
+        worker.shutdown()

tests/v1/kv_offload/test_file_mapper.py

@@ -62,6 +62,8 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
     )
     spec = MagicMock(spec=OffloadingSpec)
     spec.config = config
+    spec.compact_group_layout = kwargs.get("storage_format") is not None
+    spec.storage_format = kwargs.get("storage_format")
     return FileMapper.from_offloading_spec(
         root_dir=kwargs.get("root_dir", "/tmp/cache"),
         offloading_spec=spec,
@@ -327,3 +329,22 @@ def test_replicated_layout_run_config_tp_invariant():
     tp2 = make_mapper_from_offloading_spec(tp_size=2, world_size=2, rank=0, **shared)
     tp4 = make_mapper_from_offloading_spec(tp_size=4, world_size=4, rank=2, **shared)
     assert tp2.get_run_config() == tp4.get_run_config()
+
+
+def test_compact_groups_isolate_persisted_row_geometry():
+    common = dict(
+        groups=((16, "layer0"), (16, "layer1")),
+        canonical_layout=True,
+        is_parallelism_agnostic=True,
+        parallel_agnostic=True,
+    )
+    key = make_offload_key(bytes(range(8)), 1)
+    # Even portable payloads cannot share files when whole-row transfers have
+    # different padding, or when groups used disjoint ranges in the old layout.
+    paths = {
+        make_mapper_from_offloading_spec(**common, storage_format=layout).get_file_name(
+            key
+        )
+        for layout in (None, "grouped-v1-4096", "grouped-v1-8192")
+    }
+    assert len(paths) == 3

tests/v1/kv_offload/tiering/test_fs_tier.py

@@ -65,6 +65,7 @@ def _make_offloading_spec(
     if world_size is None:
         world_size = tp_size
     spec = MagicMock()
+    spec.storage_format = None
     spec.config = OffloadingConfig(
         groups=(),
         worker_kv_bytes_per_block=0,

tests/v1/kv_offload/tiering/test_obj_tier.py

@@ -83,6 +83,7 @@ def _make_offloading_config(
 
 
 _OFFLOADING_SPEC = SimpleNamespace(
+    storage_format=None,
     config=_make_offloading_config(enable_kv_cache_events=False),
 )
 
@@ -218,6 +219,7 @@ def _query_memory(self, queries, mem_type, agent_name):
 def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace:
     """Offloading spec stub with an explicit global KV events flag."""
     return SimpleNamespace(
+        storage_format=None,
         config=_make_offloading_config(enable_kv_cache_events),
         kv_events_config=OffloadingKVEventsConfig(
             enable_kv_cache_events=enable_kv_cache_events,
@@ -769,6 +771,7 @@ def test_ca_bundle_included_when_set(self):
 def test_obj_tier_replicated_layout_collapses_mapper_identity():
     """TP=2 and TP=4 replicated configs share the obj FileMapper namespace."""
     tp2_spec = SimpleNamespace(
+        storage_format=None,
         config=_make_offloading_config(
             False, tp_size=2, world_size=2, rank=1, replicated_layout=True
         ),
@@ -778,6 +781,7 @@ def test_obj_tier_replicated_layout_collapses_mapper_identity():
         ),
     )
     tp4_spec = SimpleNamespace(
+        storage_format=None,
         config=_make_offloading_config(
             False, tp_size=4, world_size=4, rank=3, replicated_layout=True
         ),

vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py

@@ -82,6 +82,7 @@ def build_offloading_config(
                 parallel_config.decode_context_parallel_size,
             ),
             layer_names=tuple(group.layer_names),
+            kv_bytes_per_block=_group_kv_bytes_per_block(group),
         )
         for group_id, group in selected_groups
     )

vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py

@@ -139,8 +139,10 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             ),
             None,
         )
-        if packed_layer_name is not None and len(selected_groups) == len(
-            kv_cache_config.kv_cache_groups
+        if (
+            packed_layer_name is not None
+            and len(selected_groups) == len(kv_cache_config.kv_cache_groups)
+            and not self.spec.compact_group_layout
         ):
             (tensor,) = tensors_per_block[packed_layer_name]
             num_blocks = tensor.shape[0]

vllm/v1/kv_offload/base.py

@@ -576,6 +576,10 @@ def shutdown(self) -> None:
 class OffloadingSpec(ABC):
     """Spec for an offloading connector"""
 
+    # CPU slots can overlay groups because each OffloadKey belongs to one group.
+    compact_group_layout: bool = False
+    storage_format: str | None = None
+
     @classmethod
     def build_metric_definitions(
         cls, extra_config: dict[str, Any]

vllm/v1/kv_offload/config.py

@@ -16,6 +16,8 @@ class OffloadingGroupConfig:
     layer_names: tuple[str, ...]
     # Original KVCacheConfig group index.
     group_id: int
+    # Padded physical bytes per worker block for this group. Zero when unknown.
+    kv_bytes_per_block: int = 0
 
 
 @dataclass(frozen=True)