[KV Offload] Support packed HMA KV cache layout

tests/v1/core/test_contiguous_kv_packing.py

@@ -1,16 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for contiguous KV cache packing in _get_kv_cache_config_deepseek_v4."""
+"""Tests for contiguous KV cache packing."""
 
 from unittest.mock import MagicMock
 
 import pytest
 import torch
 
-from vllm.v1.core.kv_cache_utils import _get_kv_cache_config_deepseek_v4
+from vllm import envs
+from vllm.v1.core.kv_cache_utils import (
+    _get_kv_cache_config_deepseek_v4,
+    get_kv_cache_config_from_groups,
+)
 from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
     KVCacheGroupSpec,
+    KVCacheTensor,
     MLAAttentionSpec,
+    SlidingWindowSpec,
     UniformTypeKVCacheSpecs,
 )
 
@@ -28,6 +35,25 @@ def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec:
     )
 
 
+def _make_full_spec() -> FullAttentionSpec:
+    return FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=2,
+        head_size=64,
+        dtype=torch.float16,
+    )
+
+
+def _make_sw_spec() -> SlidingWindowSpec:
+    return SlidingWindowSpec(
+        block_size=16,
+        num_kv_heads=2,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=128,
+    )
+
+
 def _make_groups(n_c4, n_c128, n_swa):
     PS_C4_MLA = 37440
     PS_C4_IDX = 8640
@@ -130,6 +156,73 @@ def test_strided_views_are_independent(self):
         for i, v in enumerate(views):
             assert (v == i + 1).all(), f"View {i} was corrupted"
 
+    def test_hma_attention_groups_keep_default_backing(self, monkeypatch):
+        monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", False, raising=False)
+        full = _make_full_spec()
+        sw = _make_sw_spec()
+        page_size = full.page_size_bytes
+        groups = [
+            KVCacheGroupSpec(["full.0", "full.1"], full),
+            KVCacheGroupSpec(["sw.0", "sw.2"], sw),
+            KVCacheGroupSpec(["sw.1", "sw.3"], sw),
+        ]
+
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
+        )
+
+        assert config.num_blocks == 32
+        assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32
+        assert config.kv_cache_tensors == [
+            KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]),
+            KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]),
+        ]
+
+    def test_hma_attention_groups_use_packed_backing_with_flag(self, monkeypatch):
+        monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", True, raising=False)
+        full = _make_full_spec()
+        sw = _make_sw_spec()
+        page_size = full.page_size_bytes
+        groups = [
+            KVCacheGroupSpec(["full.0", "full.1"], full),
+            KVCacheGroupSpec(["sw.0", "sw.2"], sw),
+            KVCacheGroupSpec(["sw.1", "sw.3"], sw),
+        ]
+
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
+        )
+
+        assert config.num_blocks == 32
+        assert {t.size for t in config.kv_cache_tensors} == {page_size * 2 * 32}
+        assert config.kv_cache_tensors == [
+            KVCacheTensor(
+                size=page_size * 2 * 32,
+                shared_by=["full.0", "sw.0", "sw.1"],
+                offset=0,
+                block_stride=page_size * 2,
+            ),
+            KVCacheTensor(
+                size=page_size * 2 * 32,
+                shared_by=["full.1", "sw.2", "sw.3"],
+                offset=page_size,
+                block_stride=page_size * 2,
+            ),
+        ]
+
+    def test_single_group_attention_keeps_unpacked_layout(self):
+        spec = _make_full_spec()
+        groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)]
+
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32
+        )
+
+        assert sum(t.size for t in config.kv_cache_tensors) == (
+            spec.page_size_bytes * 2 * 32
+        )
+        assert [t.block_stride for t in config.kv_cache_tensors] == [0, 0]
+
 
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

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

@@ -50,15 +50,16 @@ def _register_handlers(self, kv_caches: CanonicalKVCaches):
     def register_kv_caches(
         self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
     ):
-        num_blocks = self.spec.kv_cache_config.num_blocks
+        kv_cache_config = self.spec.kv_cache_config
+        num_blocks = kv_cache_config.num_blocks
 
         # layer_name -> (num_blocks, page_size_bytes) tensor
         tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}
         # layer_name -> size of (un-padded) page in bytes
         unpadded_page_size_bytes: dict[str, int] = {}
         # layer_name -> size of page in bytes
         page_size_bytes: dict[str, int] = {}
-        for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:
+        for kv_cache_group in kv_cache_config.kv_cache_groups:
             group_layer_names = kv_cache_group.layer_names
             group_kv_cache_spec = kv_cache_group.kv_cache_spec
             if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs):
@@ -122,9 +123,35 @@ def register_kv_caches(
                 else:
                     raise NotImplementedError
 
+        packed_kv_cache_tensor = next(
+            (t for t in kv_cache_config.kv_cache_tensors if t.block_stride), None
+        )
+        is_dsv4 = all(
+            isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
+            for group in kv_cache_config.kv_cache_groups
+        )
+        if packed_kv_cache_tensor is not None and not is_dsv4:
+            (tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]]
+            block_stride = tensor.stride(0)
+            packed_tensor = tensor.as_strided(
+                (num_blocks, block_stride),
+                (block_stride, 1),
+                storage_offset=0,
+            )
+            self._register_handlers(
+                CanonicalKVCaches(
+                    [CanonicalKVCacheTensor(packed_tensor, block_stride)],
+                    [
+                        [CanonicalKVCacheRef(0, block_stride)]
+                        for _ in kv_cache_config.kv_cache_groups
+                    ],
+                )
+            )
+            return
+
         block_tensors: list[CanonicalKVCacheTensor] = []
         block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list)
-        for kv_cache_tensor in self.spec.kv_cache_config.kv_cache_tensors:
+        for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
             # Filter to layers that were actually processed above.
             # _get_kv_cache_config_deepseek_v4 emits KVCacheTensor entries for
             # every (tuple_idx, page_size) slot; slots where no group has a
@@ -166,7 +193,7 @@ def register_kv_caches(
                     )
 
         group_data_refs: list[list[CanonicalKVCacheRef]] = []
-        for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:
+        for kv_cache_group in kv_cache_config.kv_cache_groups:
             group_refs: list[CanonicalKVCacheRef] = []
             for layer_name in kv_cache_group.layer_names:
                 group_refs += block_data_refs[layer_name]

vllm/envs.py

@@ -209,6 +209,7 @@
     VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300
     VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5
     VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None
+    VLLM_USE_PACKED_HMA_KV_CACHE: bool = False
     VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None
     VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
     VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[
@@ -1608,6 +1609,11 @@ def _resolve_rust_frontend_path() -> str | None:
     "VLLM_KV_CACHE_LAYOUT": env_with_choices(
         "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"]
     ),
+    # Opt into packed per-block KV cache allocation for multi-group
+    # attention-only HMA models (e.g. gpt-oss, Gemma 3/4).
+    "VLLM_USE_PACKED_HMA_KV_CACHE": lambda: bool(
+        int(os.getenv("VLLM_USE_PACKED_HMA_KV_CACHE", "0"))
+    ),
     # SSM conv state layout used for Mamba models.
     # - SD: (state_len, dim) — dim contiguous (default)
     # - DS: (dim, state_len) — TP-sharded dim on dim1,

vllm/v1/core/kv_cache_utils.py

@@ -938,9 +938,7 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
         kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
     ):
         return kv_cache_groups[0].kv_cache_spec.page_size_bytes
-    if all(
-        isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in kv_cache_groups
-    ):
+    if _use_packed_kv_cache_groups(kv_cache_groups):
         # buckets = {page_size: [[layer_names], [layer_names], ...]}
         buckets = _bucket_layers_by_page_size(kv_cache_groups)
         return sum(ps * len(slots) for ps, slots in buckets.items())
@@ -1218,16 +1216,29 @@ def _bucket_layers_by_page_size(
     return buckets
 
 
-def _get_kv_cache_config_deepseek_v4(
+def _use_packed_kv_cache_groups(
+    kv_cache_groups: list[KVCacheGroupSpec],
+) -> bool:
+    is_dsv4 = all(
+        isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
+        for group in kv_cache_groups
+    )
+    return is_dsv4 or (
+        bool(envs.VLLM_USE_PACKED_HMA_KV_CACHE) and len(kv_cache_groups) > 1
+    )
+
+
+def _get_kv_cache_config_packed(
     vllm_config: VllmConfig,
     kv_cache_groups: list[KVCacheGroupSpec],
     available_memory: int,
 ) -> tuple[int, list[KVCacheTensor]]:
-    """DeepseekV4 KV cache tensor layout planning.
+    """Plan a packed per-block KV cache tensor layout.
 
     Emit one KVCacheTensor per (slot_idx, page_size). Layers from different
     groups at the same slot share a tensor (they have independent block
-    tables so block-id namespaces never collide).
+    tables so block-id namespaces never collide). Each emitted tensor aliases
+    one physical backing allocation, with per-block data laid out contiguously.
     """
     # buckets = {page_size: [[layer_names], [layer_names], ...]}
     buckets = _bucket_layers_by_page_size(kv_cache_groups)
