@tduncan / Nvidia-Optimization-Stack-VLLM / commits / cc22621b51

[KV Offload] Support packed HMA KV cache layout (#46205)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

Lucas Wilkinson committed Jun 20, 2026 at 17:19 UTC cc22621b51207e1af96269a840108ea654af9b42

6 files changed +175 -21

tests/v1/core/test_contiguous_kv_packing.py

+95 -2

@@ -1,16 +1,23 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -"""Tests for contiguous KV cache packing in _get_kv_cache_config_deepseek_v4."""

3 +"""Tests for contiguous KV cache packing."""

5 from unittest.mock import MagicMock

7 import pytest

8 import torch

10 -from vllm.v1.core.kv_cache_utils import _get_kv_cache_config_deepseek_v4

10 +from vllm import envs

11 +from vllm.v1.core.kv_cache_utils import (

12 + _get_kv_cache_config_deepseek_v4,

13 + get_kv_cache_config_from_groups,

14 +)

15 from vllm.v1.kv_cache_interface import (

16 + FullAttentionSpec,

17 KVCacheGroupSpec,

18 + KVCacheTensor,

19 MLAAttentionSpec,

20 + SlidingWindowSpec,

21 UniformTypeKVCacheSpecs,

22 )

@@ -28,6 +35,25 @@ def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec:

35 )

38 +def _make_full_spec() -> FullAttentionSpec:

39 + return FullAttentionSpec(

40 + block_size=16,

41 + num_kv_heads=2,

42 + head_size=64,

43 + dtype=torch.float16,

44 + )

45 +

46 +

47 +def _make_sw_spec() -> SlidingWindowSpec:

48 + return SlidingWindowSpec(

49 + block_size=16,

50 + num_kv_heads=2,

51 + head_size=64,

52 + dtype=torch.float16,

53 + sliding_window=128,

54 + )

55 +

56 +

57 def _make_groups(n_c4, n_c128, n_swa):

58 PS_C4_MLA = 37440

59 PS_C4_IDX = 8640

@@ -130,6 +156,73 @@ class TestInterleavedPacking:

156 for i, v in enumerate(views):

157 assert (v == i + 1).all(), f"View {i} was corrupted"

158

159 + def test_hma_attention_groups_keep_default_backing(self, monkeypatch):

160 + monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", False, raising=False)

161 + full = _make_full_spec()

162 + sw = _make_sw_spec()

163 + page_size = full.page_size_bytes

164 + groups = [

165 + KVCacheGroupSpec(["full.0", "full.1"], full),

166 + KVCacheGroupSpec(["sw.0", "sw.2"], sw),

167 + KVCacheGroupSpec(["sw.1", "sw.3"], sw),

168 + ]

169 +

170 + config = get_kv_cache_config_from_groups(

171 + _mock_vllm_config(), groups, available_memory=page_size * 2 * 32

172 + )

173 +

174 + assert config.num_blocks == 32

175 + assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32

176 + assert config.kv_cache_tensors == [

177 + KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]),

178 + KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]),

179 + ]

180 +

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 + full = _make_full_spec()

184 + sw = _make_sw_spec()

185 + page_size = full.page_size_bytes

186 + groups = [

187 + KVCacheGroupSpec(["full.0", "full.1"], full),

188 + KVCacheGroupSpec(["sw.0", "sw.2"], sw),

189 + KVCacheGroupSpec(["sw.1", "sw.3"], sw),

190 + ]

191 +

192 + config = get_kv_cache_config_from_groups(

193 + _mock_vllm_config(), groups, available_memory=page_size * 2 * 32

194 + )

195 +

196 + assert config.num_blocks == 32

197 + assert {t.size for t in config.kv_cache_tensors} == {page_size * 2 * 32}

198 + assert config.kv_cache_tensors == [

199 + KVCacheTensor(

200 + size=page_size * 2 * 32,

201 + shared_by=["full.0", "sw.0", "sw.1"],

202 + offset=0,

203 + block_stride=page_size * 2,

204 + ),

205 + KVCacheTensor(

206 + size=page_size * 2 * 32,

207 + shared_by=["full.1", "sw.2", "sw.3"],

208 + offset=page_size,

209 + block_stride=page_size * 2,

210 + ),

211 + ]

212 +

213 + def test_single_group_attention_keeps_unpacked_layout(self):

214 + spec = _make_full_spec()

215 + groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)]

216 +

217 + config = get_kv_cache_config_from_groups(

218 + _mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32

219 + )

220 +

221 + assert sum(t.size for t in config.kv_cache_tensors) == (

222 + spec.page_size_bytes * 2 * 32

223 + )

