vllm.v1.core.kv_cache_planning
¶
KV cache planning: grouping, sizing, and memory-fit logic.
Functions:
-
check_enough_kv_cache_memory–Checks whether
available_memoryis enough for the KV cache to hold at -
create_kv_cache_group_specs–Create KVCacheGroupSpec object for each kv cache group layer.
-
estimate_max_model_len–Estimates the maximum model length that can fit in the available memory
-
get_kv_cache_config_from_groups–Generate the KV cache configuration from the KV cache groups and spec
-
get_kv_cache_configs–Generates the KV cache configurations for a model.
-
get_kv_cache_groups–Split the layers in the model into groups with the same KV cache spec.
-
get_uniform_page_size–Get the page size of the KV cache.
-
is_kv_cache_spec_uniform–Whether all layers in the given KVCacheSpec have the same KV cache spec.
-
kv_cache_groups_tp_replicas–Consecutive TP ranks holding identical KV for every layer.
-
max_memory_usage_bytes–Get the maximum memory usage in bytes for the given KV cache specs.
-
may_override_num_blocks–Override the number of kv cache blocks if
num_gpu_blocks_overrideis set. -
unify_hybrid_kv_cache_specs–This function tries to convert the KV cache specs to one type if the model
-
unify_kv_cache_spec_page_size–Unify the page size of the given KVCacheSpec. If the page size of all layers
-
validate_kv_cache_layout–Validate that the resolved layout can express this model's packing.
_annotate_eagle_groups(vllm_config, kv_cache_spec, kv_cache_groups, use_trailing_layer_fallback=False)
¶
Flag the KV cache groups that hold drafter attention layers.
Two detection rules, in order of preference:
- Spec-driven.
non_causal_multi_token_decodeis declared on MLAAttentionSpec and set by drafter attention layers that run a non-causal multi-token decode (today only Kimi-K3 DSpark). It survives MLAAttentionSpec.merge, so it still identifies a group after per-group spec merging, wherever grouping happens to land. It is sufficient but not necessary: a drafter whose spec is indistinguishable from the target's cannot be found this way. - Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose
MTP block reuses the target's own decoder layer and so carries no spec
marker. Their draft attention layers always register after every
target layer, so flag whichever group holds the last registered layer.
This rule is only valid where the groups partition exactly the layers
of
kv_cache_spec, which is re-checked here before applying it. When a drafter's trailing caches span several groups, this rule flags only the group holding the very last layer and must be generalized. FIXME(yifan): avoid/generalize this hacky check.
Parameters:
-
(vllm_config¶VllmConfig) –Config supplying the speculative method, if any.
-
(kv_cache_spec¶dict[str, KVCacheSpec]) –The kv cache spec of each attention layer, in layer registration order. Only read by rule 2.
-
(kv_cache_groups¶list[KVCacheGroupSpec]) –Groups to annotate in place.
-
(use_trailing_layer_fallback¶bool, default:False) –Enable rule 2. Callers gate this on
_uses_trailing_mtp_layers.
Source code in vllm/v1/core/kv_cache_planning.py
_approximate_gcd(values, *, lower_bound=None)
¶
Pick a chunk size that minimizes total upward padding.
Each x is rounded up to a multiple of d:
x -> ceil(x / d) * d
Total padding is:
pad(d) = sum_i (ceil(x_i / d) * d - x_i)
We brute-force d in [lower_bound, max(values)] (fine for small lists / small maxima) and return the d with minimum padding. Ties prefer larger d.
Source code in vllm/v1/core/kv_cache_planning.py
_auto_fit_max_model_len(vllm_config, projected_groups_per_worker, available_memory)
¶
When max_model_len is set to -1, this function estimates the largest context length that can be supported with the available GPU memory. It uses binary search to find the maximum length that fits across all workers.
Parameters:
-
(vllm_config¶VllmConfig) –The global VllmConfig (will be modified in-place)
-
(projected_groups_per_worker¶list[list[KVCacheGroupSpec]]) –KV cache groups projected to each worker.
-
(available_memory¶list[int]) –Memory available for KV cache in bytes for each worker.
Source code in vllm/v1/core/kv_cache_planning.py
1364 1365 1366 1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419 1420 1421 1422 1423 1424 1425 1426 1427 1428 1429 1430 1431 1432 1433 | |
_ensure_min_page_size(groups, common_page, hidden_specs)
¶
Scale up group block sizes so the common page is at least as large as the biggest hidden-state per-token cost.
This protects the hidden state extraction feature, where we store hidden states in the KV cache, and return them to clients using a custom hidden state connector.
Returns the (possibly scaled) groups and updated common page size.
Source code in vllm/v1/core/kv_cache_planning.py
_estimate_max_model_len_from_groups(vllm_config, kv_cache_groups, available_memory)
¶
Binary search for the maximum model length that fits in available memory. Returns 0 if even 1 token doesn't fit.
Source code in vllm/v1/core/kv_cache_planning.py
_get_kv_cache_bytes_per_block(kv_cache_groups, layout)
¶
Return the largest cache group's bytes per block.
Source code in vllm/v1/core/kv_cache_planning.py
_get_kv_cache_groups_uniform_page_size(kv_cache_spec)
¶
Generates the KV cache groups for hybrid models with multiple attention types but still with a uniform page size (physical memory per block per layer) for all layers.
Detailed explanation about kv cache management of hybrid models:
The layers in the models are repeated with some patterns, e.g., a model
with 10 full attention layers and 20 sliding window attention layers can be
regarded as repeating the pattern (1 * full, 2 * sw) 10 times.
The KVCacheManager allocates different block tables for each of the 3 layers
in the pattern, and repeats each of them 10 times to generate the
block_table for the 30 layers in the model.
Therefore, we can group the layers in the model into 3 kv_cache_groups, each
of which contains 10 layers in the model.
The KVCacheManager allocates the block_table for each group based on its
kv_cache spec, and the model runner applies the block table to each layer
in the group.
For example:
1. A model only uses full attention. The pattern is
(num_hidden_layers * full), so there is only one group and the block table
is shared by all layers. It is already handled by
_get_kv_cache_config_uniform_type.
2. A model with 10 full attention layers and 20 sliding window
attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so
there are 3 kv_cache_groups, each of which represents 10 layers.
To simplify the implementation, we make the following assumptions:
1. Physical memory per block: Must be the same across all KV cache groups.
Breaking this assumption is non-trivial due to memory fragmentation concerns
when allocating blocks of different sizes.
2. Tokens per block (block_size): Currently, we directly use
CacheConfig.block_size for all layers. It can be extended to vary by KV
cache group, but within each KV cache group, all layers must share the same
block size.
3. Physical memory per token per layer: This property is decided by model
config. Currently we only support models that have the same physical memory
per token per layer for all layers. Can be relaxed with a simple extension,
but still need to keep physical memory per block the same for all groups.
4. Number of layers per group: Currently assumed the same for all layers.
Can be relaxed with a simple extension, but still need to keep physical
memory per block the same for all groups.
5. Attention type within groups: All layers in a group must share the same
attention type. One exception is that, when
--disable-hybrid-kv-cache-manager is true, the single group for full
attention layers may also include attention layers using sliding window or
LLaMA 4 local attention. See unify_hybrid_kv_cache_specs for more details.
6. Support for multiple attention types: The design for most components is
general to an arbitrary number of attention types. But
find_longest_cache_hit only supports one attention type or two
types of full-attention plus exactly one another type. The general
implementation of this function is feasible but we don't know how to
implement it cleanly yet.
As we assume tokens per block, physical memory per token per layer, and number of layers per group are the same now, we can ensure that physical memory per block is the same for all groups.
Parameters:
-
(kv_cache_spec¶dict[str, KVCacheSpec]) –The KVCacheSpec of each attention layer in the model
Returns: The generated KVCacheGroupSpecs
Source code in vllm/v1/core/kv_cache_planning.py
373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 | |
_get_kv_cache_groups_uniform_spec(kv_cache_specs)
¶
Generates the KV cache configuration for a model with the same KV cache spec for all layers.
Parameters:
-
(kv_cache_specs¶dict[str, KVCacheSpec]) –The kv cache spec of each attention layer in the model
Returns:
-
list[KVCacheGroupSpec]–The generated KVCacheGroupSpecs
Source code in vllm/v1/core/kv_cache_planning.py
_get_kv_cache_groups_uniform_type(spec)
¶
Generates the KV cache configuration for a model with one type of KV cache but different hidden sizes. All layers are merged into one group.
Parameters:
-
(spec¶UniformTypeKVCacheSpecs) –The UniformTypeKVCacheSpecs of the model
Returns:
-
list[KVCacheGroupSpec]–The generated KVCacheGroupSpecs
Source code in vllm/v1/core/kv_cache_planning.py
_get_layer_block_stride(spec)
¶
A layer-outer block stride: the page padded to its alignment.
_get_packed_kv_cache_groups(vllm_config, kv_cache_spec)
¶
Group mixed-page-size layers for contiguous block-outermost packing.
Greedily buckets layers into uniform-type specs. Buckets with equal layer
counts per page size are treated as a repeating layer pattern (one layer
per page size) and split into groups covering the same number of pattern
repeats (picked by _approximate_gcd to minimize padding), so all
groups pack into the same per-block layout. Mamba buckets are additionally
split to fit the block the attention buckets already need.
Returns None when the layout is not block-outermost or all layers already
share one page size.
Source code in vllm/v1/core/kv_cache_planning.py
871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 | |
_groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups)
¶
Whether the groups cover every layer of kv_cache_spec exactly once.
The trailing-layer draft fallback is only meaningful when the last registered layer is guaranteed to appear in exactly one group.
Source code in vllm/v1/core/kv_cache_planning.py
_max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups)
¶
Calculate maximum memory usage in bytes from KV cache groups.
This correctly accounts for padding in hybrid models. For example, if a model has 8 full attention layers and 9 sliding window layers, they will be padded to 9 full + 9 sliding window for uniform group sizes.
Each group independently claims blocks from the shared pool, so a request consumes
the sum of the per-group block counts, i.e. bytes_per_block * total_blocks.
Source code in vllm/v1/core/kv_cache_planning.py
_pool_bytes_per_block(vllm_config, kv_cache_groups)
¶
Bytes consumed by one block in the worker's shared KV cache pool, mirroring
the divisor used by get_kv_cache_config_from_groups to convert
available_memory into num_blocks. Used to compute the effective KV cache
capacity once num_gpu_blocks_override is applied.
Source code in vllm/v1/core/kv_cache_planning.py
_project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec)
¶
Projects global KV cache groups onto a single worker's assigned layers.
In pipeline parallelism, each worker only owns a subset of layers. This function filters the global groups to include only layers present on the given worker, adjusting UniformTypeKVCacheSpecs accordingly.
Parameters:
-
(global_kv_cache_groups¶list[KVCacheGroupSpec]) –The global KV cache groups for the whole model.
-
(worker_spec¶dict[str, KVCacheSpec]) –The KV cache spec of each layer on this worker.
Returns:
-
list[KVCacheGroupSpec]–The projected KV cache groups containing only this worker's layers.
Source code in vllm/v1/core/kv_cache_planning.py
_promote_local_kv_cache_specs(kv_cache_spec)
¶
Use full-attention allocation for local-attention cache specs.
The returned specs affect KV cache management only. Attention modules keep their original sliding-window or chunked-local compute behavior.
Source code in vllm/v1/core/kv_cache_planning.py
_try_get_full_allocation_fallback_groups(kv_cache_spec)
¶
Try a supported full-allocation fallback for local-attention layers.
Source code in vllm/v1/core/kv_cache_planning.py
_uses_trailing_mtp_layers(vllm_config)
¶
Whether the drafter's KV layers can be located positionally.
MTP drafters register their KV layers after the target layers, without a
distinguishing spec marker. DeepseekV4/V4.1 do this with dspark too.
The annotator separately checks that the groups partition the layers exactly.
Source code in vllm/v1/core/kv_cache_planning.py
_warn_if_unannotated_eagle_mamba(vllm_config, kv_cache_groups)
¶
Warn when no KV cache group could be identified as the draft model's.
Parameters:
-
(vllm_config¶VllmConfig) –Config supplying the speculative method, if any.
-
(kv_cache_groups¶list[KVCacheGroupSpec]) –Groups as they will be handed to consumers.
Source code in vllm/v1/core/kv_cache_planning.py
check_enough_kv_cache_memory(vllm_config, kv_cache_spec, available_memory)
¶
Checks whether available_memory is enough for the KV cache to hold at
least one request with the model's max_model_len.
Parameters:
-
(vllm_config¶VllmConfig) –The global VllmConfig
-
(kv_cache_spec¶dict[str, KVCacheSpec]) –The kv cache spec of each attention layer in the model
-
(available_memory¶int) –Memory available for KV cache in bytes.
Raises:
-
ValueError–If there is not enough memory available for the KV cache.
Source code in vllm/v1/core/kv_cache_planning.py
create_kv_cache_group_specs(kv_cache_spec, grouped_layer_names)
¶
Create KVCacheGroupSpec object for each kv cache group layer. The layers in the same group should share the same KVCacheSpec.
Parameters:
-
(kv_cache_spec¶dict[str, KVCacheSpec]) –A mapping from each layer name to its corresponding KVCacheSpec.
-
(grouped_layer_names¶list[list[str]]) –A list of kv cache groups, where each element is a list of layer names that belong to the same group and should share the same KVCacheSpec.
Returns:
-
list[KVCacheGroupSpec]–A list of KVCacheGroupSpec objects, one for each group.
Source code in vllm/v1/core/kv_cache_planning.py
estimate_max_model_len(vllm_config, kv_cache_spec, available_memory)
¶
Estimates the maximum model length that can fit in the available memory using binary search.
This function temporarily modifies max_model_len during estimation but restores the original value before returning, ensuring no side effects.
Parameters:
-
(vllm_config¶VllmConfig) –The global VllmConfig
-
(kv_cache_spec¶dict[str, KVCacheSpec]) –The kv cache spec of each attention layer in the model
-
(available_memory¶int) –Memory available for KV cache in bytes.
Returns:
-
int–The estimated maximum model length that can fit in the available memory.
Source code in vllm/v1/core/kv_cache_planning.py
get_kv_cache_config_from_groups(vllm_config, kv_cache_groups, available_memory)
¶
Generate the KV cache configuration from the KV cache groups and spec of each layer.
Parameters:
-
(vllm_config¶VllmConfig) –The global VllmConfig
-
(kv_cache_groups¶list[KVCacheGroupSpec]) –The KV cache groups
-
(available_memory¶int) –Memory available for KV cache in bytes
Returns: The generated KVCacheConfig
Source code in vllm/v1/core/kv_cache_planning.py
599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 | |
get_kv_cache_configs(vllm_config, kv_cache_specs, available_memory)
¶
Generates the KV cache configurations for a model.
Since we use a shared centralized controller for all workers, we need the
kv_cache_config to be consistent across all workers to make sure
the KV cache allocation can be applied to all workers. However, different
workers may have different memory available, and different type of layers
(when pipeline parallel is enabled). To handle the difference between
workers, the current implementation is:
1. Merge the KV cache specs of all workers to get the KVCacheSpecs for
the whole model.
2. Generate the KV cache groups based on the layer ratio of the whole model.
This also handles spec unification for hybrid models.
3. Handle auto-fit max_model_len and memory checks using per-worker
projected groups to account for PP sharding.
4. Generate the KV cache configs for each worker based on the KV cache
grouping strategy. (This is reasonable because the layer ratio of
different PP stages are similar.)
5. Change the num_blocks of each worker to the smallest among all workers
and shrink tensor sizes proportionally to avoid allocating unused memory.
Parameters:
-
(vllm_config¶VllmConfig) –The global VllmConfig
-
(kv_cache_specs¶list[dict[str, KVCacheSpec]]) –List of dict[layer_name, KVCacheSpec] for each worker.
-
(available_memory¶list[int]) –Memory available for KV cache in bytes for each worker.
Returns:
-
list[KVCacheConfig]–The generated KVCacheConfigs for each worker.
Source code in vllm/v1/core/kv_cache_planning.py
1481 1482 1483 1484 1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526 1527 1528 1529 1530 1531 1532 1533 1534 1535 1536 1537 1538 1539 1540 1541 1542 1543 1544 1545 1546 1547 1548 1549 1550 1551 1552 1553 1554 1555 1556 1557 1558 1559 1560 1561 1562 1563 1564 1565 1566 1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585 1586 1587 1588 1589 1590 1591 1592 1593 1594 1595 1596 1597 1598 1599 1600 1601 1602 1603 1604 1605 1606 1607 1608 1609 1610 1611 1612 1613 1614 1615 1616 1617 1618 1619 1620 1621 1622 1623 1624 1625 1626 1627 1628 1629 1630 1631 1632 1633 1634 1635 1636 1637 | |
get_kv_cache_groups(vllm_config, kv_cache_spec)
¶
Split the layers in the model into groups with the same KV cache spec.
Parameters:
-
(vllm_config¶VllmConfig) –The global VllmConfig
-
(kv_cache_spec¶dict[str, KVCacheSpec]) –The kv cache spec of each attention layer in the model
Returns:
-
list[KVCacheGroupSpec]–The generated KVCacheGroups
Source code in vllm/v1/core/kv_cache_planning.py
1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 1261 1262 1263 1264 | |
get_uniform_page_size(kv_cache_specs)
¶
Get the page size of the KV cache.
is_kv_cache_spec_uniform(kv_cache_spec)
¶
Whether all layers in the given KVCacheSpec have the same KV cache spec. Note that we regard FullAttentionSpec with and without sliding window as the same type.
Parameters:
-
(kv_cache_spec¶dict[str, KVCacheSpec]) –The kv cache spec of each attention layer in the model
Returns:
-
bool–True if all layers have the same type, False otherwise.
Source code in vllm/v1/core/kv_cache_planning.py
kv_cache_groups_tp_replicas(groups, tp_size, dcp_size=1)
¶
Consecutive TP ranks holding identical KV for every layer.
Source code in vllm/v1/core/kv_cache_planning.py
max_memory_usage_bytes(vllm_config, kv_cache_specs)
¶
Get the maximum memory usage in bytes for the given KV cache specs.
Source code in vllm/v1/core/kv_cache_planning.py
may_override_num_blocks(vllm_config, num_blocks)
¶
Override the number of kv cache blocks if num_gpu_blocks_override is set.
The override is logged once, at the call site in get_kv_cache_configs.
Source code in vllm/v1/core/kv_cache_planning.py
unify_hybrid_kv_cache_specs(kv_cache_spec)
¶
This function tries to convert the KV cache specs to one type if the model is a hybrid model with multiple type of KV cache. It will convert all SlidingWindowSpec to FullAttentionSpec if both types are present.
Parameters:
-
(kv_cache_spec¶dict[str, KVCacheSpec]) –The kv cache spec of each attention layer in the model
Source code in vllm/v1/core/kv_cache_planning.py
unify_kv_cache_spec_page_size(kv_cache_spec)
¶
Unify the page size of the given KVCacheSpec. If the page size of all layers
are the same, return the original KVCacheSpec. If not same, unify the page
size by increasing the block size of layers with smaller page size. Two
cases cannot be unified by block size alone and pad their physical page to
the maximum instead: Mamba layers, whose page size comes from state shapes
and is independent of block size; and non-MLA attention layers whose page
does not evenly divide the maximum (the padded page is read through a
strided view). MLA is excluded because sparse MLA indexes the cache in
whole token rows (see flat_kv_row_view), so its block stride can only
be padded by its own row-aligned alignment, not to an arbitrary page
size. Raise NotImplementedError if failed to unify the page size;
get_kv_cache_groups catches it to try the full-allocation fallback
(e.g. MLA next to an incompatible sliding-window draft).
Parameters:
-
(kv_cache_spec¶dict[str, KVCacheSpec]) –The KVCacheSpec of each attention layer in the model
Returns:
-
dict[str, KVCacheSpec]–The updated KVCacheSpec with the same page_size_bytes.
Source code in vllm/v1/core/kv_cache_planning.py
validate_kv_cache_layout(layout, kv_cache_groups)
¶
Validate that the resolved layout can express this model's packing.
The layout was chosen once in the engine core from the backends' supported sets; a backend whose model packs pages side by side (e.g. the DeepSeek-V4 indexer) declares block-outermost layouts there, so an inexpressible layout reaching this point is an error.