@@ -1255,6 +1266,9 @@ def _get_kv_cache_config_deepseek_v4(
     return num_blocks, kv_cache_tensors
 
 
+_get_kv_cache_config_deepseek_v4 = _get_kv_cache_config_packed
+
+
 def get_kv_cache_config_from_groups(
     vllm_config: VllmConfig,
     kv_cache_groups: list[KVCacheGroupSpec],
@@ -1299,13 +1313,11 @@ def get_kv_cache_config_from_groups(
             )
             for layer_name in kv_cache_groups[0].layer_names
         ]
-    elif all(
-        isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
-        for group in kv_cache_groups
-    ):
-        # DeepseekV4: UniformTypeKVCacheSpecs but multiple groups.
-        # Delegate to the DeepseekV4-specific allocator.
-        num_blocks, kv_cache_tensors = _get_kv_cache_config_deepseek_v4(
+    elif _use_packed_kv_cache_groups(kv_cache_groups):
+        # DeepSeek V4 keeps the existing packed layout. Other multi-group
+        # attention-only HMA layouts can opt in with
+        # VLLM_USE_PACKED_HMA_KV_CACHE=1.
+        num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
             vllm_config, kv_cache_groups, available_memory
         )
     else:

vllm/v1/kv_offload/cpu/spec.py

@@ -58,7 +58,15 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
         self.cpu_page_size_per_worker = 0
         assert kv_cache_config is not None
         if kv_cache_config.num_blocks > 0 and world_size > 0:
-            total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors)
+            is_packed = any(t.block_stride for t in kv_cache_config.kv_cache_tensors)
+            assert not is_packed or all(
+                t.block_stride for t in kv_cache_config.kv_cache_tensors
+            )
+            total_gpu_kv_bytes = (
+                kv_cache_config.kv_cache_tensors[0].size
+                if is_packed
+                else sum(t.size for t in kv_cache_config.kv_cache_tensors)
+            )
             kv_bytes_per_block = (
                 total_gpu_kv_bytes // kv_cache_config.num_blocks
             ) * world_size

vllm/v1/simple_kv_offload/manager.py

@@ -187,14 +187,22 @@ def _derive_cpu_config(
 
         assert len(gpu_config.kv_cache_tensors) > 0
 
-        gpu_total_bytes = sum(t.size for t in gpu_config.kv_cache_tensors)
+        is_packed = any(t.block_stride for t in gpu_config.kv_cache_tensors)
+        assert not is_packed or all(t.block_stride for t in gpu_config.kv_cache_tensors)
+        gpu_total_bytes = (
+            gpu_config.kv_cache_tensors[0].size
+            if is_packed
+            else sum(t.size for t in gpu_config.kv_cache_tensors)
+        )
         num_gpu_blocks = gpu_config.num_blocks
         num_cpu_blocks = max(1, num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes)
         # Create CPU kv_cache_tensors mirroring GPU by scaling size proportionally.
         cpu_tensors = [
             KVCacheTensor(
                 size=t.size // num_gpu_blocks * num_cpu_blocks,
                 shared_by=list(t.shared_by),
+                offset=t.offset,
+                block_stride=t.block_stride,
             )
             for t in gpu_config.kv_cache_tensors
         ]