A high-throughput and memory-efficient inference and serving engine for LLMs
[1/N][KV-Cache Layout Refactor] Refactor DSV4 KV cache config construction (#44454)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Lucas Wilkinson committed
Jun 7, 2026 at 08:53 UTC
4dcd10eb0d223a3ec4b2c96deaf3a48a96c8dcaa
1 file changed
+38
-49
vllm/v1/core/kv_cache_utils.py
+38
-49
@@ -941,14 +941,9 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
941
if all(
942
isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in kv_cache_groups
943
):
944
- # DeepseekV4: shared layout sized by the largest per-page-size bucket.
945
- full_mla_spec = cast(UniformTypeKVCacheSpecs, kv_cache_groups[0].kv_cache_spec)
946
- layer_tuple_page_bytes = sum(full_mla_spec.get_page_sizes())
947
- num_layer_tuples = max(
948
- cast(UniformTypeKVCacheSpecs, g.kv_cache_spec).get_num_layer_tuples()
949
- for g in kv_cache_groups
950
- )
951
- return layer_tuple_page_bytes * num_layer_tuples
944
+ # buckets = {page_size: [[layer_names], [layer_names], ...]}
945
+ buckets = _bucket_layers_by_page_size(kv_cache_groups)
946
+ return sum(ps * len(slots) for ps, slots in buckets.items())
947
group_size = max(len(g.layer_names) for g in kv_cache_groups)
948
page_size = get_uniform_page_size([g.kv_cache_spec for g in kv_cache_groups])
949
return page_size * group_size
@@ -1198,6 +1193,31 @@ def _get_kv_cache_groups_uniform_page_size(
1193
return create_kv_cache_group_specs(kv_cache_spec, grouped_layers)
1194
1195
1196
+def _bucket_layers_by_page_size(
1197
+ kv_cache_groups: list[KVCacheGroupSpec],
1198
+) -> dict[int, list[list[str]]]:
1199
+ """Bucket layers by page size: ``result[ps][slot_idx] = [layer_names]``.
1200
+
1201
+ Layers from different groups at the same ``slot_idx`` share an underlying tensor
1202
+ (they have independent block tables so block-id namespaces never collide).
1203
+ """
1204
+ buckets: dict[int, list[list[str]]] = defaultdict(list)
1205
+ for group in kv_cache_groups:
1206
+ spec = group.kv_cache_spec
1207
+ slot_count: dict[int, int] = defaultdict(int)
1208
+ for layer_name in group.layer_names:
1209
+ if isinstance(spec, UniformTypeKVCacheSpecs):
1210
+ ps = spec.kv_cache_specs[layer_name].page_size_bytes
1211
+ else:
1212
+ ps = spec.page_size_bytes
1213
+ slot_idx = slot_count[ps]
1214
+ slot_count[ps] += 1
1215
+ if slot_idx == len(buckets[ps]):
1216
+ buckets[ps].append([])
1217
+ buckets[ps][slot_idx].append(layer_name)
1218
+ return buckets
1219
+
1220
+
1221
def _get_kv_cache_config_deepseek_v4(
1222
vllm_config: VllmConfig,
1223
kv_cache_groups: list[KVCacheGroupSpec],
@@ -1205,52 +1225,21 @@ def _get_kv_cache_config_deepseek_v4(
1225
) -> tuple[int, list[KVCacheTensor]]:
1226
"""DeepseekV4 KV cache tensor layout planning.
1227
1208
- Precondition: kv_cache_groups[0] is the full-MLA group; its page sizes
1209
- define the canonical bucket set. Non-full-MLA groups must have been
1210
- page_size-padded upstream (see _get_kv_cache_groups_uniform_groups) so
1211
- every layer's page_size matches one of the full-MLA bucket sizes.
1212
-
1213
- For each group, bucket its layers by page_size_bytes and place each
1214
- layer at tuple_idx = position-within-bucket. Emit one KVCacheTensor
1215
- per (tuple_idx, bucket) whose shared_by is the union of per-group
1216
- layers at that slot.
1228
+ Emit one KVCacheTensor per (slot_idx, page_size). Layers from different
1229
+ groups at the same slot share a tensor (they have independent block
1230
+ tables so block-id namespaces never collide).
1231
"""
1218
- full_mla_spec = kv_cache_groups[0].kv_cache_spec
1219
- assert isinstance(full_mla_spec, UniformTypeKVCacheSpecs)
1220
- page_sizes = sorted(full_mla_spec.get_page_sizes())
1221
- layer_tuple_page_bytes = sum(page_sizes)
1232
+ # buckets = {page_size: [[layer_names], [layer_names], ...]}
1233
+ buckets = _bucket_layers_by_page_size(kv_cache_groups)
1234
+ total_num_bytes_per_block = sum(ps * len(slots) for ps, slots in buckets.items())
1235
1223
- # Pre-bucket each group's layers by page_size (registration order within
1224
- # bucket). bucketed[g_idx][page_size] = [layer_name, ...].
1225
- bucketed: list[dict[int, list[str]]] = []
1226
- for group in kv_cache_groups:
1227
- assert isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
1228
- specs = group.kv_cache_spec.kv_cache_specs
1229
- b: dict[int, list[str]] = defaultdict(list)
1230
- for name in group.layer_names:
1231
- b[specs[name].page_size_bytes].append(name)
1232
- bucketed.append(b)
1233
-
1234
- # num_layer_tuples = longest bucket list across all groups. For the
1235
- # full-MLA group this equals the count of layers in the largest
1236
- # per-page-size bucket (= get_num_layer_tuples()); for SWA sub-groups
1237
- # this equals the sub-group size (each has a single page_size).
1238
- num_layer_tuples = max(len(layers) for b in bucketed for layers in b.values())
1239
-
1240
- num_blocks = available_memory // (layer_tuple_page_bytes * num_layer_tuples)
1236
+ num_blocks = available_memory // total_num_bytes_per_block
1237
num_blocks = may_override_num_blocks(vllm_config, num_blocks)
1238
1239
kv_cache_tensors: list[KVCacheTensor] = []
1244
- for tuple_idx in range(num_layer_tuples):
1245
- for ps in page_sizes:
1246
- shared_by: list[str] = []
1247
- for b in bucketed:
1248
- bucket = b.get(ps)
1249
- if bucket is not None and tuple_idx < len(bucket):
1250
- shared_by.append(bucket[tuple_idx])
1251
- kv_cache_tensors.append(
1252
- KVCacheTensor(size=ps * num_blocks, shared_by=shared_by)
1253
- )
1240
+ for ps, slots in buckets.items():
1241
+ for slot in slots:
1242
+ kv_cache_tensors.append(KVCacheTensor(size=ps * num_blocks, shared_by=slot))
1243
1244
return num_blocks, kv_cache_tensors
1245