[KV Offload] Gate packed HMA KV cache on cross-layer config (#46252)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jun 24, 2026 at 11:55 UTC
e7df23228895c32982856051aeb50a89acb5d8c8
4 files changed
+42
-33
tests/v1/core/test_contiguous_kv_packing.py
+13
-10
@@ -7,9 +7,8 @@ from unittest.mock import MagicMock
7
import pytest
8
import torch
9
10
-from vllm import envs
10
from vllm.v1.core.kv_cache_utils import (
12
- _get_kv_cache_config_deepseek_v4,
11
+ _get_kv_cache_config_packed,
12
get_kv_cache_config_from_groups,
13
)
14
from vllm.v1.kv_cache_interface import (
@@ -84,15 +83,19 @@ def _make_groups(n_c4, n_c128, n_swa):
83
return [mla_group, swa_group]
84
85
87
-def _mock_vllm_config():
86
+def _mock_vllm_config(kv_connector_extra_config: dict[str, str] | None = None):
87
config = MagicMock()
88
config.cache_config.num_gpu_blocks_override = None
89
+ config.kv_transfer_config = None
90
+ if kv_connector_extra_config is not None:
91
+ config.kv_transfer_config = MagicMock()
92
+ config.kv_transfer_config.kv_connector_extra_config = kv_connector_extra_config
93
return config
94
95
96
def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024):
97
groups = _make_groups(n_c4, n_c128, n_swa)
95
- return _get_kv_cache_config_deepseek_v4(_mock_vllm_config(), groups, mem)
98
+ return _get_kv_cache_config_packed(_mock_vllm_config(), groups, mem)
99
100
101
def _page_sizes_by_layer(
@@ -135,7 +138,7 @@ class TestInterleavedPacking:
138
def test_strided_views_are_independent(self):
139
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
140
page_sizes = _page_sizes_by_layer(groups)
138
- num_blocks, tensors = _get_kv_cache_config_deepseek_v4(
141
+ num_blocks, tensors = _get_kv_cache_config_packed(
142
_mock_vllm_config(), groups, 100 * 1024 * 1024
143
)
144
backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
@@ -156,8 +159,7 @@ class TestInterleavedPacking:
159
for i, v in enumerate(views):
160
assert (v == i + 1).all(), f"View {i} was corrupted"
161
159
- def test_hma_attention_groups_keep_default_backing(self, monkeypatch):
160
- monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", False, raising=False)
162
+ def test_hma_attention_groups_keep_default_backing(self):
163
full = _make_full_spec()
164
sw = _make_sw_spec()
165
page_size = full.page_size_bytes
@@ -178,8 +180,7 @@ class TestInterleavedPacking:
180
KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]),
181
]
182
181
- def test_hma_attention_groups_use_packed_backing_with_flag(self, monkeypatch):
182
- monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", True, raising=False)
183
+ def test_hma_attention_groups_use_packed_backing_with_enable_cross_layers(self):
184
full = _make_full_spec()
185
sw = _make_sw_spec()
186
page_size = full.page_size_bytes
@@ -190,7 +191,9 @@ class TestInterleavedPacking:
191
]
192
193
config = get_kv_cache_config_from_groups(
193
- _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
194
+ _mock_vllm_config({"enable_cross_layers_blocks": "True"}),
195
+ groups,
196
+ available_memory=page_size * 2 * 32,
197
)
198
199
assert config.num_blocks == 32
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+8
-7
@@ -124,13 +124,14 @@ class OffloadingConnectorWorker:
124
raise NotImplementedError
125
126
packed_kv_cache_tensor = next(
127
- (t for t in kv_cache_config.kv_cache_tensors if t.block_stride), None
127
+ (
128
+ t
129
+ for t in kv_cache_config.kv_cache_tensors
130
+ if t.block_stride and t.shared_by
131
+ ),
132
+ None,
133
)
129
- is_dsv4 = all(
130
- isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
131
- for group in kv_cache_config.kv_cache_groups
132
- )
133
- if packed_kv_cache_tensor is not None and not is_dsv4:
134
+ if packed_kv_cache_tensor is not None:
135
(tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]]
136
block_stride = tensor.stride(0)
137
packed_tensor = tensor.as_strided(
@@ -153,7 +154,7 @@ class OffloadingConnectorWorker:
154
block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list)
155
for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
156
# Filter to layers that were actually processed above.
156
- # _get_kv_cache_config_deepseek_v4 emits KVCacheTensor entries for
157
+ # Packed KV allocation emits KVCacheTensor entries for
158
# every (tuple_idx, page_size) slot; slots where no group has a
159
# layer at that index produce an empty shared_by (reserved memory
160
# with no corresponding model layer).
vllm/envs.py
-6
@@ -210,7 +210,6 @@ if TYPE_CHECKING:
210
VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300
211
VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5
212
VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None
213
- VLLM_USE_PACKED_HMA_KV_CACHE: bool = False
213
VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None
214
VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
215
VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[
@@ -1614,11 +1613,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
1613
"VLLM_KV_CACHE_LAYOUT": env_with_choices(
1614
"VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"]
1615
),
1617
- # Opt into packed per-block KV cache allocation for multi-group
1618
- # attention-only HMA models (e.g. gpt-oss, Gemma 3/4).
1619
- "VLLM_USE_PACKED_HMA_KV_CACHE": lambda: bool(
1620
- int(os.getenv("VLLM_USE_PACKED_HMA_KV_CACHE", "0"))
1621
- ),
1616
# SSM conv state layout used for Mamba models.
1617
# - SD: (state_len, dim) — dim contiguous (default)
1618
# - DS: (dim, state_len) — TP-sharded dim on dim1,
vllm/v1/core/kv_cache_utils.py
+21
-10
@@ -947,7 +947,9 @@ def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int:
947
return num_blocks
948
949
950
-def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
950
+def _pool_bytes_per_block(
951
+ vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec]
952
+) -> int:
953
"""
954
Bytes consumed by one block in the worker's shared KV cache pool, mirroring
955
the divisor used by `get_kv_cache_config_from_groups` to convert
@@ -958,7 +960,7 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
960
kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
961
):
962
return kv_cache_groups[0].kv_cache_spec.page_size_bytes
961
- if _use_packed_kv_cache_groups(kv_cache_groups):
963
+ if _use_packed_kv_cache_config(vllm_config, kv_cache_groups):
964
# buckets = {page_size: [[layer_names], [layer_names], ...]}
965
buckets = _bucket_layers_by_page_size(kv_cache_groups)
966
return sum(ps * len(slots) for ps, slots in buckets.items())
@@ -1250,16 +1252,26 @@ def _bucket_layers_by_page_size(
1252
return buckets
1253
1254
1253
-def _use_packed_kv_cache_groups(
1255
+def _use_packed_kv_cache_config(
1256
+ vllm_config: VllmConfig,
1257
kv_cache_groups: list[KVCacheGroupSpec],
1258
) -> bool:
1259
is_dsv4 = all(
1260
isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
1261
for group in kv_cache_groups
1262
)
1260
- return is_dsv4 or (
1261
- bool(envs.VLLM_USE_PACKED_HMA_KV_CACHE) and len(kv_cache_groups) > 1
1263
+ kv_transfer_config = vllm_config.kv_transfer_config
1264
+ extra_config = (
1265
+ kv_transfer_config.kv_connector_extra_config
1266
+ if kv_transfer_config is not None
1267
+ else {}
1268
+ )
1269
+ # NOTE: enable_cross_layers_blocks is an experimental API and subject to change with
1270
+ # https://github.com/vllm-project/vllm/issues/42082
1271
+ enable_cross_layers = (
1272
+ str(extra_config.get("enable_cross_layers_blocks", "False")).lower() == "true"
1273
)
1274
+ return is_dsv4 or (enable_cross_layers and len(kv_cache_groups) > 1)
1275
1276
1277
def _get_kv_cache_config_packed(
@@ -1347,10 +1359,9 @@ def get_kv_cache_config_from_groups(
1359
)
1360
for layer_name in kv_cache_groups[0].layer_names
1361
]
1350
- elif _use_packed_kv_cache_groups(kv_cache_groups):
1351
- # DeepSeek V4 keeps the existing packed layout. Other multi-group
1352
- # attention-only HMA layouts can opt in with
1353
- # VLLM_USE_PACKED_HMA_KV_CACHE=1.
1362
+ elif _use_packed_kv_cache_config(vllm_config, kv_cache_groups):
1363
+ # DeepSeek V4 uses the packed layout by default. Other multi-group
1364
+ # layouts can opt in with --enable-cross-layers.
1365
num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
1366
vllm_config, kv_cache_groups, available_memory
1367
)
@@ -2069,7 +2080,7 @@ def get_kv_cache_configs(
2080
if not groups:
2081
adjusted_memory.append(avail_mem)
2082
continue
2072
- bytes_per_block = _pool_bytes_per_block(groups)
2083
+ bytes_per_block = _pool_bytes_per_block(vllm_config, groups)
2084
logger.info(
2085
"Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d",
2086
avail_mem // bytes_per_block,