224 + assert [t.block_stride for t in config.kv_cache_tensors] == [0, 0]

225 +

226

227 if __name__ == "__main__":

228 pytest.main([__file__, "-v"])

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

+31 -4

@@ -50,7 +50,8 @@ class OffloadingConnectorWorker:

50 def register_kv_caches(

51 self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]

52 ):

53 - num_blocks = self.spec.kv_cache_config.num_blocks

53 + kv_cache_config = self.spec.kv_cache_config

54 + num_blocks = kv_cache_config.num_blocks

56 # layer_name -> (num_blocks, page_size_bytes) tensor

57 tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}

@@ -58,7 +59,7 @@ class OffloadingConnectorWorker:

59 unpadded_page_size_bytes: dict[str, int] = {}

60 # layer_name -> size of page in bytes

61 page_size_bytes: dict[str, int] = {}

61 - for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:

62 + for kv_cache_group in kv_cache_config.kv_cache_groups:

63 group_layer_names = kv_cache_group.layer_names

64 group_kv_cache_spec = kv_cache_group.kv_cache_spec

65 if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs):

@@ -122,9 +123,35 @@ class OffloadingConnectorWorker:

123 else:

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

128 + )

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 + (tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]]

135 + block_stride = tensor.stride(0)

136 + packed_tensor = tensor.as_strided(

137 + (num_blocks, block_stride),

138 + (block_stride, 1),

139 + storage_offset=0,

140 + )

141 + self._register_handlers(

142 + CanonicalKVCaches(

143 + [CanonicalKVCacheTensor(packed_tensor, block_stride)],

144 + [

145 + [CanonicalKVCacheRef(0, block_stride)]

146 + for _ in kv_cache_config.kv_cache_groups

147 + ],

148 + )

149 + )

150 + return

151 +

152 block_tensors: list[CanonicalKVCacheTensor] = []

153 block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list)

127 - for kv_cache_tensor in self.spec.kv_cache_config.kv_cache_tensors:

154 + for kv_cache_tensor in kv_cache_config.kv_cache_tensors:

155 # Filter to layers that were actually processed above.

156 # _get_kv_cache_config_deepseek_v4 emits KVCacheTensor entries for

157 # every (tuple_idx, page_size) slot; slots where no group has a

@@ -166,7 +193,7 @@ class OffloadingConnectorWorker:

193 )

194

195 group_data_refs: list[list[CanonicalKVCacheRef]] = []

169 - for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:

196 + for kv_cache_group in kv_cache_config.kv_cache_groups:

197 group_refs: list[CanonicalKVCacheRef] = []

198 for layer_name in kv_cache_group.layer_names:

199 group_refs += block_data_refs[layer_name]

vllm/envs.py

@@ -209,6 +209,7 @@ if TYPE_CHECKING:

209 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300

210 VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5

211 VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None

212 + 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[

@@ -1608,6 +1609,11 @@ environment_variables: dict[str, Callable[[], Any]] = {

1609 "VLLM_KV_CACHE_LAYOUT": env_with_choices(

1610 "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"]

1611 ),

1612 + # Opt into packed per-block KV cache allocation for multi-group

1613 + # attention-only HMA models (e.g. gpt-oss, Gemma 3/4).

1614 + "VLLM_USE_PACKED_HMA_KV_CACHE": lambda: bool(

1615 + int(os.getenv("VLLM_USE_PACKED_HMA_KV_CACHE", "0"))

1616 + ),

1617 # SSM conv state layout used for Mamba models.

1618 # - SD: (state_len, dim) — dim contiguous (default)

1619 # - DS: (dim, state_len) — TP-sharded dim on dim1,

vllm/v1/core/kv_cache_utils.py

+25 -13

@@ -938,9 +938,7 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:

938 kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs

939 ):

940 return kv_cache_groups[0].kv_cache_spec.page_size_bytes

941 - if all(

942 - isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in kv_cache_groups

943 - ):

941 + if _use_packed_kv_cache_groups(kv_cache_groups):

942 # buckets = {page_size: [[layer_names], [layer_names], ...]}

943 buckets = _bucket_layers_by_page_size(kv_cache_groups)

944 return sum(ps * len(slots) for ps, slots in buckets.items())

