[DSv4] Pack KV caches into contiguous per-block allocations for DeepSeek V4 (#44577)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Tyler Michael Smith committed
Jun 19, 2026 at 12:55 UTC
01192139bf022bec84e2cca3a3e36e8bb5293b5c
7 files changed
+344
-22
tests/v1/core/test_contiguous_kv_packing.py
new
+135
@@ -0,0 +1,135 @@
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."""
4
+
5
+from unittest.mock import MagicMock
6
+
7
+import pytest
8
+import torch
9
+
10
+from vllm.v1.core.kv_cache_utils import _get_kv_cache_config_deepseek_v4
11
+from vllm.v1.kv_cache_interface import (
12
+ KVCacheGroupSpec,
13
+ MLAAttentionSpec,
14
+ UniformTypeKVCacheSpecs,
15
+)
16
+
17
+
18
+def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec:
19
+ return MLAAttentionSpec(
20
+ block_size=block_size,
21
+ num_kv_heads=1,
22
+ head_size=512,
23
+ dtype=torch.uint8,
24
+ page_size_padded=page_size,
25
+ cache_dtype_str="fp8_ds_mla",
26
+ model_version="deepseek_v4",
27
+ alignment=576,
28
+ )
29
+
30
+
31
+def _make_groups(n_c4, n_c128, n_swa):
32
+ PS_C4_MLA = 37440
33
+ PS_C4_IDX = 8640
34
+ PS_C128 = 1728
35
+ PS_SWA = 37440
36
+
37
+ mla_specs = {}
38
+ for i in range(n_c4):
39
+ mla_specs[f"c4_mla.{i}"] = _make_mla_spec(PS_C4_MLA)
40
+ mla_specs[f"c4_idx.{i}"] = _make_mla_spec(PS_C4_IDX)
41
+ for i in range(n_c128):
42
+ mla_specs[f"c128_mla.{i}"] = _make_mla_spec(PS_C128)
43
+
44
+ mla_group = KVCacheGroupSpec(
45
+ layer_names=list(mla_specs.keys()),
46
+ kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=mla_specs),
47
+ )
48
+
49
+ swa_specs = {}
50
+ for i in range(n_swa):
51
+ swa_specs[f"swa.{i}"] = _make_mla_spec(PS_SWA)
52
+
53
+ swa_group = KVCacheGroupSpec(
54
+ layer_names=list(swa_specs.keys()),
55
+ kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=swa_specs),
56
+ )
57
+
58
+ return [mla_group, swa_group]
59
+
60
+
61
+def _mock_vllm_config():
62
+ config = MagicMock()
63
+ config.cache_config.num_gpu_blocks_override = None
64
+ return config
65
+
66
+
67
+def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024):
68
+ groups = _make_groups(n_c4, n_c128, n_swa)
69
+ return _get_kv_cache_config_deepseek_v4(_mock_vllm_config(), groups, mem)
70
+
71
+
72
+def _page_sizes_by_layer(
73
+ groups: list[KVCacheGroupSpec],
74
+) -> dict[str, int]:
75
+ page_sizes = {}
76
+ for group in groups:
77
+ specs = group.kv_cache_spec.kv_cache_specs
78
+ for layer_name in group.layer_names:
79
+ page_sizes[layer_name] = specs[layer_name].page_size_bytes
80
+ return page_sizes
81
+
82
+
83
+class TestInterleavedPacking:
84
+ def test_all_tensors_have_block_stride(self):
85
+ _, tensors = _run()
86
+ for t in tensors:
87
+ assert t.block_stride > 0
88
+
89
+ def test_all_tensors_share_same_size(self):
90
+ _, tensors = _run()
91
+ sizes = set(t.size for t in tensors)
92
+ assert len(sizes) == 1
93
+ assert sizes.pop() > 0
94
+
95
+ def test_offsets_within_one_block(self):
96
+ _, tensors = _run()
97
+ for t in tensors:
98
+ assert t.offset < t.block_stride
99
+
100
+ def test_all_layers_accounted_for(self):
101
+ n_c4, n_c128, n_swa = 5, 4, 7
102
+ _, tensors = _run(n_c4=n_c4, n_c128=n_c128, n_swa=n_swa)
103
+ all_names = set()
104
+ for t in tensors:
105
+ all_names.update(t.shared_by)
106
+ expected = n_c4 * 2 + n_c128 + n_swa
107
+ assert len(all_names) == expected
108
+
109
+ def test_strided_views_are_independent(self):
110
+ groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
111
+ page_sizes = _page_sizes_by_layer(groups)
112
+ num_blocks, tensors = _get_kv_cache_config_deepseek_v4(
113
+ _mock_vllm_config(), groups, 100 * 1024 * 1024
114
+ )
115
+ backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
116
+ views = []
117
+ for t in tensors:
118
+ page_size = page_sizes[t.shared_by[0]]
119
+ v = torch.as_strided(
120
+ backing,
121
+ size=(num_blocks, page_size),
122
+ stride=(t.block_stride, 1),
123
+ storage_offset=t.offset,
124
+ )
125
+ views.append(v)
126
+
127
+ for i, v in enumerate(views):
128
+ v.fill_(i + 1)
129
+
130
+ for i, v in enumerate(views):
131
+ assert (v == i + 1).all(), f"View {i} was corrupted"
132
+
133
+
134
+if __name__ == "__main__":
135
+ pytest.main([__file__, "-v"])
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+98
@@ -841,8 +841,106 @@ class NixlBaseConnectorWorker:
841
# Forwarding a real layer name rather than a synthetic key
842
self.register_kv_caches({first_layer: kv_cache})
843
844
+ def _register_packed_kv_cache(
845
+ self,
846
+ storage: torch.UntypedStorage,
847
+ ) -> None:
848
+ """Register a packed KV cache as a single NIXL region.
849
+
850
+ The packed allocation interleaves all layers per block, so each
851
+ block_stride-byte chunk is one logical block. We register 1
852
+ NIXL region and create 1 descriptor per block.
853
+ """
854
+ self.transfer_topo = TransferTopology(
855
+ tp_rank=self.tp_rank,
856
+ tp_size=self.world_size,
857
+ block_size=self.block_size,
858
+ engine_id=self.engine_id,
859
+ is_mla=self.use_mla,
860
+ total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
861
+ attn_backends=self.attn_backends,
862
+ tensor_shape=None,
863
+ is_mamba=self._has_mamba,
864
+ )
865
+ self.compat_hash = compute_nixl_compatibility_hash(
866
+ self.vllm_config,
867
+ self.backend_name,
868
+ self.transfer_topo.cross_layers_blocks,
869
+ )
870
+
871
+ total_size = storage.nbytes()
872
+ block_stride = total_size // self.num_blocks
873
+ base_addr = storage.data_ptr()
874
+ device_id = storage.device.index
875
+ assert device_id is not None
876
+
877
+ logger.info(
878
+ "Registering packed KV cache: total_size=%s, block_stride=%s, "
879
+ "num_blocks=%s, num_regions=1",
880
+ total_size,
881
+ block_stride,
882
+ self.num_blocks,
883
+ )
884
+
885
+ self.device_id = device_id
886
+ caches_data = [(base_addr, total_size, self.device_id, "")]
887
+
888
+ self.block_len_per_layer = [block_stride]
889
+ self.num_regions = 1
890
+ self.num_descs = self.num_blocks
891
+ self.kv_caches_base_addr[self.engine_id][self.tp_rank] = [base_addr]
892
+
893
+ descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
894
+ self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
895
+ self._registered_descs.append(descs)
896
+
897
+ self.dst_num_blocks[self.engine_id] = self.num_blocks
898
+
899
+ self.src_xfer_handles_by_block_size[self.block_size], (self.src_blocks_data) = (
900
+ self.register_local_xfer_handler(self.block_size)
901
+ )
902
+
903
+ agent_metadata = NixlAgentMetadata(
904
+ engine_id=self.engine_id,
905
+ agent_metadata=self.nixl_wrapper.get_agent_metadata(),
906
+ device_id=self.device_id,
907
+ kv_caches_base_addr=(
908
+ self.kv_caches_base_addr[self.engine_id][self.tp_rank]
909
+ ),
910
+ num_blocks=self.num_blocks,
911
+ block_lens=self.block_len_per_layer,
912
+ kv_cache_layout=self.kv_cache_layout,
913
+ block_size=self.block_size,
914
+ ssm_sizes=self._mamba_ssm_size,
915
+ attn_backend_name=self.backend_name,
916
+ physical_blocks_per_logical_kv_block=(
917
+ self._physical_blocks_per_logical_kv_block
918
+ ),
919
+ )
920
+ assert self.compat_hash is not None
921
+ encoder = msgspec.msgpack.Encoder()
922
+ self.xfer_handshake_metadata = NixlHandshakePayload(
923
+ compatibility_hash=self.compat_hash,
924
+ agent_metadata_bytes=encoder.encode(agent_metadata),
925
+ )
926
+
927
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
928
"""Register the KV Cache data in nixl."""
929
+
930
+ # Detect packed allocation: all tensors are strided views into the
931
+ # same backing storage (different data_ptr but same storage).
932
+ # This happens with DSv4-style contiguous per-block packing.
933
+ if len(kv_caches) > 1 and not self._has_mamba:
934
+ storage = next(iter(kv_caches.values())).untyped_storage()
935
+ storage_ptrs = {
936
+ cache.untyped_storage().data_ptr() for cache in kv_caches.values()
937
+ }
938
+ data_ptrs = {cache.data_ptr() for cache in kv_caches.values()}
939
+ if len(storage_ptrs) == 1 and len(data_ptrs) > 1:
940
+ self._register_packed_kv_cache(storage)
941
+ self.device_kv_caches = kv_caches
942
+ return
943
+
944
self.transfer_topo = TransferTopology(
945
tp_rank=self.tp_rank,
946
tp_size=self.world_size,
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+9
-5
@@ -72,18 +72,22 @@ class OffloadingConnectorWorker:
72
if isinstance(layer_kv_cache_spec, AttentionSpec):
73
layer_kv_cache = kv_caches[layer_name]
74
assert isinstance(layer_kv_cache, torch.Tensor)
75
- assert layer_kv_cache.storage_offset() == 0
75
77
- storage = layer_kv_cache.untyped_storage()
76
page = layer_kv_cache_spec.page_size_bytes
77
+ elem_size = layer_kv_cache.element_size()
78
+ byte_offset = layer_kv_cache.storage_offset() * elem_size
79
+ block_stride_bytes = layer_kv_cache.stride(0) * elem_size
80
tensors_per_block[layer_name] = (
81
torch.tensor(
82
[],
83
dtype=torch.int8,
84
device=layer_kv_cache.device,
84
- )
85
- .set_(storage)
86
- .view(num_blocks, page),
85
+ ).set_(
86
+ layer_kv_cache.untyped_storage(),
87
+ byte_offset,
88
+ (num_blocks, page),
89
+ (block_stride_bytes, 1),
90
+ ),
91
)
92
page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
93
unpadded_page_size_bytes[layer_name] = (
vllm/v1/core/kv_cache_utils.py
+12
-1
@@ -1236,10 +1236,21 @@ def _get_kv_cache_config_deepseek_v4(
1236
num_blocks = available_memory // total_num_bytes_per_block
1237
num_blocks = may_override_num_blocks(vllm_config, num_blocks)
1238
1239
+ total_size = total_num_bytes_per_block * num_blocks
1240
+
1241
kv_cache_tensors: list[KVCacheTensor] = []
1242
+ byte_offset = 0
1243
for ps, slots in buckets.items():
1244
for slot in slots:
1242
- kv_cache_tensors.append(KVCacheTensor(size=ps * num_blocks, shared_by=slot))
1245
+ kv_cache_tensors.append(
1246
+ KVCacheTensor(
1247
+ size=total_size,
1248
+ shared_by=slot,
1249
+ offset=byte_offset,
1250
+ block_stride=total_num_bytes_per_block,
1251
+ )
1252
+ )
1253
+ byte_offset += ps
1254
1255
return num_blocks, kv_cache_tensors
1256
vllm/v1/kv_cache_interface.py
+2
@@ -847,6 +847,8 @@ class KVCacheTensor:
847
848
size: int # size of the KV cache tensor in bytes
849
shared_by: list[str] # layer names that share the same KV cache tensor
850
+ offset: int = 0 # byte offset of this layer within a contiguous block
851
+ block_stride: int = 0 # total bytes per block in a packed layout (0 = not packed)
852
853
854
@dataclass
vllm/v1/worker/gpu/attn_utils.py
+41
-7
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable, Sequence
4
from dataclasses import dataclass
5
+from math import prod
6
from typing import Any, cast
7
8
import torch
@@ -155,8 +156,17 @@ def _allocate_kv_cache(
156
kv_cache_config: KVCacheConfig, shared_layers: dict[str, str], device: torch.device
157
):
158
kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
159
+ packed_backing: torch.Tensor | None = None
160
for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
159
- tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device)
161
+ if kv_cache_tensor.block_stride > 0:
162
+ # Allocate once; all packed tensors alias the same backing.
163
+ if packed_backing is None:
164
+ packed_backing = torch.zeros(
165
+ kv_cache_tensor.size, dtype=torch.int8, device=device
166
+ )
167
+ tensor = packed_backing
168
+ else:
169
+ tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device)
170
for layer_name in kv_cache_tensor.shared_by:
171
kv_cache_raw_tensors[layer_name] = tensor
172
@@ -176,10 +186,18 @@ def _reshape_kv_cache(
186
cache_dtype: str,
187
kernel_block_sizes: list[int],
188
shared_kv_cache_layers: dict[str, str],
189
+ kv_cache_config: "KVCacheConfig | None" = None,
190
) -> dict[str, Any]:
191
kv_caches: dict[str, Any] = {}
192
has_attn, has_mamba = False, False
193
194
+ layer_packing: dict[str, tuple[int, int]] = {}
195
+ if kv_cache_config is not None:
196
+ for kv_tensor in kv_cache_config.kv_cache_tensors:
197
+ if kv_tensor.block_stride > 0:
198
+ for ln in kv_tensor.shared_by:
199
+ layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride)
200
+
201
for group in attn_groups:
202
if group.kv_cache_group_id >= len(kernel_block_sizes):
203
continue
@@ -198,8 +216,13 @@ def _reshape_kv_cache(
216
continue
217
218
kv_raw_tensor = kv_cache_raw_tensors[layer_name]
201
- assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
202
- num_blocks = kv_raw_tensor.numel() // kv_cache_spec.page_size_bytes
219
+ packing = layer_packing.get(layer_name)
220
+ if packing is not None:
221
+ _, blk_stride = packing
222
+ num_blocks = kv_raw_tensor.numel() // blk_stride
223
+ else:
224
+ assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
225
+ num_blocks = kv_raw_tensor.numel() // kv_cache_spec.page_size_bytes
226
227
if isinstance(kv_cache_spec, AttentionSpec):
228
has_attn = True
@@ -232,8 +255,18 @@ def _reshape_kv_cache(
255
]
256
257
dtype = kv_cache_spec.dtype
235
- kv_tensor = kv_raw_tensor.view(dtype)
236
- if kv_cache_spec.page_size_padded is not None:
258
+ if packing is not None:
259
+ offset, block_stride = packing
260
+ assert inv_order[0] == 0
261
+ page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype)
262
+ kv_cache = (
263
+ kv_raw_tensor.view(-1, block_stride)[
264
+ :, offset : offset + page_bytes
265
+ ]
266
+ .view(dtype)
267
+ .view(kv_cache_shape)
268
+ )
269
+ elif kv_cache_spec.page_size_padded is not None:
270
# Use strided view to handle page_size_bytes that
271
# include padding. This follows the same pattern as
272
# MambaSpec handling in gpu_model_runner.py.
@@ -246,13 +279,13 @@ def _reshape_kv_cache(
279
strides = list(torch.empty(kv_cache_shape).stride())
280
strides[inv_order[0]] = page_stride
281
kv_cache = torch.as_strided(
249
- kv_tensor,
282
+ kv_raw_tensor.view(dtype),
283
size=kv_cache_shape,
284
stride=tuple(strides),
285
)
286
else:
287
# No padding — safe to use a contiguous view.
255
- kv_cache = kv_tensor.view(kv_cache_shape)
288
+ kv_cache = kv_raw_tensor.view(dtype).view(kv_cache_shape)
289
kv_caches[layer_name] = kv_cache.permute(*inv_order)
290
291
elif isinstance(kv_cache_spec, MambaSpec):
@@ -365,6 +398,7 @@ def init_kv_cache(
398
kernel_block_sizes=kernel_block_sizes,
399
cache_dtype=cache_dtype,
400
shared_kv_cache_layers=shared_kv_cache_layers,
401
+ kv_cache_config=kv_cache_config,
402
)
403
bind_kv_cache(kv_caches, forward_context, runner_kv_caches)
404
return kv_caches
vllm/v1/worker/gpu_model_runner.py
+47
-9
@@ -12,6 +12,7 @@ from contextlib import contextmanager
12
from copy import copy, deepcopy
13
from dataclasses import dataclass, replace
14
from functools import reduce
15
+from math import prod
16
from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast
17
18
import numpy as np
@@ -7029,10 +7030,21 @@ class GPUModelRunner(
7030
corresponding memory buffer for KV cache.
7031
"""
7032
kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
7033
+ packed_backing: torch.Tensor | None = None
7034
for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
7033
- tensor = torch.zeros(
7034
- kv_cache_tensor.size, dtype=torch.int8, device=self.device
7035
- )
7035
+ if kv_cache_tensor.block_stride > 0:
7036
+ # Allocate once; all packed tensors alias the same backing.
7037
+ if packed_backing is None:
7038
+ packed_backing = torch.zeros(
7039
+ kv_cache_tensor.size,
7040
+ dtype=torch.int8,
7041
+ device=self.device,
7042
+ )
7043
+ tensor = packed_backing
7044
+ else:
7045
+ tensor = torch.zeros(
7046
+ kv_cache_tensor.size, dtype=torch.int8, device=self.device
7047
+ )
7048
for layer_name in kv_cache_tensor.shared_by:
7049
kv_cache_raw_tensors[layer_name] = tensor
7050
@@ -7074,6 +7086,14 @@ class GPUModelRunner(
7086
"""
7087
kv_caches: dict[str, torch.Tensor] = {}
7088
has_attn, has_mamba = False, False
7089
+
7090
+ # Map layer names to (offset, block_stride) within the packed
7091
+ # backing tensor so we can create strided views per layer.
7092
+ layer_packing: dict[str, tuple[int, int]] = {}
7093
+ for kv_tensor in self.kv_cache_config.kv_cache_tensors:
7094
+ if kv_tensor.block_stride > 0:
7095
+ for ln in kv_tensor.shared_by:
7096
+ layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride)
7097
for group in self._kv_cache_spec_attn_group_iterator():
7098
kv_cache_spec = group.kv_cache_spec
7099
attn_backend = group.backend
@@ -7085,8 +7105,13 @@ class GPUModelRunner(
7105
if layer_name in self.runner_only_attn_layers:
7106
continue
7107
raw_tensor = kv_cache_raw_tensors[layer_name]
7088
- assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
7089
- num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes
7108
+ packing = layer_packing.get(layer_name)
7109
+ if packing is not None:
7110
+ _, blk_stride = packing
7111
+ num_blocks = raw_tensor.numel() // blk_stride
7112
+ else:
7113
+ assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
7114
+ num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes
7115
if isinstance(kv_cache_spec, AttentionSpec):
7116
has_attn = True
7117
num_blocks_per_kv_block = (
@@ -7127,8 +7152,17 @@ class GPUModelRunner(
7152
for i in range(len(kv_cache_stride_order))
7153
]
7154
7130
- raw_tensor = kv_cache_raw_tensors[layer_name].view(dtype)
7131
- if kv_cache_spec.page_size_padded is not None:
7155
+ if packing is not None:
7156
+ offset, block_stride = packing
7157
+ assert inv_order[0] == 0
7158
+ page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype)
7159
+ kv_cache = (
7160
+ kv_cache_raw_tensors[layer_name]
7161
+ .view(-1, block_stride)[:, offset : offset + page_bytes]
7162
+ .view(dtype)
7163
+ .view(kv_cache_shape)
7164
+ )
7165
+ elif kv_cache_spec.page_size_padded is not None:
7166
# Use strided view to handle page_size_bytes that
7167
# include padding. This follows
7168
# the same pattern as MambaSpec handling below.
@@ -7142,13 +7176,17 @@ class GPUModelRunner(
7176
strides = list(torch.empty(kv_cache_shape).stride())
7177
strides[inv_order[0]] = page_stride
7178
kv_cache = torch.as_strided(
7145
- raw_tensor,
7179
+ kv_cache_raw_tensors[layer_name].view(dtype),
7180
size=kv_cache_shape,
7181
stride=tuple(strides),
7182
)
7183
else:
7184
# No padding — safe to use a contiguous view.
7151
- kv_cache = raw_tensor.view(kv_cache_shape)
7185
+ kv_cache = (
7186
+ kv_cache_raw_tensors[layer_name]
7187
+ .view(dtype)
7188
+ .view(kv_cache_shape)
7189
+ )
7190
kv_caches[layer_name] = kv_cache.permute(*inv_order)
7191
7192
elif isinstance(kv_cache_spec, MambaSpec):