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

[Core][DSV4] Compact MXFP4 indexer KV cache and packed group overlays (#48993)

Summer Yang committed Jul 22, 2026 at 16:58 UTC f3a920a0764012555fa0d3b6559ac0ddeb746200
3 files changed +216 -84
tests/v1/core/test_contiguous_kv_packing.py
+170 -22
@@ -9,6 +9,7 @@ import torch
9
10 from vllm.v1.core.kv_cache_utils import (
11 _get_kv_cache_config_packed,
12 + _get_kv_cache_groups_uniform_groups,
13 get_kv_cache_config_from_groups,
14 )
15 from vllm.v1.kv_cache_interface import (
@@ -16,6 +17,7 @@ from vllm.v1.kv_cache_interface import (
17 KVCacheGroupSpec,
18 KVCacheTensor,
19 MLAAttentionSpec,
20 + SlidingWindowMLASpec,
21 SlidingWindowSpec,
22 UniformTypeKVCacheSpecs,
23 )
@@ -109,7 +111,132 @@ def _page_sizes_by_layer(
111 return page_sizes
112
113
114 +def _packing_by_layer(
115 + tensors: list[KVCacheTensor],
116 +) -> dict[str, tuple[int, int]]:
117 + return {
118 + layer_name: (tensor.offset, tensor.block_stride)
119 + for tensor in tensors
120 + for layer_name in tensor.shared_by
121 + }
122 +
123 +
124 +def _make_views(
125 + groups: list[KVCacheGroupSpec],
126 + num_blocks: int,
127 + tensors: list[KVCacheTensor],
128 +) -> dict[str, torch.Tensor]:
129 + page_sizes = _page_sizes_by_layer(groups)
130 + packing = _packing_by_layer(tensors)
131 + backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
132 + return {
133 + layer_name: torch.as_strided(
134 + backing,
135 + size=(num_blocks, page_size),
136 + stride=(packing[layer_name][1], 1),
137 + storage_offset=packing[layer_name][0],
138 + )
139 + for layer_name, page_size in page_sizes.items()
140 + }
141 +
142 +
143 +def _make_page_group(prefix: str, page_sizes: list[int]) -> KVCacheGroupSpec:
144 + specs = {
145 + f"{prefix}.{i}": MagicMock(page_size_bytes=page_size)
146 + for i, page_size in enumerate(page_sizes)
147 + }
148 + return KVCacheGroupSpec(
149 + layer_names=list(specs),
150 + kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=specs),
151 + )
152 +
153 +
154 class TestInterleavedPacking:
155 + def test_compact_cache_overlays_fp32_state_group(self):
156 + full_specs = {}
157 + state_specs = {}
158 + for i in range(2):
159 + full_specs[f"mla.{i}"] = MLAAttentionSpec(
160 + block_size=256,
161 + num_kv_heads=1,
162 + head_size=512,
163 + dtype=torch.uint8,
164 + page_size_padded=32768,
165 + indexes_kv_by_block_stride=True,
166 + compress_ratio=4,
167 + )
168 + full_specs[f"indexer.{i}"] = MLAAttentionSpec(
169 + block_size=256,
170 + num_kv_heads=1,
171 + head_size=68,
172 + dtype=torch.uint8,
173 + page_size_padded=4608,
174 + compress_ratio=4,
175 + )
176 + state_specs[f"mla_state.{i}"] = SlidingWindowMLASpec(
177 + block_size=4,
178 + num_kv_heads=1,
179 + head_size=2048,
180 + dtype=torch.float32,
181 + sliding_window=8,
182 + indexes_kv_by_block_stride=True,
183 + )
184 + state_specs[f"indexer_state.{i}"] = SlidingWindowMLASpec(
185 + block_size=4,
186 + num_kv_heads=1,
187 + head_size=512,
188 + dtype=torch.float32,
189 + sliding_window=8,
190 + indexes_kv_by_block_stride=True,
191 + )
192 +
193 + grouped_specs = [
194 + UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=full_specs),
195 + UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=state_specs),
196 + ]
197 + groups = _get_kv_cache_groups_uniform_groups(grouped_specs)
198 +
199 + assert len(groups) == 2
200 + assert {full_specs[f"indexer.{i}"].page_size_bytes for i in range(2)} == {4608}
201 + assert {full_specs[f"indexer.{i}"].real_page_size_bytes for i in range(2)} == {
202 + 4352
203 + }
204 + assert {
205 + state_specs[f"indexer_state.{i}"].page_size_bytes for i in range(2)
206 + } == {8192}
207 +
208 + full_group_bytes = 2 * (32768 + 4608)
209 + state_group_bytes = 2 * (32768 + 8192)
210 + bytes_per_block = max(full_group_bytes, state_group_bytes)
211 + num_blocks, tensors = _get_kv_cache_config_packed(
212 + _mock_vllm_config(), groups, bytes_per_block * 32
213 + )
214 + assert num_blocks == 32
215 + assert {tensor.block_stride for tensor in tensors} == {bytes_per_block}
216 +
217 + packing = _packing_by_layer(tensors)
218 + assert packing["mla.0"][0] == packing["mla_state.0"][0] == 0
219 + assert packing["indexer.0"][0] == 32768
220 + assert packing["indexer_state.0"][0] == 32768
221 +
222 + def test_deepseek_v4_pro_stride(self):
223 + groups = [
224 + _make_page_group("full", [32768, 4608] * 30 + [1024] * 31),
225 + _make_page_group("c4_state", [32768, 8192] * 30),
226 + _make_page_group("c128_state", [32768] * 31),
227 + _make_page_group("swa.0", [32768] * 31),
228 + _make_page_group("swa.1", [32768] * 30),
229 + ]
230 + expected_stride = 1_228_800
231 +
232 + num_blocks, tensors = _get_kv_cache_config_packed(
233 + _mock_vllm_config(), groups, expected_stride * 32
234 + )
235 +
236 + assert num_blocks == 32
237 + assert {tensor.block_stride for tensor in tensors} == {expected_stride}
238 + assert {tensor.size for tensor in tensors} == {expected_stride * 32}
239 +
240 def test_all_tensors_have_block_stride(self):
241 _, tensors = _run()
242 for t in tensors:
@@ -122,9 +249,30 @@ class TestInterleavedPacking:
249 assert sizes.pop() > 0
250
251 def test_offsets_within_one_block(self):
125 - _, tensors = _run()
126 - for t in tensors:
127 - assert t.offset < t.block_stride
252 + groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
253 + _, tensors = _get_kv_cache_config_packed(
254 + _mock_vllm_config(), groups, 100 * 1024 * 1024
255 + )
256 + page_sizes = _page_sizes_by_layer(groups)
257 + packing = _packing_by_layer(tensors)
258 + for layer_name, page_size in page_sizes.items():
259 + offset, block_stride = packing[layer_name]
260 + assert offset + page_size <= block_stride
261 +
262 + def test_layouts_are_disjoint_within_each_group(self):
263 + groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
264 + _, tensors = _get_kv_cache_config_packed(
265 + _mock_vllm_config(), groups, 100 * 1024 * 1024
266 + )
267 + page_sizes = _page_sizes_by_layer(groups)
268 + packing = _packing_by_layer(tensors)
269 +
270 + for group in groups:
271 + ranges = sorted(
272 + (packing[name][0], packing[name][0] + page_sizes[name])
273 + for name in group.layer_names
274 + )
275 + assert all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:]))
276
277 def test_all_layers_accounted_for(self):
278 n_c4, n_c128, n_swa = 5, 4, 7
@@ -135,29 +283,29 @@ class TestInterleavedPacking:
283 expected = n_c4 * 2 + n_c128 + n_swa
284 assert len(all_names) == expected
285
138 - def test_strided_views_are_independent(self):
286 + def test_group_owned_blocks_do_not_alias(self):
287 groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
140 - page_sizes = _page_sizes_by_layer(groups)
288 num_blocks, tensors = _get_kv_cache_config_packed(
289 _mock_vllm_config(), groups, 100 * 1024 * 1024
290 )
144 - backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
145 - views = []
146 - for t in tensors:
147 - page_size = page_sizes[t.shared_by[0]]
148 - v = torch.as_strided(
149 - backing,
150 - size=(num_blocks, page_size),
151 - stride=(t.block_stride, 1),
152 - storage_offset=t.offset,
153 - )
154 - views.append(v)
155 -
156 - for i, v in enumerate(views):
157 - v.fill_(i + 1)
158 -
159 - for i, v in enumerate(views):
160 - assert (v == i + 1).all(), f"View {i} was corrupted"
291 + views = _make_views(groups, num_blocks, tensors)
292 +
293 + expected = {}
294 + value = 1
295 + for block_id, group in enumerate(groups):
296 + for layer_name in group.layer_names:
297 + views[layer_name][block_id].fill_(value)
298 + expected[layer_name] = (block_id, value)
299 + value += 1
300 +
301 + for layer_name, (block_id, value) in expected.items():
302 + assert (views[layer_name][block_id] == value).all()
303 +
304 + # Once the first group releases its block, another group may reuse it.
305 + for layer_name in groups[1].layer_names:
306 + views[layer_name][0].fill_(255)
307 + for layer_name in groups[1].layer_names:
308 + assert (views[layer_name][0] == 255).all()
309
310 def test_hma_attention_groups_keep_default_backing(self):
311 full = _make_full_spec()
vllm/models/deepseek_v4/attention.py
+11 -5
@@ -26,6 +26,7 @@ from vllm.models.deepseek_v4.common.ops import (
26 fused_indexer_q_rope_quant,
27 fused_q_kv_rmsnorm,
28 )
29 +from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
30
31 if TYPE_CHECKING:
32 from vllm.v1.attention.backends.mla.sparse_swa import (
@@ -727,11 +728,16 @@ class DeepseekV4Indexer(nn.Module):
728 )
729
730 assert cache_config is not None, "Deepseek V4 indexer requires cache_config"
730 - # NOTE(yifan): FP8 indxer cache use the same layout as V3.2:
731 - # head_dim bytes = 128 fp8 + 4 fp32 scale = 132.
732 - # For FP4 indexer cache, we still allocate the same amount of memory as FP8,
733 - # but only use the first half of the memory.
734 - k_cache_head_dim = self.head_dim + self.head_dim // self.quant_block_size * 4
731 + if self.use_fp4_kv:
732 + # MXFP4 stores two values per byte plus one UE8M0 byte per 32 values.
733 + # head_dim bytes = 64 packed values + 4 UE8M0 scales = 68.
734 + k_cache_head_dim = self.head_dim // 2 + self.head_dim // MXFP4_BLOCK_SIZE
735 + else:
736 + # NOTE(yifan): FP8 indexer cache uses the same layout as V3.2:
737 + # head_dim bytes = 128 fp8 + 4 fp32 scale = 132.
738 + k_cache_head_dim = (
739 + self.head_dim + self.head_dim // self.quant_block_size * 4
740 + )
741 self.k_cache = DeepseekV4IndexerCache(
742 head_dim=k_cache_head_dim,
743 dtype=torch.uint8,
vllm/v1/core/kv_cache_utils.py
+35 -57
@@ -983,9 +983,8 @@ def _pool_bytes_per_block(
983 ):
984 return kv_cache_groups[0].kv_cache_spec.page_size_bytes
985 if _use_packed_kv_cache_config(vllm_config, kv_cache_groups):
986 - # buckets = {page_size: [[layer_names], [layer_names], ...]}
987 - buckets = _bucket_layers_by_page_size(kv_cache_groups)
988 - return sum(ps * len(slots) for ps, slots in buckets.items())
986 + block_stride, _ = _get_packed_kv_cache_layout(kv_cache_groups)
987 + return block_stride
988 group_size = max(len(g.layer_names) for g in kv_cache_groups)
989 page_size = get_uniform_page_size([g.kv_cache_spec for g in kv_cache_groups])
990 return page_size * group_size
@@ -1260,29 +1259,29 @@ def _get_kv_cache_groups_uniform_page_size(
1259 return create_kv_cache_group_specs(kv_cache_spec, grouped_layers)
1260
1261
1263 -def _bucket_layers_by_page_size(
1262 +def _get_packed_kv_cache_layout(
1263 kv_cache_groups: list[KVCacheGroupSpec],
1265 -) -> dict[int, list[list[str]]]:
1266 - """Bucket layers by page size: ``result[ps][slot_idx] = [layer_names]``.
1264 +) -> tuple[int, dict[int, list[str]]]:
1265 + """Lay out each cache group densely in one shared block slab.
1266
1268 - Layers from different groups at the same ``slot_idx`` share an underlying tensor
1269 - (they have independent block tables so block-id namespaces never collide).
1267 + A block ID is owned by one cache group at a time, so layouts from different
1268 + groups may overlap. Layers within a group remain disjoint.
1269 """
1271 - buckets: dict[int, list[list[str]]] = defaultdict(list)
1270 + layers_by_offset: dict[int, list[str]] = defaultdict(list)
1271 + block_stride = 0
1272 for group in kv_cache_groups:
1273 spec = group.kv_cache_spec
1274 - slot_count: dict[int, int] = defaultdict(int)
1274 + byte_offset = 0
1275 for layer_name in group.layer_names:
1276 if isinstance(spec, UniformTypeKVCacheSpecs):
1277 - ps = spec.kv_cache_specs[layer_name].page_size_bytes
1277 + page_size = spec.kv_cache_specs[layer_name].page_size_bytes
1278 else:
1279 - ps = spec.page_size_bytes
1280 - slot_idx = slot_count[ps]
1281 - slot_count[ps] += 1
1282 - if slot_idx == len(buckets[ps]):
1283 - buckets[ps].append([])
1284 - buckets[ps][slot_idx].append(layer_name)
1285 - return buckets
1279 + page_size = spec.page_size_bytes
1280 + layers_by_offset[byte_offset].append(layer_name)
1281 + byte_offset += page_size
1282 + block_stride = max(block_stride, byte_offset)
1283 + assert block_stride > 0
1284 + return block_stride, layers_by_offset
1285
1286
1287 def _use_packed_kv_cache_config(
@@ -1314,33 +1313,26 @@ def _get_kv_cache_config_packed(
1313 ) -> tuple[int, list[KVCacheTensor]]:
1314 """Plan a packed per-block KV cache tensor layout.
1315
1317 - Emit one KVCacheTensor per (slot_idx, page_size). Layers from different
1318 - groups at the same slot share a tensor (they have independent block
1319 - tables so block-id namespaces never collide). Each emitted tensor aliases
1320 - one physical backing allocation, with per-block data laid out contiguously.
1316 + Cache groups use dense, overlapping layouts within one block slab. Each
1317 + emitted tensor aliases the same physical backing allocation.
1318 """
1322 - # buckets = {page_size: [[layer_names], [layer_names], ...]}
1323 - buckets = _bucket_layers_by_page_size(kv_cache_groups)
1324 - total_num_bytes_per_block = sum(ps * len(slots) for ps, slots in buckets.items())
1319 + block_stride, layers_by_offset = _get_packed_kv_cache_layout(kv_cache_groups)
1320
1326 - num_blocks = available_memory // total_num_bytes_per_block
1321 + num_blocks = available_memory // block_stride
1322 num_blocks = may_override_num_blocks(vllm_config, num_blocks)
1323
1329 - total_size = total_num_bytes_per_block * num_blocks
1324 + total_size = block_stride * num_blocks
1325
1326 kv_cache_tensors: list[KVCacheTensor] = []
1332 - byte_offset = 0
1333 - for ps, slots in buckets.items():
1334 - for slot in slots:
1335 - kv_cache_tensors.append(
1336 - KVCacheTensor(
1337 - size=total_size,
1338 - shared_by=slot,
1339 - offset=byte_offset,
1340 - block_stride=total_num_bytes_per_block,
1341 - )
1327 + for byte_offset in sorted(layers_by_offset):
1328 + kv_cache_tensors.append(
1329 + KVCacheTensor(
1330 + size=total_size,
1331 + shared_by=layers_by_offset[byte_offset],
1332 + offset=byte_offset,
1333 + block_stride=block_stride,
1334 )
1343 - byte_offset += ps
1335 + )
1336
1337 return num_blocks, kv_cache_tensors
1338
@@ -1651,29 +1643,15 @@ def _get_kv_cache_groups_uniform_groups(
1643 for spec in group.kv_cache_specs.values()
1644 )
1645
1654 - # Split each SWA UniformKV group into smaller groups to align their #(layer tuples)
1655 - # Possibly padding layer tuples for this.
1656 - # Additionally, we also pad KV blocks in each SWA layer, to align the page size
1657 - # with the corresponding layer in the full-MLA group.
1658 - all_page_sizes = full_mla_spec.get_page_sizes()
1646 + # Split each SWA UniformKV group into smaller groups to align their
1647 + # numbers of layer tuples. The packed block planner overlays groups, so
1648 + # their page sizes do not need to match.
1649 swa_mla_groups = []
1650 for sm_spec in swa_mla_specs:
1661 - sm_page_sizes = sm_spec.get_page_sizes()
1651 layers_per_size: dict[int, list[str]] = defaultdict(list)
1663 - assert max(sm_page_sizes) <= max(all_page_sizes)
1664 -
1665 - # Unify page size by padding layers' page_size to the nearest larger page_size.
1666 - # Compute candidate (nearest larger page_size) for each unique page size.
1667 - size_to_candidate: dict[int, int] = {}
1668 - for ps in sm_page_sizes:
1669 - size_to_candidate[ps] = min(x for x in all_page_sizes if x >= ps)
1670 - # Pad and collect layer names per page size.
1652 +
1653 for layer_name, layer_spec in sm_spec.kv_cache_specs.items():
1672 - current_size = layer_spec.page_size_bytes
1673 - candidate = size_to_candidate[current_size]
1674 - if current_size < candidate:
1675 - object.__setattr__(layer_spec, "page_size_padded", candidate)
1676 - layers_per_size[candidate].append(layer_name)
1654 + layers_per_size[layer_spec.page_size_bytes].append(layer_name)
1655 # NOTE(yifan): for now, inside a UniformKV group, each page_size should
1656 # have the same number of layers. This also means we don't need to pad layers
1657 # inside a partial-full layer tuple.