@@ -1218,16 +1216,29 @@ def _bucket_layers_by_page_size(

1216 return buckets

1217

1218

1221 -def _get_kv_cache_config_deepseek_v4(

1219 +def _use_packed_kv_cache_groups(

1220 + kv_cache_groups: list[KVCacheGroupSpec],

1221 +) -> bool:

1222 + is_dsv4 = all(

1223 + isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)

1224 + for group in kv_cache_groups

1225 + )

1226 + return is_dsv4 or (

1227 + bool(envs.VLLM_USE_PACKED_HMA_KV_CACHE) and len(kv_cache_groups) > 1

1228 + )

1229 +

1230 +

1231 +def _get_kv_cache_config_packed(

1232 vllm_config: VllmConfig,

1233 kv_cache_groups: list[KVCacheGroupSpec],

1234 available_memory: int,

1235 ) -> tuple[int, list[KVCacheTensor]]:

1226 - """DeepseekV4 KV cache tensor layout planning.

1236 + """Plan a packed per-block KV cache tensor layout.

1237

1238 Emit one KVCacheTensor per (slot_idx, page_size). Layers from different

1239 groups at the same slot share a tensor (they have independent block

1230 - tables so block-id namespaces never collide).

1240 + tables so block-id namespaces never collide). Each emitted tensor aliases

1241 + one physical backing allocation, with per-block data laid out contiguously.

1242 """

1243 # buckets = {page_size: [[layer_names], [layer_names], ...]}

1244 buckets = _bucket_layers_by_page_size(kv_cache_groups)

@@ -1255,6 +1266,9 @@ def _get_kv_cache_config_deepseek_v4(

1266 return num_blocks, kv_cache_tensors

1267

1268

1269 +_get_kv_cache_config_deepseek_v4 = _get_kv_cache_config_packed

1270 +

1271 +

1272 def get_kv_cache_config_from_groups(

1273 vllm_config: VllmConfig,

1274 kv_cache_groups: list[KVCacheGroupSpec],

@@ -1299,13 +1313,11 @@ def get_kv_cache_config_from_groups(

1313 )

1314 for layer_name in kv_cache_groups[0].layer_names

1315 ]

1302 - elif all(

1303 - isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)

1304 - for group in kv_cache_groups

1305 - ):

1306 - # DeepseekV4: UniformTypeKVCacheSpecs but multiple groups.

1307 - # Delegate to the DeepseekV4-specific allocator.

1308 - num_blocks, kv_cache_tensors = _get_kv_cache_config_deepseek_v4(

1316 + elif _use_packed_kv_cache_groups(kv_cache_groups):

1317 + # DeepSeek V4 keeps the existing packed layout. Other multi-group

1318 + # attention-only HMA layouts can opt in with

1319 + # VLLM_USE_PACKED_HMA_KV_CACHE=1.

1320 + num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(

1321 vllm_config, kv_cache_groups, available_memory

1322 )

1323 else:

vllm/v1/kv_offload/cpu/spec.py

+9 -1

@@ -58,7 +58,15 @@ class CPUOffloadingSpec(OffloadingSpec):

58 self.cpu_page_size_per_worker = 0

59 assert kv_cache_config is not None

60 if kv_cache_config.num_blocks > 0 and world_size > 0:

61 - total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors)

61 + is_packed = any(t.block_stride for t in kv_cache_config.kv_cache_tensors)

62 + assert not is_packed or all(

63 + t.block_stride for t in kv_cache_config.kv_cache_tensors

64 + )

65 + total_gpu_kv_bytes = (

66 + kv_cache_config.kv_cache_tensors[0].size

67 + if is_packed

68 + else sum(t.size for t in kv_cache_config.kv_cache_tensors)

69 + )

70 kv_bytes_per_block = (

71 total_gpu_kv_bytes // kv_cache_config.num_blocks

72 ) * world_size

vllm/v1/simple_kv_offload/manager.py

+9 -1

@@ -187,7 +187,13 @@ class SimpleCPUOffloadScheduler:

187

188 assert len(gpu_config.kv_cache_tensors) > 0

189

190 - gpu_total_bytes = sum(t.size for t in gpu_config.kv_cache_tensors)

190 + is_packed = any(t.block_stride for t in gpu_config.kv_cache_tensors)

191 + assert not is_packed or all(t.block_stride for t in gpu_config.kv_cache_tensors)

192 + gpu_total_bytes = (

193 + gpu_config.kv_cache_tensors[0].size

194 + if is_packed

195 + else sum(t.size for t in gpu_config.kv_cache_tensors)

196 + )

197 num_gpu_blocks = gpu_config.num_blocks

198 num_cpu_blocks = max(1, num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes)

199 # Create CPU kv_cache_tensors mirroring GPU by scaling size proportionally.

@@ -195,6 +201,8 @@ class SimpleCPUOffloadScheduler:

201 KVCacheTensor(

202 size=t.size // num_gpu_blocks * num_cpu_blocks,

203 shared_by=list(t.shared_by),

204 + offset=t.offset,

205 + block_stride=t.block_stride,

206 )

207 for t in gpu_config.kv_cache_tensors

208 ]