diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index d2702e7d4f..5145f5a272 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -4321,6 +4321,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4354,6 +4355,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4387,6 +4389,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4420,6 +4423,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4453,6 +4457,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index f6a4a89f11..96e5d4fc5e 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -371,6 +371,7 @@ steps: - vllm/tokenizers/ - vllm/v1/attention/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/v1/spec_decode/metrics.py - vllm/v1/worker/ - vllm/_aiter_ops.py @@ -516,6 +517,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -559,6 +561,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -602,6 +605,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -645,6 +649,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -688,6 +693,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index b2156af346..81019421ed 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -394,7 +394,7 @@ def test_k3_dspark_mla_kv_cache_spec_groups_with_target_mla(): """The draft's MLA layers must share a KV cache group with the target's.""" from vllm.model_executor.layers.attention.mla_attention import MLAAttention from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention - from vllm.v1.core.kv_cache_utils import _get_kv_cache_groups_uniform_page_size + from vllm.v1.core.kv_cache_planning import _get_kv_cache_groups_uniform_page_size vllm_config = SimpleNamespace( model_config=None, cache_config=SimpleNamespace(block_size=64) diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index 073f616af8..ee6dfee7ec 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -14,10 +14,8 @@ from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.mem_constants import GiB_bytes from vllm.v1.attention.backends.utils import resolve_kv_cache_layout -from vllm.v1.core.kv_cache_utils import ( - generate_scheduler_kv_cache_config, - get_kv_cache_configs, -) +from vllm.v1.core.kv_cache_planning import get_kv_cache_configs +from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config from vllm.v1.engine.core import EngineCore as V1EngineCore from ..utils import ( diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py index d869141df8..45ba1ac0dc 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -23,7 +23,7 @@ from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, ) from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts -from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups +from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs from vllm.v1.kv_cache_layout import KVCacheLayout from vllm.v1.worker.utils import allocate_kv_cache, customize_attention_spec diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 6a705b6e5e..64cd550540 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -18,13 +18,15 @@ import torch from vllm.config import CacheConfig from vllm.v1.core.kv_cache_manager import KVCacheManager -from vllm.v1.core.kv_cache_utils import ( +from vllm.v1.core.kv_cache_planning import ( _get_kv_cache_bytes_per_block, _get_packed_kv_cache_groups, _pool_bytes_per_block, - generate_scheduler_kv_cache_config, get_kv_cache_config_from_groups, get_kv_cache_groups, +) +from vllm.v1.core.kv_cache_utils import ( + generate_scheduler_kv_cache_config, resolve_kv_cache_block_sizes, ) from vllm.v1.hisparse.layout import _build_hisparse_kv_cache_tensors diff --git a/tests/v1/core/test_kv_cache_planning.py b/tests/v1/core/test_kv_cache_planning.py new file mode 100644 index 0000000000..6292868081 --- /dev/null +++ b/tests/v1/core/test_kv_cache_planning.py @@ -0,0 +1,2284 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import mmap +import subprocess +import sys +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +import vllm.v1.core.kv_cache_planning as kv_cache_planning +import vllm.v1.core.kv_cache_utils as kv_cache_utils +import vllm.v1.hisparse.runtime as hisparse_runtime_module +from tests.v1.core.utils import ( + new_kv_cache_spec, + new_mamba_spec, + new_mla_spec, + new_sliding_window_spec, +) +from vllm.config import ( + CacheConfig, + ModelConfig, + SchedulerConfig, + VllmConfig, +) +from vllm.config.attention import HiSparseConfig +from vllm.utils.mem_constants import GiB_bytes +from vllm.v1.core.kv_cache_planning import ( + check_enough_kv_cache_memory, + estimate_max_model_len, + get_kv_cache_configs, + get_kv_cache_groups, + is_kv_cache_spec_uniform, + kv_cache_groups_tp_replicas, +) +from vllm.v1.core.kv_cache_utils import ( + generate_scheduler_kv_cache_config, +) +from vllm.v1.hisparse.layout import ( + get_hisparse_kv_cache_groups, +) +from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, + CircularBufferSpec, + FullAttentionSpec, + HiddenStateCacheSpec, + HiSparseHotSpec, + HiSparseResidentSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, + SparseCacheRole, + UniformTypeKVCacheSpecs, + iter_layer_specs, +) +from vllm.v1.kv_cache_layout import KVCacheLayout + +pytestmark = pytest.mark.cpu_test + + +@pytest.mark.parametrize("gpu_block_size", [32, 64]) +@pytest.mark.parametrize("shared_host_pool", [False, True]) +def test_hisparse_hma_uses_resolved_gpu_block_size( + monkeypatch, gpu_block_size, shared_host_pool +): + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) + specs = { + "model.layers.0.self_attn": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + is_index_group_leader=True, + ), + "model.layers.0.self_attn.indexer": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ), + } + config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), + model_config=SimpleNamespace( + hf_config=SimpleNamespace(index_topk=128), + max_model_len=gpu_block_size, + ), + parallel_config=SimpleNamespace( + tensor_parallel_size=2 if shared_host_pool else 1, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + decode_context_parallel_size=1, + world_size=2 if shared_host_pool else 1, + distributed_executor_backend="mp", + nnodes_within_dp=1, + ), + cache_config=SimpleNamespace( + num_gpu_blocks_override=7, + prefix_cache_retention_interval=None, + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC, + ), + ) + monkeypatch.setattr( + kv_cache_planning, "get_hisparse_host_pool_bytes", lambda _: 2**30 + ) + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( + config, get_hisparse_kv_cache_groups(config, specs), available_memory=2**30 + ) + assert cache_config.num_blocks == 7 + assert cache_config.hisparse_host_num_blocks is not None + assert cache_config.hisparse_host_num_blocks > 7 + + host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups + assert cache_config.hisparse_shared_host_pool is shared_host_pool + host_page = host_group.kv_cache_spec.page_size_bytes + alignment = mmap.PAGESIZE if shared_host_pool else 1 + expected_host_stride = (host_page + alignment - 1) // alignment * alignment + assert cache_config.hisparse_host_block_stride == expected_host_stride + assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride + assert host_group.host_resident + assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) + host_layers = set(host_group.layer_names) + for tensor in cache_config.kv_cache_tensors: + assert all( + (name in host_layers) == tensor.host_resident for name in tensor.layers + ) + host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident) + host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]] + assert host_tensor.block_stride == host_spec.page_size_bytes + assert host_tensor.layer_stride == ( + host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks + ) + assert host_group.kv_cache_spec.block_size == gpu_block_size + assert indexer_group.kv_cache_spec.block_size == gpu_block_size + host_specs = host_group.kv_cache_spec.kv_cache_specs + gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs + assert set(host_specs) == {"model.layers.0.self_attn"} + assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"} + assert indexer_group.kv_cache_spec.prefix_cacheable + assert host_group.enable_kv_transfer + auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups] + assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs) + assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs) + assert all( + spec.block_size == gpu_block_size + for spec in auxiliary_specs + if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)) + ) + scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes( + cache_config, + SimpleNamespace( + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + prefix_match_unit=None, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + kv_transfer_config=object(), + ), + ) + assert scheduler_block_size == hash_block_size == gpu_block_size + + +@pytest.mark.parametrize( + "max_model_len,num_gpu_blocks,ok", + [ + (8192, 160, False), + (8192, 1000, True), + (32768, 656, False), + (32768, 657, True), + (-1, 657, True), + ], +) +def test_hisparse_pool_must_fit_max_model_len( + monkeypatch, max_model_len, num_gpu_blocks, ok +): + """Each resident group takes its own blocks from HiSparse's shared GPU pool, + so a pool that fits only the indexer pages must be rejected at startup. + Resident pages are capped at the in-flight window: at 32768 tokens with 2048 + in flight, one request needs 512 indexer + 4 x 32 resident + 16 hot blocks, + plus the null block. Auto-fit (-1) must pick a length that passes the same + check.""" + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) + monkeypatch.setattr( + kv_cache_planning, "get_hisparse_host_pool_bytes", lambda _: 2**30 + ) + specs: dict[str, KVCacheSpec] = {} + for i in range(4): + specs[f"model.layers.{i}.self_attn"] = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + is_index_group_leader=True, + ) + specs[f"model.layers.{i}.self_attn.indexer"] = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ) + config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + config.scheduler_config.max_num_batched_tokens = 2048 + # Async scheduling doubles the in-flight window, and with it the bound. + config.scheduler_config.async_scheduling = False + config.attention_config.hisparse_config = HiSparseConfig() + config.model_config.hf_config.index_topk = 128 + config.cache_config.num_gpu_blocks_override = num_gpu_blocks + config.cache_config.kv_cache_layout = "BLHNC" + # Workers report the resident/hot caches HiSparse derives. + specs |= { + name: group.kv_cache_spec + for group in get_hisparse_kv_cache_groups(config, specs) or [] + if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec)) + for name in group.layer_names + } + + if ok: + kv_cache_planning.get_kv_cache_configs(config, [specs], [2**34]) + else: + with pytest.raises(ValueError, match="max seq len"): + kv_cache_planning.get_kv_cache_configs(config, [specs], [2**34]) + + +def new_chunked_local_attention_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + attention_chunk_size=4, +): + return ChunkedLocalAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + attention_chunk_size=attention_chunk_size, + ) + + +def test_get_kv_cache_configs_multiple_workers(): + model_config = ModelConfig(max_model_len=16) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + ref_kv_cache_spec = new_kv_cache_spec() + same_kv_cache_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + ] + + # Basic case. All things are the same. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + same_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + expected = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + ], + ) + assert kv_cache_configs == [expected, expected] + + # Different available memory. This is the case for TP. + # Use the smallest memory available. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + same_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 20, + ], + ) + assert kv_cache_configs == [expected, expected] + + # Different KV cache specs. This is the case for PP. + different_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + }, + { + "layer2": new_kv_cache_spec(), + "layer3": new_kv_cache_spec(), + }, + ] + + # Different workers have different layers. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1"], new_kv_cache_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer2", "layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer2", "layer3"], new_kv_cache_spec()), + ], + ), + ] + + # Some layers are the same, some are different. This is the case for TP+PP + tp_pp_kv_cache_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer3": new_kv_cache_spec(), + }, + { + "layer3": new_kv_cache_spec(), + }, + ] + + kv_cache_configs = get_kv_cache_configs( + vllm_config, + tp_pp_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + expected_12 = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + ], + ) + expected_3 = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), + ], + ) + assert kv_cache_configs == [ + expected_12, + expected_12, + expected_3, + expected_3, + ] + + # Different workers have different types of layers. This is the case for + # hybrid models + PP. + different_type_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer3": new_sliding_window_spec(), + "layer4": new_sliding_window_spec(), + }, + ] + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_type_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + KVCacheGroupSpec([], new_sliding_window_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer3", "layer4"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec([], ref_kv_cache_spec), + KVCacheGroupSpec(["layer3", "layer4"], new_sliding_window_spec()), + ], + ), + ] + + # When divided into multiple KVCacheGroups, need to ensure the number of + # layers per group is similar. + different_type_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_sliding_window_spec(), + "layer3": new_sliding_window_spec(), + }, + { + "layer4": new_kv_cache_spec(), + "layer5": new_sliding_window_spec(), + "layer6": new_sliding_window_spec(), + }, + ] + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_type_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 10, + ref_kv_cache_spec.page_size_bytes * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1"], ref_kv_cache_spec), + KVCacheGroupSpec(["layer2"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer3"], new_sliding_window_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer4"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer5"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer6"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer4"], ref_kv_cache_spec), + KVCacheGroupSpec(["layer5"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer6"], new_sliding_window_spec()), + ], + ), + ] + + # Have conflicting layers. Need to raise an error. + conflicting_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + }, + { + "layer1": new_sliding_window_spec(), + }, + ] + with pytest.raises(AssertionError): + get_kv_cache_configs( + vllm_config, + conflicting_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + + +@pytest.mark.parametrize( + "asymmetric_memory", + [False, True], + ids=["symmetric", "asymmetric"], +) +def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): + model_config = ModelConfig(max_model_len=512) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + ref_kv_cache_spec = new_kv_cache_spec() + pp_kv_cache_specs = [ + {"layer1": ref_kv_cache_spec}, + {"layer2": ref_kv_cache_spec}, + ] + + expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1 + avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks + + # With per-worker validation, each worker only needs memory for its own + # layers. Worker 2 having more memory shouldn't affect worker 1's config. + available_memory = ( + [avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2 + ) + + kv_cache_configs = get_kv_cache_configs( + vllm_config, + pp_kv_cache_specs, + available_memory, + ) + + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=expected_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes + * expected_num_blocks, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)], + ), + KVCacheConfig( + num_blocks=expected_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, + layers=["layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes + * expected_num_blocks, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)], + ), + ] + + +def test_project_kv_cache_groups_to_worker(): + spec_a = new_kv_cache_spec() + spec_b = new_kv_cache_spec(num_kv_heads=4) + + global_groups = [ + KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a), + ] + worker_spec = {"layer1": spec_a, "layer2": spec_a} + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups, worker_spec + ) + assert len(projected) == 1 + assert projected[0].layer_names == ["layer1", "layer2"] + assert projected[0].kv_cache_spec is spec_a + + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups, {"layer4": spec_a} + ) + assert len(projected) == 1 + assert projected[0].layer_names == [] + assert projected[0].kv_cache_spec is spec_a + + uniform_spec = UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a}, + ) + global_groups_uniform = [ + KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec), + ] + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups_uniform, {"layer1": spec_a, "layer3": spec_a} + ) + assert len(projected) == 1 + assert projected[0].layer_names == ["layer1", "layer3"] + proj_spec = projected[0].kv_cache_spec + assert isinstance(proj_spec, UniformTypeKVCacheSpecs) + assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"} + + +@pytest.mark.parametrize("with_draft", [False, True]) +@pytest.mark.parametrize("indexer_alignment", [None, 512]) +def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment): + from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view + from vllm.v1.worker.utils import allocate_kv_cache + + config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) + config.cache_config.kv_cache_layout = "BLHNC" + config.parallel_config.decode_context_parallel_size = 4 + common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8) + specs = { + "target": MLAAttentionSpec( + **common, + head_size=576, + state_content_bytes=656, + block_stride_alignment=656, + cache_dtype_str="fp8_ds_mla", + is_index_group_leader=True, + ), + "indexer": MLAAttentionSpec( + **common, + head_size=132, + cache_role=SparseCacheRole.INDEXER, + block_stride_alignment=indexer_alignment, + ), + "draft": SlidingWindowSpec( + block_size=64, + num_kv_heads=64, + head_size=64, + dtype=torch.bfloat16, + sliding_window=2048, + dcp_sharded=False, + ), + } + if not with_draft: + del specs["draft"] + groups = get_kv_cache_groups(config, specs) + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( + config, groups, available_memory=8 * 1024 * 1024 + ) + caches = allocate_kv_cache( + cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups) + ) + cache = caches["target"].squeeze(1) + if indexer_alignment is not None: + assert caches["indexer"].stride(0) % indexer_alignment == 0 + rows, stride_rows = flat_kv_row_view(cache, 64) + cache[1, 0].fill_(7) + torch.testing.assert_close(rows[stride_rows], cache[1, 0]) + assert torch.all(rows[stride_rows] == 7) + + +def test_is_kv_cache_spec_uniform(): + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_kv_cache_spec(num_kv_heads=32), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_kv_cache_spec(num_kv_heads=32, sliding_window=1), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + } + assert not is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=2), + } + assert not is_kv_cache_spec_uniform(kv_cache_spec) + + script = """ +import sys + +from vllm.v1.core.kv_cache_planning import is_kv_cache_spec_uniform +from vllm.v1.kv_cache_interface import KVCacheSpec + +if sys.flags.optimize < 1: + raise RuntimeError("subprocess is not running with optimization enabled") +specs = { + "a": KVCacheSpec(block_size=1), + "b": KVCacheSpec(block_size=2), +} +if is_kv_cache_spec_uniform(specs): + raise RuntimeError("different specs were treated as uniform") +""" + subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=Path(__file__).resolve().parents[3], + check=True, + timeout=60, + ) + + +@pytest.mark.parametrize( + ("model_id", "max_model_len", "want_estimated_max_len"), + [ + ("Qwen/Qwen1.5-7B", 16385, 16384), + ("Qwen/Qwen1.5-7B", 16383, 16383), + ], +) +def test_estimate_max_model_len(model_id, max_model_len, want_estimated_max_len): + # Create a VllmConfig + model_config = ModelConfig( + model_id, + runner="generate", + dtype="float16", + max_model_len=max_model_len, + ) + scheduler_config = SchedulerConfig( + max_num_batched_tokens=32768, + max_model_len=model_config.max_model_len, + is_encoder_decoder=model_config.is_encoder_decoder, + ) + + vllm_config = VllmConfig( + model_config=model_config, + scheduler_config=scheduler_config, + ) + + # Create KV cache specs + kv_cache_spec = {} + for i in range(32): + layer_name = f"layer_{i}" + kv_cache_spec[layer_name] = FullAttentionSpec( + block_size=16, + num_kv_heads=32, + head_size=128, + dtype=torch.float16, + ) + # Estimate the maximum model length, 16384 model_len need 8GB + estimated_max_len = estimate_max_model_len( + vllm_config, kv_cache_spec, 8 * GiB_bytes + ) + assert estimated_max_len == want_estimated_max_len + + +def test_get_kv_cache_config_one_worker(): + # pass max_model_len to pass check_enough_kv_cache_memory + model_config = ModelConfig(max_model_len=16) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + # all layers are full attention -> single group + kv_cache_specs_full = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + kv_cache_config_full = get_kv_cache_configs( + vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] + )[0] + print(kv_cache_config_full) + assert kv_cache_config_full == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], + ) + + # all layers are sliding window -> single group + kv_cache_specs_sliding = { + "layer_1": new_sliding_window_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_sliding = get_kv_cache_configs( + vllm_config, [kv_cache_specs_sliding], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_sliding == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec()) + ], + ) + + # full + sliding, but disable_hybrid_kv_cache_manager + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2"], new_kv_cache_spec(sliding_window=1) + ), + ], + ) + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False + + # full + sliding, with hybrid_kv_cache_manager + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=64, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 64, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 64, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 64, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 64, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1"], new_kv_cache_spec()), + KVCacheGroupSpec(["layer_2"], new_sliding_window_spec()), + ], + ) + + # 2 full + 4 sliding, 2 layers per group + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_sliding_window_spec(), + "layer_4": new_sliding_window_spec(), + "layer_5": new_sliding_window_spec(), + "layer_6": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_3", "layer_5"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_4", "layer_6"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()), + KVCacheGroupSpec(["layer_3", "layer_5"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer_4", "layer_6"], new_sliding_window_spec()), + ], + ) + + # 3 full + 7 sliding, pad to 3 full + 9 sliding + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_kv_cache_spec(), + "layer_4": new_sliding_window_spec(), + "layer_5": new_sliding_window_spec(), + "layer_6": new_sliding_window_spec(), + "layer_7": new_sliding_window_spec(), + "layer_8": new_sliding_window_spec(), + "layer_9": new_sliding_window_spec(), + "layer_10": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_1", "layer_2", "layer_3"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_4", "layer_7", "layer_10"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_5", "layer_8"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_6", "layer_9"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"], new_kv_cache_spec()), + KVCacheGroupSpec( + ["layer_4", "layer_7", "layer_10"], new_sliding_window_spec() + ), + KVCacheGroupSpec(["layer_5", "layer_8"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer_6", "layer_9"], new_sliding_window_spec()), + ], + ) + + # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss + # eagle where there is only one more full attention layer than sliding window layers + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_kv_cache_spec(), + "layer_4": new_kv_cache_spec(), + "layer_5": new_kv_cache_spec(), + "layer_6": new_kv_cache_spec(), + "layer_7": new_sliding_window_spec(), + "layer_8": new_sliding_window_spec(), + "layer_9": new_sliding_window_spec(), + "layer_10": new_sliding_window_spec(), + "layer_11": new_sliding_window_spec(), + } + + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32] + )[0] + print(kv_cache_config_hybrid) + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 6, + layers=[ + "layer_1", + "layer_2", + "layer_3", + "layer_4", + "layer_5", + "layer_6", + ], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 6, + layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6"], + new_kv_cache_spec(), + ), + KVCacheGroupSpec( + ["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], + new_sliding_window_spec(), + ), + ], + ) + + # different hidden size but same type, use UniformTypeKVCacheSpecs + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=128), + "layer_2": new_kv_cache_spec(head_size=64), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] + )[0] + # Layers of different page sizes pack densely into one allocation: the + # 2x-larger layer_1 takes the first region, layer_2 the next. + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 2 * 32, + block_stride=mem_per_block_per_layer * 2, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + offset=mem_per_block_per_layer * 2 * 32, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2"], + UniformTypeKVCacheSpecs( + block_size=16, kv_cache_specs=kv_cache_specs_hybrid + ), + ) + ], + ) + + # Different hidden size and different type, align by different block size + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=64), + "layer_2": new_sliding_window_spec(head_size=32), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)), + KVCacheGroupSpec( + ["layer_2"], new_sliding_window_spec(head_size=32, block_size=32) + ), + ], + ) + + # different hidden size that cannot be aligned by using different block size, + # but can be aligned by padding the smaller physical page. + swa_spec = new_sliding_window_spec(head_size=96) + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=64), + "layer_2": swa_spec, + } + + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + padded_page_size = swa_spec.page_size_bytes + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=42, + kv_cache_tensors=[ + KVCacheTensor( + size=padded_page_size * 42, + layers=["layer_1"], + layer_stride=padded_page_size * 42, + block_stride=padded_page_size, + ), + KVCacheTensor( + size=padded_page_size * 42, + layers=["layer_2"], + layer_stride=padded_page_size * 42, + block_stride=padded_page_size, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1"], + new_kv_cache_spec( + head_size=64, + page_size_padded=padded_page_size, + ), + ), + KVCacheGroupSpec( + ["layer_2"], + new_sliding_window_spec(head_size=96), + ), + ], + ) + + # Test num_gpu_blocks_override + vllm_config.cache_config.num_gpu_blocks_override = 16 + kv_cache_config_override_blocks = get_kv_cache_configs( + vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_override_blocks == KVCacheConfig( + num_blocks=16, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 16 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 16, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], + ) + + +def test_get_kv_cache_configs_attention_free(): + kv_cache_specs: dict[str, KVCacheSpec] = {} + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) + vllm_config.cache_config.prefix_cache_retention_interval = None + kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0]) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[], + ) + ] + + +def test_packed_groups_glm5_like_hybrid(): + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=8192)) + vllm_config.cache_config.kv_cache_layout = "BLHNC" + kpool, block_size = 4, 1024 + kv_cache_spec: dict[str, KVCacheSpec] = {} + for i in range(11): + kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + block_stride_alignment=32 * 576 * 2, + ) + kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + tokens_per_state=kpool, + block_stride_alignment=64 * 132, + ) + kv_cache_spec[f"layers.{i}.tail"] = CircularBufferSpec( + block_size=kpool, + num_kv_heads=2, + head_size=128, + head_size_v=0, + dtype=torch.bfloat16, + ) + for i in range(34): + kv_cache_spec[f"layers.{i}.kda"] = MambaSpec( + shapes=((3, 6144), (16, 128, 128)), + dtypes=(torch.bfloat16, torch.float32), + block_size=block_size, + ) + + groups = kv_cache_planning.get_kv_cache_groups(vllm_config, kv_cache_spec) + attn = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, MLAAttentionSpec) + ] + rings = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, CircularBufferSpec) + ] + mamba = [ + g + for g in groups + if isinstance( + getattr(g.kv_cache_spec, "first_spec", g.kv_cache_spec), MambaSpec + ) + ] + assert len(attn) == 1 and len(attn[0].layer_names) == 22 + assert len(rings) == 1 and len(rings[0].layer_names) == 11 + # 11 * (1.18 MB MLA + 34 KB indexer) per block fits 12 fp32-SSM KDA states. + assert len(mamba) == 3 + assert sum(len(g.layer_names) for g in mamba) == 34 + + kv_cache_config = kv_cache_planning.get_kv_cache_config_from_groups( + vllm_config, groups, available_memory=2 * 1024**3 + ) + strides = {t.block_stride for t in kv_cache_config.kv_cache_tensors} + assert len(strides) == 1 + (block_stride,) = strides + # Blocks sit a whole number of 32-row MLA pages and 64-state indexer pages + # apart so both caches can be re-paged for their kernels. + assert block_stride % (32 * 576 * 2) == 0 + assert block_stride % (64 * 132) == 0 + + +def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): + return SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + max_tp_shards=1, + head_size=head_size, + dtype=torch.float32, + sliding_window=sliding_window, + model_version=model_version, + ) + + +def new_indexer_mla_spec(block_size=16): + # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller + # page size than the main MLA attention (uint8, small head), so their pages + # cannot be unified. + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + ) + + +def test_mixed_page_size_groups_use_spec_compatibility(): + specs = {} + for i in range(3): + specs[f"mla.{i}"] = new_mla_spec() + specs[f"indexer.{i}"] = new_indexer_mla_spec() + specs.update({f"swa.{i}": new_swa_mla_spec(head_size=1024) for i in range(5)}) + + config = _grouping_config() + config.cache_config = CacheConfig() + config.cache_config.kv_cache_layout = "BLNHC" + groups = get_kv_cache_groups(config, specs) + + assert len(groups) == 3 + assert {name for group in groups for name in group.layer_names} == set(specs) + assert sorted(len(group.layer_names) for group in groups) == [2, 3, 6] + + +def _grouping_config(): + cache_config = CacheConfig() + cache_config.kv_cache_layout = "LBNHC" + return SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + speculative_config=None, + cache_config=cache_config, + ) + + +def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): + block_size = 544 + full_spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + # FullAttentionSpec stores both K and V, so this produces a + # 544 * 1 * (512 + 512) * 2 = 1,114,112-byte page. + head_size=512, + dtype=torch.bfloat16, + ) + mamba_spec = MambaSpec( + block_size=block_size, + shapes=((557056,),), + dtypes=(torch.bfloat16,), + mamba_cache_mode="align", + ) + hidden_spec = HiddenStateCacheSpec( + block_size=block_size, + num_kv_heads=3, + head_size=1024, + dtype=torch.bfloat16, + ) + assert full_spec.page_size_bytes == mamba_spec.page_size_bytes + + groups = get_kv_cache_groups( + _grouping_config(), + { + "model.full_attn": full_spec, + "model.mamba": mamba_spec, + "cache_only_layers.0": hidden_spec, + }, + ) + + hidden_group = next( + group + for group in groups + if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) + ) + assert hidden_group.kv_cache_spec.block_size == 136 + + kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=groups, + ) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + prefix_match_unit=None, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + kv_transfer_config=object(), + ) + assert kv_cache_utils.resolve_kv_cache_block_sizes( + kv_cache_config, vllm_config + ) == (544, 136) + + +def test_multi_run_layer_compact_strides_place_hoisted_heads(): + """A layer-compact run region is its own dense allocation: under LHBNC the head + groups sit between the layers and the blocks, so a run's block stride is one head + group's slice, not the whole page (regression test for setStorage out-of-bounds + on ROCm hybrid models).""" + full = new_kv_cache_spec() + swa = new_sliding_window_spec(sliding_window=full.block_size * 2) + page = full.page_size_bytes + assert swa.page_size_bytes == page + + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) + vllm_config.cache_config.kv_cache_layout = "LHBNC" + config = kv_cache_planning.get_kv_cache_config_from_groups( + vllm_config, + [ + KVCacheGroupSpec(["full.0", "full.1"], full), + KVCacheGroupSpec(["swa.0", "swa.1"], swa), + ], + available_memory=4 * page * 8, + ) + + num_blocks = config.num_blocks + head_group = full.block_size * full.state_content_size_bytes + for tensor in config.kv_cache_tensors: + # One head group per block, head groups between the layers and blocks. + assert tensor.block_stride == head_group + assert tensor.layer_stride == full.num_heads * num_blocks * head_group + # The last block of the last layer's last head group must stay within the + # allocation (the old page-based stride overran it). + end = ( + tensor.offset + + tensor.layer_stride * (len(tensor.layers) - 1) + + (full.num_heads - 1) * num_blocks * head_group + + tensor.block_stride * (num_blocks - 1) + + head_group + ) + assert end <= tensor.size + + +def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): + specs = { + "target.0.attn": new_mla_spec(), + "draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288), + } + assert len({spec.page_size_bytes for spec in specs.values()}) == 1 + + groups = get_kv_cache_groups(_grouping_config(), specs) + + assert len(groups) == 2 + assert all( + not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups + ) + + +def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): + # Sparse MLA pages cannot be padded safely. Keeping the draft's attention + # compute sliding-window while promoting only its allocation semantics lets + # every layer share the target's block table and remain contiguous. + draft = new_sliding_window_spec(block_size=16) + specs = { + "target.0.attn": new_mla_spec(block_size=64), + "target.0.indexer": new_indexer_mla_spec(block_size=64), + "draft.0": draft, + } + + groups = get_kv_cache_groups(_grouping_config(), specs) + assert len(groups) == 1 + assert set(groups[0].layer_names) == set(specs) + group_spec = groups[0].kv_cache_spec + assert isinstance(group_spec, UniformTypeKVCacheSpecs) + assert group_spec.block_size == 64 + promoted_draft = group_spec.kv_cache_specs["draft.0"] + assert isinstance(promoted_draft, FullAttentionSpec) + assert not isinstance(promoted_draft, SlidingWindowSpec) + assert promoted_draft.block_size == 64 + assert promoted_draft.sliding_window == draft.sliding_window + assert specs["draft.0"] is draft + assert "attention compute is unchanged" in caplog_vllm.text + + +def test_hidden_states_with_tp_scales_page_size(): + """When TP shrinks KV pages below the hidden-state per-token cost, + get_kv_cache_groups must scale up target block sizes so that the + common page accommodates the unsharded hidden states.""" + # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. + # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) + # = 16 * 2 * 64 * 2 = 4096 bytes. + kv_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.bfloat16, + ) + # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512) + # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page. + hs_spec = HiddenStateCacheSpec( + block_size=16, + num_kv_heads=6, + head_size=512, + dtype=torch.bfloat16, + ) + specs = { + "target.0.attn": kv_spec, + "target.1.attn": kv_spec, + "cache_only_layers.48": hs_spec, + } + + groups = get_kv_cache_groups(_grouping_config(), specs) + + # The hidden-state layer should be present and no assertion should fire. + all_layers = {name for g in groups for name in g.layer_names} + assert "cache_only_layers.48" in all_layers + + # The target group block sizes must have been scaled up. + for g in groups: + if "cache_only_layers.48" not in g.layer_names: + assert g.kv_cache_spec.block_size > kv_spec.block_size + + +def test_auto_fit_max_model_len(): + """Test that max_model_len=-1 auto-fits to available GPU memory.""" + # Create config with original_max_model_len=-1 to trigger auto-fit + model_config = ModelConfig(max_model_len=1024) + # Simulate the user passing -1 by setting original_max_model_len + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + + # With enough memory, max_model_len stays at the derived max + large_available_memory = mem_per_block_per_layer * 2 * 1024 # plenty of memory + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [large_available_memory] + ) + assert vllm_config.model_config.max_model_len == 1024 + + # Reset for next test + model_config = ModelConfig(max_model_len=1024) + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + # With limited memory, max_model_len should be reduced + # Need memory for at least max_model_len tokens + # 32 blocks worth of memory for 2 layers = can fit 32*16=512 tokens + limited_memory = mem_per_block_per_layer * 2 * 32 + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [limited_memory] + ) + # Should be reduced to fit in memory + assert vllm_config.model_config.max_model_len < 1024 + assert vllm_config.model_config.max_model_len > 0 + + +def test_auto_fit_max_model_len_with_hybrid(): + """Test that auto-fit works with hybrid KV cache specs.""" + # Create config with original_max_model_len=-1 to trigger auto-fit + model_config = ModelConfig(max_model_len=8192) + # Simulate the user passing -1 by setting original_max_model_len + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer + gamma = 2 + kv_cache_specs = { + "layer_1": new_mamba_spec(num_speculative_blocks=gamma), + "layer_2": new_kv_cache_spec(), + } + + # One extra block on top of what a 1024-token request needs: the pool + # reserves one block as the null block. + available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1) + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [available_memory] + ) + assert vllm_config.model_config.max_model_len == 1024 + + +def test_auto_fit_max_model_len_not_triggered(): + """Test that auto-fit is not triggered when original_max_model_len is not -1.""" + model_config = ModelConfig(max_model_len=16) + # original_max_model_len should be None by default, not -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + + # This should work normally without auto-fit + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32] + ) + assert vllm_config.model_config.max_model_len == 16 + + +def test_auto_fit_max_model_len_respects_num_gpu_blocks_override(): + """Auto-fit must size max_model_len against the override-clamped pool, not + the raw `available_memory`. Without this, auto-fit could pick a + max_model_len that no longer fits once `num_gpu_blocks_override` is applied. + """ + model_config = ModelConfig(max_model_len=16384) + model_config.original_max_model_len = -1 # request auto-fit + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + # Cap the cache to 32 blocks regardless of available memory. + vllm_config.cache_config.num_gpu_blocks_override = 32 + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), # block_size=16 + "layer_2": new_kv_cache_spec(), + } + # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384). + large_available_memory = mem_per_block_per_layer * 2 * 1024 + + get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) + + # 32 blocks * block_size 16 = 512 token slots, so max_model_len must + # auto-fit at or below that. + assert 0 < vllm_config.model_config.max_model_len <= 32 * 16 + + +def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override(): + """Admission check must use the override-clamped pool size, not raw + `available_memory`. Without this, startup could accept a max_model_len + that does not actually fit in `num_gpu_blocks_override` blocks. + """ + model_config = ModelConfig(max_model_len=16384) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + # 32 blocks is far too small for max_model_len=16384 (would need 1024). + vllm_config.cache_config.num_gpu_blocks_override = 32 + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + # Plenty of raw memory: a bytes-only check against this would pass. + large_available_memory = mem_per_block_per_layer * 2 * 1024 + + with pytest.raises(ValueError, match="max seq len"): + get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) + + +def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): + """DFlash drafters can have a smaller head size than the target model. + + For example, MiMo uses 192-dim target KV heads while its DFlash draft uses + 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer + block-size multiple, so the smaller page must be padded instead. + """ + # Attention layers read padded pages through a strided view, so padding is + # allowed. + target_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=1, + head_size=192, + dtype=torch.bfloat16, + ) + draft_spec = new_sliding_window_spec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + sliding_window=1024, + ) + + unified_specs = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "target_attn": target_spec, + "draft_attn": draft_spec, + } + ) + + assert unified_specs["target_attn"] == target_spec + unified_draft_spec = unified_specs["draft_attn"] + assert unified_draft_spec.block_size == draft_spec.block_size + assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes + assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes + assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes + + +def test_unify_hybrid_kv_cache_specs(): + # 1. has_full_attention and has_sliding_window + before_spec_1 = new_kv_cache_spec(block_size=64) + before_spec_2 = new_sliding_window_spec( + block_size=16, page_size_padded=32 * 1024, sliding_window=1024 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + } + kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False) + kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False) + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec(block_size=64) + expected_spec_2 = new_kv_cache_spec( + block_size=64, page_size_padded=64 * 1024, sliding_window=1024 + ) + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024 + assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False) + + # 2. has_full_attention and has_chunked_local_attention + before_spec_1 = new_kv_cache_spec() + before_spec_2 = new_chunked_local_attention_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + } + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec() + expected_spec_2 = new_kv_cache_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + + # 3. has_full_attention, has_sliding_window and has_chunked_local_attention + before_spec_1 = new_kv_cache_spec() + before_spec_2 = new_sliding_window_spec( + page_size_padded=32 * 1024, sliding_window=1024 + ) + before_spec_3 = new_chunked_local_attention_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + "layer_3": before_spec_3, + } + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec() + expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024) + expected_spec_3 = new_kv_cache_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + assert kv_cache_spec["layer_3"] == expected_spec_3 + + # 4. No FullAttentionSpec, should not convert + kv_cache_spec = { + "layer_1": new_sliding_window_spec(sliding_window=1024), + "layer_2": new_chunked_local_attention_spec(attention_chunk_size=512), + } + + with pytest.raises(ValueError): + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + + # Replicated Mamba state still requires the hybrid cache manager. + kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()} + with pytest.raises(ValueError): + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + + +def test_unify_kv_cache_spec_page_size_mamba(): + """Regression test for https://github.com/vllm-project/vllm/issues/43626. + + MambaSpec's page_size_bytes is determined by its state shapes and does not + change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba + page instead of scaling its block_size. This situation arises when a layer + with a page larger than the (already platform-aligned) Mamba page joins the + specs, e.g. a dense draft model with more KV heads than the hybrid main + model. + """ + # 1. Hybrid main model (Mamba + full attention, pages already aligned at + # 16KB) plus a dense draft model layer with a 2x larger page (32KB). + # Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so + # the old code scaled the Mamba block_size, which left page_size_bytes + # unchanged at 16384. + mamba_spec = new_mamba_spec() # page_size_bytes = 16384 + main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384 + draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768 + assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384 + assert draft_attn_spec.page_size_bytes == 32768 + + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": mamba_spec, + "main_attn_layer": main_attn_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + # Mamba page is padded; block_size (caching granularity) is unchanged. + assert unified["mamba_layer"].page_size_bytes == 32768 + assert unified["mamba_layer"].page_size_padded == 32768 + assert unified["mamba_layer"].block_size == mamba_spec.block_size + # Attention layer with smaller page still unifies by scaling block_size. + assert unified["main_attn_layer"].page_size_bytes == 32768 + assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size + assert unified["main_attn_layer"].page_size_padded is None + # Layer already at max page size is unchanged. + assert unified["draft_attn_layer"] == draft_attn_spec + + # 2. Mamba page already padded by the platform (state smaller than the + # padded page); the padding is re-applied at the new maximum. + padded_mamba_spec = new_mamba_spec( + shapes=((2, 256), (3, 32, 32)), page_size_padded=16384 + ) + assert padded_mamba_spec.page_size_bytes == 16384 + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": padded_mamba_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + assert unified["mamba_layer"].page_size_bytes == 32768 + assert unified["mamba_layer"].page_size_padded == 32768 + + # 3. Mamba page that does not evenly divide the maximum page size is + # padded as well (the divisibility constraint only applies to block_size + # scaling). + odd_mamba_spec = new_mamba_spec(shapes=((6144,),)) + assert odd_mamba_spec.page_size_bytes == 24576 + assert 32768 % odd_mamba_spec.page_size_bytes != 0 + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": odd_mamba_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + assert unified["mamba_layer"].page_size_bytes == 32768 + + # 4. Attention layers with non-divisible page sizes are padded too: every + # backend reads a padded page through the view's block stride, so there is + # no longer a case that must raise. + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "attn_layer": new_kv_cache_spec(block_size=24), # 24576 + "draft_attn_layer": draft_attn_spec, # 32768 + } + ) + assert unified["attn_layer"].page_size_padded == 32768 + assert unified["attn_layer"].page_size_bytes == 32768 + assert unified["attn_layer"].block_size == 24 + + # 5. Uniform page sizes are returned unchanged. + specs = { + "mamba_layer": new_mamba_spec(), + "attn_layer": new_kv_cache_spec(), + } + assert kv_cache_planning.unify_kv_cache_spec_page_size(specs) == specs + + +@pytest.mark.parametrize("use_override", [True, False]) +def test_kv_cache_reserves_null_block_for_max_model_len(use_override): + """A KV cache sized to exactly the blocks a max_model_len request needs must + be rejected. BlockPool keeps one block as the null block, so only + num_blocks - 1 are usable; accepting num_blocks == needed would leave the + request unschedulable and hang the engine. Covers both the + num_gpu_blocks_override path and the memory-derived block count. + """ + block_size = 16 + max_model_len = 512 # needs 512 / 16 = 32 blocks + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + spec = new_kv_cache_spec(block_size=block_size) + + # 32 blocks -> only 31 usable after the null block: one short -> reject. + if use_override: + # Ample raw memory; the override alone constrains the block count. + vllm_config.cache_config.num_gpu_blocks_override = 32 + available_memory = [spec.page_size_bytes * 1024] + else: + available_memory = [spec.page_size_bytes * 32] + + with pytest.raises(ValueError, match="max seq len"): + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + +def test_auto_fit_max_model_len_reserves_null_block(): + """Auto-fit (max_model_len=-1) must size max_model_len against usable + blocks, not the total pool. With memory for exactly 64 blocks, one is the + null block, so auto-fit must settle on 63 * block_size; picking the full + pool would leave a max-length request unschedulable and hang the engine. + """ + block_size = 16 + model_config = ModelConfig(max_model_len=1024) + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + spec = new_kv_cache_spec(block_size=block_size) + + # Exactly the 1024 / 16 = 64 blocks a full-length request would need. + available_memory = [spec.page_size_bytes * 64] + + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + assert vllm_config.model_config.max_model_len == 63 * block_size + + +def test_check_enough_kv_cache_memory_reserves_null_block(): + """The public admission check must reject a KV cache sized to exactly the + blocks a max_model_len request needs. BlockPool keeps one block as the + null block, so only num_blocks - 1 are usable; accepting + num_blocks == needed would leave the request unschedulable and hang the + engine. + """ + block_size = 16 + max_model_len = 512 # needs 512 / 16 = 32 blocks + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + spec = new_kv_cache_spec(block_size=block_size) + + # 32 blocks -> only 31 usable after the null block: one short -> reject. + with pytest.raises(ValueError, match="max seq len"): + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 32 + ) + + # 33 blocks -> 32 usable after the null block -> accept. + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 33 + ) + + +def _spec_decode_grouping_config(method="dspark", model_type=None): + """Grouping config with an EAGLE-family speculative method enabled.""" + return SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + cache_config=SimpleNamespace( + get_resolved_kv_cache_layout=lambda: SimpleNamespace( + is_block_outermost=True + ) + ), + model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)), + speculative_config=SimpleNamespace( + method=method, + use_eagle=lambda: True, + use_eagle_block_drop=lambda: True, + ), + ) + + +def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False): + """A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a + DSpark-style draft MLA layer marked non_causal_multi_token_decode. + + The target's fp8 KV dtype is one difference from the draft on Kimi-K3 + (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass + draft_shares_target_spec to match that dtype while retaining the draft + marker as a group boundary. + """ + target_dtype = None if draft_shares_target_spec else "fp8_e4m3" + specs = { + "target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), + "target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), + "target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"), + "target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"), + } + if draft: + draft_spec = new_mla_spec(block_size=64) + specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True) + return specs + + +def test_draft_group_annotated_on_hybrid_general_path(): + # A drafter's MLA layer carries non_causal_multi_token_decode, so its group + # is identifiable without keying off a model version. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "draft.attn.0" in flagged[0].layer_names + + +def test_mamba_groups_never_flagged_even_when_draft_shares_a_group(): + # Packed uniform-type groups can contain distinct target and draft layer + # specs; the combined group still holds volatile draft KV and must be + # flagged. What must never happen is a Mamba group being flagged: that + # widens its lookup window to two consecutive chunks, which align-mode + # checkpointing never produces, zeroing every lookup. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), + _hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True), + ) + + for group in groups: + if "draft.attn.0" in group.layer_names: + assert group.is_eagle_group + if isinstance(group.kv_cache_spec, MambaSpec): + assert not group.is_eagle_group + + +def test_draft_group_not_annotated_without_spec_decode(): + # The marker alone must not flag anything; the eagle semantics only apply + # when a speculative method is actually enabled. + config = _spec_decode_grouping_config() + config.speculative_config = None + groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True)) + + assert not any(g.is_eagle_group for g in groups) + + +def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): + # No group carries the draft marker, so consumers fall back to + # conservative behavior that silently breaks reuse for Mamba groups. + # That must at least be visible. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) + ) + + assert not any(g.is_eagle_group for g in groups) + assert "no KV cache group could be identified as the draft model's" in ( + caplog_vllm.text + ) + + +def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): + # Pure-attention models degrade gracefully under the consumers' + # conservative fallback (a one-block hit drop at most), so the warning + # stays silent to avoid noise on every unannotated EAGLE deployment. + specs = { + "target.attn.0": new_mla_spec(block_size=64), + "target.attn.1": new_mla_spec(block_size=64), + } + groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs) + + assert not any(g.is_eagle_group for g in groups) + assert "could be identified as the draft model's" not in caplog_vllm.text + + +def test_no_warning_when_draft_group_is_identified(caplog_vllm): + get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) + ) + + assert "could be identified as the draft model's" not in caplog_vllm.text + + +def _deepseek_v4_specs(model_version="deepseek_v4"): + """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers + at differing page sizes, with the MTP draft layer registered last.""" + return { + "model.layers.0.self_attn.attn": new_mla_spec(), + "model.layers.1.self_attn.attn": new_mla_spec(), + "model.layers.2.self_attn.attn": new_swa_mla_spec( + head_size=1024, model_version=model_version + ), + # The MTP block registers last, and its sliding-window size differs, so + # it lands in a group of its own. + "model.layers.3.self_attn.attn": new_swa_mla_spec( + head_size=1024, sliding_window=256, model_version=model_version + ), + } + + +@pytest.mark.parametrize( + ("method", "model_type"), + [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")], +) +def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type): + # DeepseekV4's MTP block reuses the target's decoder layer, so its spec + # carries no draft marker and only the positional rule can find it. This + # pins the pre-existing behaviour that the unified annotator must preserve. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method=method, model_type=model_type), + _deepseek_v4_specs(model_version=None), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "model.layers.3.self_attn.attn" in flagged[0].layer_names + + +def test_trailing_layer_fallback_applies_to_any_mtp_model(): + # The positional rule is sound for every MTP drafter, not just DeepseekV4: + # MTP blocks reuse the target's decoder layer (no spec marker) and always + # register after every target layer. The model_type must not gate it. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type="other"), + _deepseek_v4_specs(), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "model.layers.3.self_attn.attn" in flagged[0].layer_names + + +def _qwen3_5_hybrid_specs(with_mtp_layer: bool): + """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with + the MTP drafter's full-attn layer (spec-identical to the target's) + registered last.""" + specs = {} + idx = 0 + for _ in range(2): + for _ in range(3): + specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec( + mamba_cache_mode="align" + ) + idx += 1 + specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec() + idx += 1 + if with_mtp_layer: + specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec() + return specs + + +def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm): + # A hybrid mamba + full-attention model with an MTP drafter that is + # spec-indistinguishable from the target reaches the general multi-group + # path. The trailing-layer rule must locate the draft group there so the + # Mamba groups are not swept up by the flag-all consumer fallback. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"), + _qwen3_5_hybrid_specs(with_mtp_layer=True), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names + for group in groups: + if any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ): + assert not group.is_eagle_group + assert "could be identified as the draft model's" not in caplog_vllm.text + + +@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"]) +def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method): + # Other drafters are not covered by the trailing-layer rule, so an + # unidentifiable hybrid draft must still warn. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method=method, model_type="qwen3_5"), + _qwen3_5_hybrid_specs(with_mtp_layer=True), + ) + + assert not any(g.is_eagle_group for g in groups) + assert "could be identified as the draft model's" in caplog_vllm.text + + +def test_trailing_layer_fallback_requires_exact_partition(): + # If the groups do not partition the layers exactly (e.g. a caller that + # dropped or duplicated layers), the positional rule is meaningless and + # must not fire. + from vllm.v1.core.kv_cache_planning import _annotate_eagle_groups + + specs = _qwen3_5_hybrid_specs(with_mtp_layer=True) + config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5") + groups = get_kv_cache_groups(config, specs) + for g in groups: + g.is_eagle_group = False + # Remove one layer from its group: no longer an exact partition. + trimmed = [ + KVCacheGroupSpec( + [n for n in g.layer_names if n != "model.layers.0.linear_attn"], + g.kv_cache_spec, + ) + for g in groups + ] + _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) + + assert not any(g.is_eagle_group for g in trimmed) + + +_GQA_SPEC = FullAttentionSpec( + block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32 +) + + +_GQA_SWA_SPEC = SlidingWindowSpec( + block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128 +) + + +@pytest.mark.parametrize( + "specs,expected", + [ + pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"), + pytest.param( + { + "l.0": new_mla_spec(), + "l.1": new_mla_spec(), + "s.0": new_swa_mla_spec(), + "s.1": new_swa_mla_spec(), + }, + 4, + id="mla-and-swa-mla", + ), + pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"), + pytest.param( + {"l.0": new_mla_spec(), "l.1": _GQA_SPEC}, + 1, + id="mla-and-gqa-uniform-group", + ), + pytest.param( + { + "l.0": new_mla_spec(), + "l.1": new_mla_spec(), + "s.0": _GQA_SWA_SPEC, + "s.1": _GQA_SWA_SPEC, + }, + 1, + id="mla-and-gqa-swa", + ), + ], +) +def test_kv_tp_replicas(monkeypatch, specs, expected): + """Resolved per layer before scheduler flattening hides mixed groups.""" + from vllm.config import ParallelConfig + from vllm.platforms import current_platform + + monkeypatch.setattr(current_platform, "device_count", lambda: 4) + parallel_config = ParallelConfig(tensor_parallel_size=4) + vllm_config = VllmConfig( + model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config + ) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + mem = sum(s.page_size_bytes for s in specs.values()) * 10 + configs = get_kv_cache_configs(vllm_config, [specs], [mem]) + assert configs[0].kv_tp_replicas == expected + scheduler = generate_scheduler_kv_cache_config(configs) + assert scheduler.kv_tp_replicas == expected + + +@pytest.mark.parametrize( + "specs,tp_size,dcp_size,expected", + [ + pytest.param([new_mla_spec()], 8, 1, 8, id="mla"), + pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"), + pytest.param( + [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset" + ), + pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"), + pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"), + pytest.param( + [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)], + 8, + 1, + 2, + id="mla-and-gqa-partial", + ), + pytest.param( + [ + HiddenStateCacheSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32 + ) + ], + 8, + 1, + 1, + id="hidden-state", + ), + ], +) +def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected): + """Replicas are the gcd of each layer's tp_size // max_tp_shards.""" + groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)] + assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index b06479f1ca..1489315ec2 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3,20 +3,23 @@ import copy import hashlib import importlib -import mmap -import subprocess -import sys from collections.abc import Callable from dataclasses import replace -from pathlib import Path from types import SimpleNamespace from typing import Any import pytest import torch +import vllm.v1.core.kv_cache_planning as kv_cache_planning import vllm.v1.core.kv_cache_utils as kv_cache_utils import vllm.v1.hisparse.runtime as hisparse_runtime_module +from tests.v1.core.utils import ( + new_kv_cache_spec, + new_mamba_spec, + new_mla_spec, + new_sliding_window_spec, +) from vllm.config import ( CacheConfig, ModelConfig, @@ -33,25 +36,21 @@ from vllm.multimodal.inputs import ( ) from vllm.sampling_params import SamplingParams from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor -from vllm.utils.mem_constants import GiB_bytes from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager +from vllm.v1.core.kv_cache_planning import ( + get_kv_cache_groups, +) from vllm.v1.core.kv_cache_utils import ( BlockHash, FreeKVCacheBlockQueue, KVCacheBlock, - check_enough_kv_cache_memory, - estimate_max_model_len, generate_block_hash_extra_keys, generate_scheduler_kv_cache_config, get_kv_cache_capacity, - get_kv_cache_configs, - get_kv_cache_groups, get_max_concurrency_for_kv_cache_config, get_request_block_hasher, hash_block_tokens, init_none_hash, - is_kv_cache_spec_uniform, - kv_cache_groups_tp_replicas, make_block_hash_with_group_id, tensor_data, to_event_extra_keys, @@ -62,10 +61,8 @@ from vllm.v1.hisparse.layout import ( get_hisparse_steady_state_concurrency, ) from vllm.v1.kv_cache_interface import ( - ChunkedLocalAttentionSpec, CircularBufferSpec, FullAttentionSpec, - HiddenStateCacheSpec, HiSparseHotSpec, HiSparseResidentSpec, KVCacheConfig, @@ -93,110 +90,6 @@ from vllm.v1.request import Request pytestmark = pytest.mark.cpu_test -@pytest.mark.parametrize("gpu_block_size", [32, 64]) -@pytest.mark.parametrize("shared_host_pool", [False, True]) -def test_hisparse_hma_uses_resolved_gpu_block_size( - monkeypatch, gpu_block_size, shared_host_pool -): - monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True - ) - specs = { - "model.layers.0.self_attn": MLAAttentionSpec( - block_size=gpu_block_size, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - is_index_group_leader=True, - ), - "model.layers.0.self_attn.indexer": MLAAttentionSpec( - block_size=gpu_block_size, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - cache_role=SparseCacheRole.INDEXER, - ), - } - config = SimpleNamespace( - attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), - model_config=SimpleNamespace( - hf_config=SimpleNamespace(index_topk=128), - max_model_len=gpu_block_size, - ), - parallel_config=SimpleNamespace( - tensor_parallel_size=2 if shared_host_pool else 1, - pipeline_parallel_size=1, - prefill_context_parallel_size=1, - decode_context_parallel_size=1, - world_size=2 if shared_host_pool else 1, - distributed_executor_backend="mp", - nnodes_within_dp=1, - ), - cache_config=SimpleNamespace( - num_gpu_blocks_override=7, - prefix_cache_retention_interval=None, - get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC, - ), - ) - monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - config, get_hisparse_kv_cache_groups(config, specs), available_memory=2**30 - ) - assert cache_config.num_blocks == 7 - assert cache_config.hisparse_host_num_blocks is not None - assert cache_config.hisparse_host_num_blocks > 7 - - host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups - assert cache_config.hisparse_shared_host_pool is shared_host_pool - host_page = host_group.kv_cache_spec.page_size_bytes - alignment = mmap.PAGESIZE if shared_host_pool else 1 - expected_host_stride = (host_page + alignment - 1) // alignment * alignment - assert cache_config.hisparse_host_block_stride == expected_host_stride - assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride - assert host_group.host_resident - assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) - host_layers = set(host_group.layer_names) - for tensor in cache_config.kv_cache_tensors: - assert all( - (name in host_layers) == tensor.host_resident for name in tensor.layers - ) - host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident) - host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]] - assert host_tensor.block_stride == host_spec.page_size_bytes - assert host_tensor.layer_stride == ( - host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks - ) - assert host_group.kv_cache_spec.block_size == gpu_block_size - assert indexer_group.kv_cache_spec.block_size == gpu_block_size - host_specs = host_group.kv_cache_spec.kv_cache_specs - gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs - assert set(host_specs) == {"model.layers.0.self_attn"} - assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"} - assert indexer_group.kv_cache_spec.prefix_cacheable - assert host_group.enable_kv_transfer - auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups] - assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs) - assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs) - assert all( - spec.block_size == gpu_block_size - for spec in auxiliary_specs - if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)) - ) - scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes( - cache_config, - SimpleNamespace( - cache_config=SimpleNamespace( - block_size=16, - enable_prefix_caching=True, - prefix_match_unit=None, - ), - parallel_config=SimpleNamespace(decode_context_parallel_size=1), - kv_transfer_config=object(), - ), - ) - assert scheduler_block_size == hash_block_size == gpu_block_size - - def test_hisparse_steady_state_concurrency_excludes_spilled_resident_pages(): """Running requests pin only their resident tail; the newest still needs its full in-flight window to be admitted.""" @@ -281,68 +174,6 @@ def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok create_hisparse_layout(config, groups, host_budget=host_budget) -@pytest.mark.parametrize( - "max_model_len,num_gpu_blocks,ok", - [ - (8192, 160, False), - (8192, 1000, True), - (32768, 656, False), - (32768, 657, True), - (-1, 657, True), - ], -) -def test_hisparse_pool_must_fit_max_model_len( - monkeypatch, max_model_len, num_gpu_blocks, ok -): - """Each resident group takes its own blocks from HiSparse's shared GPU pool, - so a pool that fits only the indexer pages must be rejected at startup. - Resident pages are capped at the in-flight window: at 32768 tokens with 2048 - in flight, one request needs 512 indexer + 4 x 32 resident + 16 hot blocks, - plus the null block. Auto-fit (-1) must pick a length that passes the same - check.""" - monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True - ) - monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) - specs: dict[str, KVCacheSpec] = {} - for i in range(4): - specs[f"model.layers.{i}.self_attn"] = MLAAttentionSpec( - block_size=64, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - is_index_group_leader=True, - ) - specs[f"model.layers.{i}.self_attn.indexer"] = MLAAttentionSpec( - block_size=64, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - cache_role=SparseCacheRole.INDEXER, - ) - config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) - config.scheduler_config.max_num_batched_tokens = 2048 - # Async scheduling doubles the in-flight window, and with it the bound. - config.scheduler_config.async_scheduling = False - config.attention_config.hisparse_config = HiSparseConfig() - config.model_config.hf_config.index_topk = 128 - config.cache_config.num_gpu_blocks_override = num_gpu_blocks - config.cache_config.kv_cache_layout = "BLHNC" - # Workers report the resident/hot caches HiSparse derives. - specs |= { - name: group.kv_cache_spec - for group in get_hisparse_kv_cache_groups(config, specs) or [] - if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec)) - for name in group.layer_names - } - - if ok: - kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34]) - else: - with pytest.raises(ValueError, match="max seq len"): - kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34]) - - def test_hisparse_derived_specs_do_not_affect_kv_cache_layout(monkeypatch): """Workers report HiSparse's resident/hot caches, which lay out their own raw backing; KV cache layout resolution must skip them.""" @@ -496,28 +327,6 @@ def test_effective_attention_block_size_matches_events(dcp): ] == [block_size] -def new_kv_cache_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - sliding_window=None, - attention_chunk_size=None, - kv_quant_mode=KVQuantMode.NONE, -): - return FullAttentionSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - sliding_window=sliding_window, - attention_chunk_size=attention_chunk_size, - kv_quant_mode=kv_quant_mode, - ) - - def test_kv_cache_config_selects_only_transferable_groups(): """Connectors must see a stable projection of transfer-eligible groups.""" groups = [ @@ -581,60 +390,6 @@ def test_kv_cache_blocks_selects_requested_groups(): assert blocks.get_block_ids(group_ids=(0, 2)) == ([1], [3]) -def new_sliding_window_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - sliding_window=1, -): - return SlidingWindowSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - sliding_window=sliding_window, - ) - - -def new_chunked_local_attention_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - attention_chunk_size=4, -): - return ChunkedLocalAttentionSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - attention_chunk_size=attention_chunk_size, - ) - - -def new_mamba_spec( - block_size=16, - shapes=((2, 512), (3, 32, 32)), - dtypes=(torch.float32, torch.float32), - num_speculative_blocks=2, - mamba_cache_mode="none", - page_size_padded=None, -): - return MambaSpec( - block_size=block_size, - shapes=shapes, - dtypes=dtypes, - page_size_padded=page_size_padded, - mamba_cache_mode=mamba_cache_mode, - num_speculative_blocks=num_speculative_blocks, - ) - - @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_none_hash(monkeypatch, hash_fn): import vllm.v1.core.kv_cache_utils @@ -1441,430 +1196,6 @@ def test_metrics_empty_stats(): assert metrics.hit_rate == 0.75 -def test_get_kv_cache_configs_multiple_workers(): - model_config = ModelConfig(max_model_len=16) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - ref_kv_cache_spec = new_kv_cache_spec() - same_kv_cache_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - ] - - # Basic case. All things are the same. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - same_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - expected = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ) - assert kv_cache_configs == [expected, expected] - - # Different available memory. This is the case for TP. - # Use the smallest memory available. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - same_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 20, - ], - ) - assert kv_cache_configs == [expected, expected] - - # Different KV cache specs. This is the case for PP. - different_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - }, - { - "layer2": new_kv_cache_spec(), - "layer3": new_kv_cache_spec(), - }, - ] - - # Different workers have different layers. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1"], new_kv_cache_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer2", "layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer2", "layer3"], new_kv_cache_spec()), - ], - ), - ] - - # Some layers are the same, some are different. This is the case for TP+PP - tp_pp_kv_cache_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer3": new_kv_cache_spec(), - }, - { - "layer3": new_kv_cache_spec(), - }, - ] - - kv_cache_configs = get_kv_cache_configs( - vllm_config, - tp_pp_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - expected_12 = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ) - expected_3 = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), - ], - ) - assert kv_cache_configs == [ - expected_12, - expected_12, - expected_3, - expected_3, - ] - - # Different workers have different types of layers. This is the case for - # hybrid models + PP. - different_type_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer3": new_sliding_window_spec(), - "layer4": new_sliding_window_spec(), - }, - ] - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_type_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - KVCacheGroupSpec([], new_sliding_window_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer3", "layer4"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec([], ref_kv_cache_spec), - KVCacheGroupSpec(["layer3", "layer4"], new_sliding_window_spec()), - ], - ), - ] - - # When divided into multiple KVCacheGroups, need to ensure the number of - # layers per group is similar. - different_type_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_sliding_window_spec(), - "layer3": new_sliding_window_spec(), - }, - { - "layer4": new_kv_cache_spec(), - "layer5": new_sliding_window_spec(), - "layer6": new_sliding_window_spec(), - }, - ] - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_type_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 10, - ref_kv_cache_spec.page_size_bytes * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1"], ref_kv_cache_spec), - KVCacheGroupSpec(["layer2"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer3"], new_sliding_window_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer4"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer5"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer6"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer4"], ref_kv_cache_spec), - KVCacheGroupSpec(["layer5"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer6"], new_sliding_window_spec()), - ], - ), - ] - - # Have conflicting layers. Need to raise an error. - conflicting_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - }, - { - "layer1": new_sliding_window_spec(), - }, - ] - with pytest.raises(AssertionError): - get_kv_cache_configs( - vllm_config, - conflicting_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - - -@pytest.mark.parametrize( - "asymmetric_memory", - [False, True], - ids=["symmetric", "asymmetric"], -) -def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): - model_config = ModelConfig(max_model_len=512) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - ref_kv_cache_spec = new_kv_cache_spec() - pp_kv_cache_specs = [ - {"layer1": ref_kv_cache_spec}, - {"layer2": ref_kv_cache_spec}, - ] - - expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1 - avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks - - # With per-worker validation, each worker only needs memory for its own - # layers. Worker 2 having more memory shouldn't affect worker 1's config. - available_memory = ( - [avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2 - ) - - kv_cache_configs = get_kv_cache_configs( - vllm_config, - pp_kv_cache_specs, - available_memory, - ) - - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=expected_num_blocks, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes - * expected_num_blocks, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)], - ), - KVCacheConfig( - num_blocks=expected_num_blocks, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, - layers=["layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes - * expected_num_blocks, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)], - ), - ] - - -def test_project_kv_cache_groups_to_worker(): - spec_a = new_kv_cache_spec() - spec_b = new_kv_cache_spec(num_kv_heads=4) - - global_groups = [ - KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a), - ] - worker_spec = {"layer1": spec_a, "layer2": spec_a} - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups, worker_spec - ) - assert len(projected) == 1 - assert projected[0].layer_names == ["layer1", "layer2"] - assert projected[0].kv_cache_spec is spec_a - - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups, {"layer4": spec_a} - ) - assert len(projected) == 1 - assert projected[0].layer_names == [] - assert projected[0].kv_cache_spec is spec_a - - uniform_spec = UniformTypeKVCacheSpecs( - block_size=16, - kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a}, - ) - global_groups_uniform = [ - KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec), - ] - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups_uniform, {"layer1": spec_a, "layer3": spec_a} - ) - assert len(projected) == 1 - assert projected[0].layer_names == ["layer1", "layer3"] - proj_spec = projected[0].kv_cache_spec - assert isinstance(proj_spec, UniformTypeKVCacheSpecs) - assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"} - - @pytest.mark.parametrize("sliding_window", [None, 256]) @pytest.mark.parametrize("disable_hybrid", [False, True]) @pytest.mark.parametrize("pcp_size", [1, 4]) @@ -1901,7 +1232,7 @@ def test_dcp_target_allocates_replicated_draft_independently( widths = [g.kv_cache_spec.max_num_blocks_per_req(config, 1024) for g in groups] assert widths == [16, 64] - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( config, groups, available_memory=16 * 1024 * 1024 ) scheduler_config = generate_scheduler_kv_cache_config([cache_config]) @@ -1971,58 +1302,6 @@ def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla): type(draft).merge([draft, replace(draft, dcp_sharded=True)]) -@pytest.mark.parametrize("with_draft", [False, True]) -@pytest.mark.parametrize("indexer_alignment", [None, 512]) -def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment): - from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view - from vllm.v1.worker.utils import allocate_kv_cache - - config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) - config.cache_config.kv_cache_layout = "BLHNC" - config.parallel_config.decode_context_parallel_size = 4 - common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8) - specs = { - "target": MLAAttentionSpec( - **common, - head_size=576, - state_content_bytes=656, - block_stride_alignment=656, - cache_dtype_str="fp8_ds_mla", - is_index_group_leader=True, - ), - "indexer": MLAAttentionSpec( - **common, - head_size=132, - cache_role=SparseCacheRole.INDEXER, - block_stride_alignment=indexer_alignment, - ), - "draft": SlidingWindowSpec( - block_size=64, - num_kv_heads=64, - head_size=64, - dtype=torch.bfloat16, - sliding_window=2048, - dcp_sharded=False, - ), - } - if not with_draft: - del specs["draft"] - groups = get_kv_cache_groups(config, specs) - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - config, groups, available_memory=8 * 1024 * 1024 - ) - caches = allocate_kv_cache( - cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups) - ) - cache = caches["target"].squeeze(1) - if indexer_alignment is not None: - assert caches["indexer"].stride(0) % indexer_alignment == 0 - rows, stride_rows = flat_kv_row_view(cache, 64) - cache[1, 0].fill_(7) - torch.testing.assert_close(rows[stride_rows], cache[1, 0]) - assert torch.all(rows[stride_rows] == 7) - - @pytest.mark.parametrize( "layer_type,dcp_size,expected_width", [ @@ -2115,103 +1394,6 @@ def test_merge_kv_cache_spec(): assert merged_layer_spec.sliding_window == 1 -def test_is_kv_cache_spec_uniform(): - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_kv_cache_spec(num_kv_heads=32), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_kv_cache_spec(num_kv_heads=32, sliding_window=1), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - } - assert not is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=2), - } - assert not is_kv_cache_spec_uniform(kv_cache_spec) - - script = """ -import sys - -from vllm.v1.core.kv_cache_utils import is_kv_cache_spec_uniform -from vllm.v1.kv_cache_interface import KVCacheSpec - -if sys.flags.optimize < 1: - raise RuntimeError("subprocess is not running with optimization enabled") -specs = { - "a": KVCacheSpec(block_size=1), - "b": KVCacheSpec(block_size=2), -} -if is_kv_cache_spec_uniform(specs): - raise RuntimeError("different specs were treated as uniform") -""" - subprocess.run( - [sys.executable, "-O", "-c", script], - cwd=Path(__file__).resolve().parents[3], - check=True, - timeout=60, - ) - - -@pytest.mark.parametrize( - ("model_id", "max_model_len", "want_estimated_max_len"), - [ - ("Qwen/Qwen1.5-7B", 16385, 16384), - ("Qwen/Qwen1.5-7B", 16383, 16383), - ], -) -def test_estimate_max_model_len(model_id, max_model_len, want_estimated_max_len): - # Create a VllmConfig - model_config = ModelConfig( - model_id, - runner="generate", - dtype="float16", - max_model_len=max_model_len, - ) - scheduler_config = SchedulerConfig( - max_num_batched_tokens=32768, - max_model_len=model_config.max_model_len, - is_encoder_decoder=model_config.is_encoder_decoder, - ) - - vllm_config = VllmConfig( - model_config=model_config, - scheduler_config=scheduler_config, - ) - - # Create KV cache specs - kv_cache_spec = {} - for i in range(32): - layer_name = f"layer_{i}" - kv_cache_spec[layer_name] = FullAttentionSpec( - block_size=16, - num_kv_heads=32, - head_size=128, - dtype=torch.float16, - ) - # Estimate the maximum model length, 16384 model_len need 8GB - estimated_max_len = estimate_max_model_len( - vllm_config, kv_cache_spec, 8 * GiB_bytes - ) - assert estimated_max_len == want_estimated_max_len - - def test_get_max_concurrency_for_kv_cache_config(): # Create a VllmConfig model_id = "Qwen/Qwen1.5-7B" @@ -2425,411 +1607,6 @@ def test_allocate_with_lookahead(): assert len(blocks.get_block_ids()[0]) == 2 -def test_get_kv_cache_config_one_worker(): - # pass max_model_len to pass check_enough_kv_cache_memory - model_config = ModelConfig(max_model_len=16) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - # all layers are full attention -> single group - kv_cache_specs_full = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - kv_cache_config_full = get_kv_cache_configs( - vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] - )[0] - print(kv_cache_config_full) - assert kv_cache_config_full == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], - ) - - # all layers are sliding window -> single group - kv_cache_specs_sliding = { - "layer_1": new_sliding_window_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_sliding = get_kv_cache_configs( - vllm_config, [kv_cache_specs_sliding], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_sliding == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec()) - ], - ) - - # full + sliding, but disable_hybrid_kv_cache_manager - vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2"], new_kv_cache_spec(sliding_window=1) - ), - ], - ) - vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False - - # full + sliding, with hybrid_kv_cache_manager - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=64, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 64, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 64, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 64, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 64, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1"], new_kv_cache_spec()), - KVCacheGroupSpec(["layer_2"], new_sliding_window_spec()), - ], - ) - - # 2 full + 4 sliding, 2 layers per group - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_sliding_window_spec(), - "layer_4": new_sliding_window_spec(), - "layer_5": new_sliding_window_spec(), - "layer_6": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_3", "layer_5"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_4", "layer_6"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()), - KVCacheGroupSpec(["layer_3", "layer_5"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer_4", "layer_6"], new_sliding_window_spec()), - ], - ) - - # 3 full + 7 sliding, pad to 3 full + 9 sliding - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_kv_cache_spec(), - "layer_4": new_sliding_window_spec(), - "layer_5": new_sliding_window_spec(), - "layer_6": new_sliding_window_spec(), - "layer_7": new_sliding_window_spec(), - "layer_8": new_sliding_window_spec(), - "layer_9": new_sliding_window_spec(), - "layer_10": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_1", "layer_2", "layer_3"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_4", "layer_7", "layer_10"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_5", "layer_8"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_6", "layer_9"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"], new_kv_cache_spec()), - KVCacheGroupSpec( - ["layer_4", "layer_7", "layer_10"], new_sliding_window_spec() - ), - KVCacheGroupSpec(["layer_5", "layer_8"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer_6", "layer_9"], new_sliding_window_spec()), - ], - ) - - # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss - # eagle where there is only one more full attention layer than sliding window layers - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_kv_cache_spec(), - "layer_4": new_kv_cache_spec(), - "layer_5": new_kv_cache_spec(), - "layer_6": new_kv_cache_spec(), - "layer_7": new_sliding_window_spec(), - "layer_8": new_sliding_window_spec(), - "layer_9": new_sliding_window_spec(), - "layer_10": new_sliding_window_spec(), - "layer_11": new_sliding_window_spec(), - } - - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32] - )[0] - print(kv_cache_config_hybrid) - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 6, - layers=[ - "layer_1", - "layer_2", - "layer_3", - "layer_4", - "layer_5", - "layer_6", - ], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 6, - layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6"], - new_kv_cache_spec(), - ), - KVCacheGroupSpec( - ["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], - new_sliding_window_spec(), - ), - ], - ) - - # different hidden size but same type, use UniformTypeKVCacheSpecs - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=128), - "layer_2": new_kv_cache_spec(head_size=64), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] - )[0] - # Layers of different page sizes pack densely into one allocation: the - # 2x-larger layer_1 takes the first region, layer_2 the next. - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 2 * 32, - block_stride=mem_per_block_per_layer * 2, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - offset=mem_per_block_per_layer * 2 * 32, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2"], - UniformTypeKVCacheSpecs( - block_size=16, kv_cache_specs=kv_cache_specs_hybrid - ), - ) - ], - ) - - # Different hidden size and different type, align by different block size - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=64), - "layer_2": new_sliding_window_spec(head_size=32), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)), - KVCacheGroupSpec( - ["layer_2"], new_sliding_window_spec(head_size=32, block_size=32) - ), - ], - ) - - # different hidden size that cannot be aligned by using different block size, - # but can be aligned by padding the smaller physical page. - swa_spec = new_sliding_window_spec(head_size=96) - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=64), - "layer_2": swa_spec, - } - - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - padded_page_size = swa_spec.page_size_bytes - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=42, - kv_cache_tensors=[ - KVCacheTensor( - size=padded_page_size * 42, - layers=["layer_1"], - layer_stride=padded_page_size * 42, - block_stride=padded_page_size, - ), - KVCacheTensor( - size=padded_page_size * 42, - layers=["layer_2"], - layer_stride=padded_page_size * 42, - block_stride=padded_page_size, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1"], - new_kv_cache_spec( - head_size=64, - page_size_padded=padded_page_size, - ), - ), - KVCacheGroupSpec( - ["layer_2"], - new_sliding_window_spec(head_size=96), - ), - ], - ) - - # Test num_gpu_blocks_override - vllm_config.cache_config.num_gpu_blocks_override = 16 - kv_cache_config_override_blocks = get_kv_cache_configs( - vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_override_blocks == KVCacheConfig( - num_blocks=16, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 16 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 16, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], - ) - - -def test_get_kv_cache_configs_attention_free(): - kv_cache_specs: dict[str, KVCacheSpec] = {} - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) - vllm_config.cache_config.prefix_cache_retention_interval = None - kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0]) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=[], - ) - ] - - def test_generate_uniform_type_kv_cache_specs(): # All layers are full attention, can be merged kv_cache_specs = { @@ -2903,201 +1680,6 @@ def test_generate_scheduler_kv_cache_config(): ) -def test_packed_groups_glm5_like_hybrid(): - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=8192)) - vllm_config.cache_config.kv_cache_layout = "BLHNC" - kpool, block_size = 4, 1024 - kv_cache_spec: dict[str, KVCacheSpec] = {} - for i in range(11): - kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - block_stride_alignment=32 * 576 * 2, - ) - kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=132, - dtype=torch.uint8, - tokens_per_state=kpool, - block_stride_alignment=64 * 132, - ) - kv_cache_spec[f"layers.{i}.tail"] = CircularBufferSpec( - block_size=kpool, - num_kv_heads=2, - head_size=128, - head_size_v=0, - dtype=torch.bfloat16, - ) - for i in range(34): - kv_cache_spec[f"layers.{i}.kda"] = MambaSpec( - shapes=((3, 6144), (16, 128, 128)), - dtypes=(torch.bfloat16, torch.float32), - block_size=block_size, - ) - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - attn = [ - g - for g in groups - if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) - and isinstance(g.kv_cache_spec.first_spec, MLAAttentionSpec) - ] - rings = [ - g - for g in groups - if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) - and isinstance(g.kv_cache_spec.first_spec, CircularBufferSpec) - ] - mamba = [ - g - for g in groups - if isinstance( - getattr(g.kv_cache_spec, "first_spec", g.kv_cache_spec), MambaSpec - ) - ] - assert len(attn) == 1 and len(attn[0].layer_names) == 22 - assert len(rings) == 1 and len(rings[0].layer_names) == 11 - # 11 * (1.18 MB MLA + 34 KB indexer) per block fits 12 fp32-SSM KDA states. - assert len(mamba) == 3 - assert sum(len(g.layer_names) for g in mamba) == 34 - - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, available_memory=2 * 1024**3 - ) - strides = {t.block_stride for t in kv_cache_config.kv_cache_tensors} - assert len(strides) == 1 - (block_stride,) = strides - # Blocks sit a whole number of 32-row MLA pages and 64-state indexer pages - # apart so both caches can be re-paged for their kernels. - assert block_stride % (32 * 576 * 2) == 0 - assert block_stride % (64 * 132) == 0 - - -def new_mla_spec(cache_dtype_str=None, block_size: int = 16): - # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576 - return MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - max_tp_shards=1, - head_size=576, - dtype=torch.float32, - cache_dtype_str=cache_dtype_str, - ) - - -def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): - return SlidingWindowMLASpec( - block_size=16, - num_kv_heads=1, - max_tp_shards=1, - head_size=head_size, - dtype=torch.float32, - sliding_window=sliding_window, - model_version=model_version, - ) - - -def new_indexer_mla_spec(block_size=16): - # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller - # page size than the main MLA attention (uint8, small head), so their pages - # cannot be unified. - return MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=132, - dtype=torch.uint8, - ) - - -def test_mixed_page_size_groups_use_spec_compatibility(): - specs = {} - for i in range(3): - specs[f"mla.{i}"] = new_mla_spec() - specs[f"indexer.{i}"] = new_indexer_mla_spec() - specs.update({f"swa.{i}": new_swa_mla_spec(head_size=1024) for i in range(5)}) - - config = _grouping_config() - config.cache_config = CacheConfig() - config.cache_config.kv_cache_layout = "BLNHC" - groups = get_kv_cache_groups(config, specs) - - assert len(groups) == 3 - assert {name for group in groups for name in group.layer_names} == set(specs) - assert sorted(len(group.layer_names) for group in groups) == [2, 3, 6] - - -def _grouping_config(): - cache_config = CacheConfig() - cache_config.kv_cache_layout = "LBNHC" - return SimpleNamespace( - scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), - speculative_config=None, - cache_config=cache_config, - ) - - -def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): - block_size = 544 - full_spec = FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - # FullAttentionSpec stores both K and V, so this produces a - # 544 * 1 * (512 + 512) * 2 = 1,114,112-byte page. - head_size=512, - dtype=torch.bfloat16, - ) - mamba_spec = MambaSpec( - block_size=block_size, - shapes=((557056,),), - dtypes=(torch.bfloat16,), - mamba_cache_mode="align", - ) - hidden_spec = HiddenStateCacheSpec( - block_size=block_size, - num_kv_heads=3, - head_size=1024, - dtype=torch.bfloat16, - ) - assert full_spec.page_size_bytes == mamba_spec.page_size_bytes - - groups = get_kv_cache_groups( - _grouping_config(), - { - "model.full_attn": full_spec, - "model.mamba": mamba_spec, - "cache_only_layers.0": hidden_spec, - }, - ) - - hidden_group = next( - group - for group in groups - if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) - ) - assert hidden_group.kv_cache_spec.block_size == 136 - - kv_cache_config = KVCacheConfig( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=groups, - ) - vllm_config = SimpleNamespace( - cache_config=SimpleNamespace( - block_size=16, - enable_prefix_caching=True, - prefix_match_unit=None, - ), - parallel_config=SimpleNamespace(decode_context_parallel_size=1), - kv_transfer_config=object(), - ) - assert kv_cache_utils.resolve_kv_cache_block_sizes( - kv_cache_config, vllm_config - ) == (544, 136) - - def test_resolve_dcp_kv_block_size_unwraps_uniform_type_specs(): attention = FullAttentionSpec( block_size=16, @@ -3133,125 +1715,6 @@ def test_resolve_dcp_kv_block_size_unwraps_uniform_type_specs(): assert kv_cache_utils.resolve_dcp_kv_cache_spec(wrapped_mamba, 4) is wrapped_mamba -def test_multi_run_layer_compact_strides_place_hoisted_heads(): - """A layer-compact run region is its own dense allocation: under LHBNC the head - groups sit between the layers and the blocks, so a run's block stride is one head - group's slice, not the whole page (regression test for setStorage out-of-bounds - on ROCm hybrid models).""" - full = new_kv_cache_spec() - swa = new_sliding_window_spec(sliding_window=full.block_size * 2) - page = full.page_size_bytes - assert swa.page_size_bytes == page - - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) - vllm_config.cache_config.kv_cache_layout = "LHBNC" - config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, - [ - KVCacheGroupSpec(["full.0", "full.1"], full), - KVCacheGroupSpec(["swa.0", "swa.1"], swa), - ], - available_memory=4 * page * 8, - ) - - num_blocks = config.num_blocks - head_group = full.block_size * full.state_content_size_bytes - for tensor in config.kv_cache_tensors: - # One head group per block, head groups between the layers and blocks. - assert tensor.block_stride == head_group - assert tensor.layer_stride == full.num_heads * num_blocks * head_group - # The last block of the last layer's last head group must stay within the - # allocation (the old page-based stride overran it). - end = ( - tensor.offset - + tensor.layer_stride * (len(tensor.layers) - 1) - + (full.num_heads - 1) * num_blocks * head_group - + tensor.block_stride * (num_blocks - 1) - + head_group - ) - assert end <= tensor.size - - -def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): - specs = { - "target.0.attn": new_mla_spec(), - "draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288), - } - assert len({spec.page_size_bytes for spec in specs.values()}) == 1 - - groups = get_kv_cache_groups(_grouping_config(), specs) - - assert len(groups) == 2 - assert all( - not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups - ) - - -def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): - # Sparse MLA pages cannot be padded safely. Keeping the draft's attention - # compute sliding-window while promoting only its allocation semantics lets - # every layer share the target's block table and remain contiguous. - draft = new_sliding_window_spec(block_size=16) - specs = { - "target.0.attn": new_mla_spec(block_size=64), - "target.0.indexer": new_indexer_mla_spec(block_size=64), - "draft.0": draft, - } - - groups = get_kv_cache_groups(_grouping_config(), specs) - assert len(groups) == 1 - assert set(groups[0].layer_names) == set(specs) - group_spec = groups[0].kv_cache_spec - assert isinstance(group_spec, UniformTypeKVCacheSpecs) - assert group_spec.block_size == 64 - promoted_draft = group_spec.kv_cache_specs["draft.0"] - assert isinstance(promoted_draft, FullAttentionSpec) - assert not isinstance(promoted_draft, SlidingWindowSpec) - assert promoted_draft.block_size == 64 - assert promoted_draft.sliding_window == draft.sliding_window - assert specs["draft.0"] is draft - assert "attention compute is unchanged" in caplog_vllm.text - - -def test_hidden_states_with_tp_scales_page_size(): - """When TP shrinks KV pages below the hidden-state per-token cost, - get_kv_cache_groups must scale up target block sizes so that the - common page accommodates the unsharded hidden states.""" - # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. - # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) - # = 16 * 2 * 64 * 2 = 4096 bytes. - kv_spec = new_kv_cache_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.bfloat16, - ) - # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512) - # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page. - hs_spec = HiddenStateCacheSpec( - block_size=16, - num_kv_heads=6, - head_size=512, - dtype=torch.bfloat16, - ) - specs = { - "target.0.attn": kv_spec, - "target.1.attn": kv_spec, - "cache_only_layers.48": hs_spec, - } - - groups = get_kv_cache_groups(_grouping_config(), specs) - - # The hidden-state layer should be present and no assertion should fire. - all_layers = {name for g in groups for name in g.layer_names} - assert "cache_only_layers.48" in all_layers - - # The target group block sizes must have been scaled up. - for g in groups: - if "cache_only_layers.48" not in g.layer_names: - assert g.kv_cache_spec.block_size > kv_spec.block_size - - def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses(): assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION @@ -3525,179 +1988,6 @@ def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes assert block_hashes[1] == expected_hash2 -def test_auto_fit_max_model_len(): - """Test that max_model_len=-1 auto-fits to available GPU memory.""" - # Create config with original_max_model_len=-1 to trigger auto-fit - model_config = ModelConfig(max_model_len=1024) - # Simulate the user passing -1 by setting original_max_model_len - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - - # With enough memory, max_model_len stays at the derived max - large_available_memory = mem_per_block_per_layer * 2 * 1024 # plenty of memory - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [large_available_memory] - ) - assert vllm_config.model_config.max_model_len == 1024 - - # Reset for next test - model_config = ModelConfig(max_model_len=1024) - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - # With limited memory, max_model_len should be reduced - # Need memory for at least max_model_len tokens - # 32 blocks worth of memory for 2 layers = can fit 32*16=512 tokens - limited_memory = mem_per_block_per_layer * 2 * 32 - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [limited_memory] - ) - # Should be reduced to fit in memory - assert vllm_config.model_config.max_model_len < 1024 - assert vllm_config.model_config.max_model_len > 0 - - -def test_auto_fit_max_model_len_with_hybrid(): - """Test that auto-fit works with hybrid KV cache specs.""" - # Create config with original_max_model_len=-1 to trigger auto-fit - model_config = ModelConfig(max_model_len=8192) - # Simulate the user passing -1 by setting original_max_model_len - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer - gamma = 2 - kv_cache_specs = { - "layer_1": new_mamba_spec(num_speculative_blocks=gamma), - "layer_2": new_kv_cache_spec(), - } - - # One extra block on top of what a 1024-token request needs: the pool - # reserves one block as the null block. - available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1) - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [available_memory] - ) - assert vllm_config.model_config.max_model_len == 1024 - - -def test_auto_fit_max_model_len_not_triggered(): - """Test that auto-fit is not triggered when original_max_model_len is not -1.""" - model_config = ModelConfig(max_model_len=16) - # original_max_model_len should be None by default, not -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - - # This should work normally without auto-fit - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32] - ) - assert vllm_config.model_config.max_model_len == 16 - - -def test_auto_fit_max_model_len_respects_num_gpu_blocks_override(): - """Auto-fit must size max_model_len against the override-clamped pool, not - the raw `available_memory`. Without this, auto-fit could pick a - max_model_len that no longer fits once `num_gpu_blocks_override` is applied. - """ - model_config = ModelConfig(max_model_len=16384) - model_config.original_max_model_len = -1 # request auto-fit - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - # Cap the cache to 32 blocks regardless of available memory. - vllm_config.cache_config.num_gpu_blocks_override = 32 - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), # block_size=16 - "layer_2": new_kv_cache_spec(), - } - # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384). - large_available_memory = mem_per_block_per_layer * 2 * 1024 - - get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) - - # 32 blocks * block_size 16 = 512 token slots, so max_model_len must - # auto-fit at or below that. - assert 0 < vllm_config.model_config.max_model_len <= 32 * 16 - - -def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override(): - """Admission check must use the override-clamped pool size, not raw - `available_memory`. Without this, startup could accept a max_model_len - that does not actually fit in `num_gpu_blocks_override` blocks. - """ - model_config = ModelConfig(max_model_len=16384) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - # 32 blocks is far too small for max_model_len=16384 (would need 1024). - vllm_config.cache_config.num_gpu_blocks_override = 32 - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - # Plenty of raw memory: a bytes-only check against this would pass. - large_available_memory = mem_per_block_per_layer * 2 * 1024 - - with pytest.raises(ValueError, match="max seq len"): - get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) - - -def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): - """DFlash drafters can have a smaller head size than the target model. - - For example, MiMo uses 192-dim target KV heads while its DFlash draft uses - 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer - block-size multiple, so the smaller page must be padded instead. - """ - # Attention layers read padded pages through a strided view, so padding is - # allowed. - target_spec = new_kv_cache_spec( - block_size=16, - num_kv_heads=1, - head_size=192, - dtype=torch.bfloat16, - ) - draft_spec = new_sliding_window_spec( - block_size=16, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - sliding_window=1024, - ) - - unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "target_attn": target_spec, - "draft_attn": draft_spec, - } - ) - - assert unified_specs["target_attn"] == target_spec - unified_draft_spec = unified_specs["draft_attn"] - assert unified_draft_spec.block_size == draft_spec.block_size - assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes - assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes - assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes - - def test_unpadded_page_size_includes_per_token_head_scales(): # Per-token-head quant carries inline fp32 scales that are carved from the # raw KV allocation, so they must be budgeted into the offload width. The @@ -3729,173 +2019,6 @@ def test_page_size_padded_wins(): assert spec.page_size_bytes == 65536 -def test_unify_hybrid_kv_cache_specs(): - # 1. has_full_attention and has_sliding_window - before_spec_1 = new_kv_cache_spec(block_size=64) - before_spec_2 = new_sliding_window_spec( - block_size=16, page_size_padded=32 * 1024, sliding_window=1024 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - } - kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False) - kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False) - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec(block_size=64) - expected_spec_2 = new_kv_cache_spec( - block_size=64, page_size_padded=64 * 1024, sliding_window=1024 - ) - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024 - assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False) - - # 2. has_full_attention and has_chunked_local_attention - before_spec_1 = new_kv_cache_spec() - before_spec_2 = new_chunked_local_attention_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - } - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec() - expected_spec_2 = new_kv_cache_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - - # 3. has_full_attention, has_sliding_window and has_chunked_local_attention - before_spec_1 = new_kv_cache_spec() - before_spec_2 = new_sliding_window_spec( - page_size_padded=32 * 1024, sliding_window=1024 - ) - before_spec_3 = new_chunked_local_attention_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - "layer_3": before_spec_3, - } - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec() - expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024) - expected_spec_3 = new_kv_cache_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - assert kv_cache_spec["layer_3"] == expected_spec_3 - - # 4. No FullAttentionSpec, should not convert - kv_cache_spec = { - "layer_1": new_sliding_window_spec(sliding_window=1024), - "layer_2": new_chunked_local_attention_spec(attention_chunk_size=512), - } - - with pytest.raises(ValueError): - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - - # Replicated Mamba state still requires the hybrid cache manager. - kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()} - with pytest.raises(ValueError): - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - - -def test_unify_kv_cache_spec_page_size_mamba(): - """Regression test for https://github.com/vllm-project/vllm/issues/43626. - - MambaSpec's page_size_bytes is determined by its state shapes and does not - change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba - page instead of scaling its block_size. This situation arises when a layer - with a page larger than the (already platform-aligned) Mamba page joins the - specs, e.g. a dense draft model with more KV heads than the hybrid main - model. - """ - # 1. Hybrid main model (Mamba + full attention, pages already aligned at - # 16KB) plus a dense draft model layer with a 2x larger page (32KB). - # Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so - # the old code scaled the Mamba block_size, which left page_size_bytes - # unchanged at 16384. - mamba_spec = new_mamba_spec() # page_size_bytes = 16384 - main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384 - draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768 - assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384 - assert draft_attn_spec.page_size_bytes == 32768 - - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": mamba_spec, - "main_attn_layer": main_attn_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - # Mamba page is padded; block_size (caching granularity) is unchanged. - assert unified["mamba_layer"].page_size_bytes == 32768 - assert unified["mamba_layer"].page_size_padded == 32768 - assert unified["mamba_layer"].block_size == mamba_spec.block_size - # Attention layer with smaller page still unifies by scaling block_size. - assert unified["main_attn_layer"].page_size_bytes == 32768 - assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size - assert unified["main_attn_layer"].page_size_padded is None - # Layer already at max page size is unchanged. - assert unified["draft_attn_layer"] == draft_attn_spec - - # 2. Mamba page already padded by the platform (state smaller than the - # padded page); the padding is re-applied at the new maximum. - padded_mamba_spec = new_mamba_spec( - shapes=((2, 256), (3, 32, 32)), page_size_padded=16384 - ) - assert padded_mamba_spec.page_size_bytes == 16384 - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": padded_mamba_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - assert unified["mamba_layer"].page_size_bytes == 32768 - assert unified["mamba_layer"].page_size_padded == 32768 - - # 3. Mamba page that does not evenly divide the maximum page size is - # padded as well (the divisibility constraint only applies to block_size - # scaling). - odd_mamba_spec = new_mamba_spec(shapes=((6144,),)) - assert odd_mamba_spec.page_size_bytes == 24576 - assert 32768 % odd_mamba_spec.page_size_bytes != 0 - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": odd_mamba_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - assert unified["mamba_layer"].page_size_bytes == 32768 - - # 4. Attention layers with non-divisible page sizes are padded too: every - # backend reads a padded page through the view's block stride, so there is - # no longer a case that must raise. - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "attn_layer": new_kv_cache_spec(block_size=24), # 24576 - "draft_attn_layer": draft_attn_spec, # 32768 - } - ) - assert unified["attn_layer"].page_size_padded == 32768 - assert unified["attn_layer"].page_size_bytes == 32768 - assert unified["attn_layer"].block_size == 24 - - # 5. Uniform page sizes are returned unchanged. - specs = { - "mamba_layer": new_mamba_spec(), - "attn_layer": new_kv_cache_spec(), - } - assert kv_cache_utils.unify_kv_cache_spec_page_size(specs) == specs - - def test_hma_not_disabled_when_kv_events_enabled(): """Test enabling KV events must not force disable_hybrid_kv_cache_manager to True. @@ -3982,78 +2105,6 @@ def test_resolve_block_hashes_rejects_mismatched_view(): resolve_block_hashes(mismatched, 2, 4) -@pytest.mark.parametrize("use_override", [True, False]) -def test_kv_cache_reserves_null_block_for_max_model_len(use_override): - """A KV cache sized to exactly the blocks a max_model_len request needs must - be rejected. BlockPool keeps one block as the null block, so only - num_blocks - 1 are usable; accepting num_blocks == needed would leave the - request unschedulable and hang the engine. Covers both the - num_gpu_blocks_override path and the memory-derived block count. - """ - block_size = 16 - max_model_len = 512 # needs 512 / 16 = 32 blocks - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - spec = new_kv_cache_spec(block_size=block_size) - - # 32 blocks -> only 31 usable after the null block: one short -> reject. - if use_override: - # Ample raw memory; the override alone constrains the block count. - vllm_config.cache_config.num_gpu_blocks_override = 32 - available_memory = [spec.page_size_bytes * 1024] - else: - available_memory = [spec.page_size_bytes * 32] - - with pytest.raises(ValueError, match="max seq len"): - get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) - - -def test_auto_fit_max_model_len_reserves_null_block(): - """Auto-fit (max_model_len=-1) must size max_model_len against usable - blocks, not the total pool. With memory for exactly 64 blocks, one is the - null block, so auto-fit must settle on 63 * block_size; picking the full - pool would leave a max-length request unschedulable and hang the engine. - """ - block_size = 16 - model_config = ModelConfig(max_model_len=1024) - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - spec = new_kv_cache_spec(block_size=block_size) - - # Exactly the 1024 / 16 = 64 blocks a full-length request would need. - available_memory = [spec.page_size_bytes * 64] - - get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) - - assert vllm_config.model_config.max_model_len == 63 * block_size - - -def test_check_enough_kv_cache_memory_reserves_null_block(): - """The public admission check must reject a KV cache sized to exactly the - blocks a max_model_len request needs. BlockPool keeps one block as the - null block, so only num_blocks - 1 are usable; accepting - num_blocks == needed would leave the request unschedulable and hang the - engine. - """ - block_size = 16 - max_model_len = 512 # needs 512 / 16 = 32 blocks - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - spec = new_kv_cache_spec(block_size=block_size) - - # 32 blocks -> only 31 usable after the null block: one short -> reject. - with pytest.raises(ValueError, match="max seq len"): - check_enough_kv_cache_memory( - vllm_config, {"layer1": spec}, spec.page_size_bytes * 32 - ) - - # 33 blocks -> 32 usable after the null block -> accept. - check_enough_kv_cache_memory( - vllm_config, {"layer1": spec}, spec.page_size_bytes * 33 - ) - - def test_is_full_attention_spec_unwraps_uniform_type_specs(): """``UniformTypeKVCacheSpecs`` is not a ``FullAttentionSpec``, so callers scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs.""" @@ -4105,340 +2156,3 @@ def test_wrapped_mamba_group_requires_block_zeroing(): ) assert config.needs_kv_cache_zeroing - - -def _spec_decode_grouping_config(method="dspark", model_type=None): - """Grouping config with an EAGLE-family speculative method enabled.""" - return SimpleNamespace( - scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), - cache_config=SimpleNamespace( - get_resolved_kv_cache_layout=lambda: SimpleNamespace( - is_block_outermost=True - ) - ), - model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)), - speculative_config=SimpleNamespace( - method=method, - use_eagle=lambda: True, - use_eagle_block_drop=lambda: True, - ), - ) - - -def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False): - """A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a - DSpark-style draft MLA layer marked non_causal_multi_token_decode. - - The target's fp8 KV dtype is one difference from the draft on Kimi-K3 - (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass - draft_shares_target_spec to match that dtype while retaining the draft - marker as a group boundary. - """ - target_dtype = None if draft_shares_target_spec else "fp8_e4m3" - specs = { - "target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), - "target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), - "target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"), - "target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"), - } - if draft: - draft_spec = new_mla_spec(block_size=64) - specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True) - return specs - - -def test_draft_group_annotated_on_hybrid_general_path(): - # A drafter's MLA layer carries non_causal_multi_token_decode, so its group - # is identifiable without keying off a model version. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "draft.attn.0" in flagged[0].layer_names - - -def test_mamba_groups_never_flagged_even_when_draft_shares_a_group(): - # Packed uniform-type groups can contain distinct target and draft layer - # specs; the combined group still holds volatile draft KV and must be - # flagged. What must never happen is a Mamba group being flagged: that - # widens its lookup window to two consecutive chunks, which align-mode - # checkpointing never produces, zeroing every lookup. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), - _hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True), - ) - - for group in groups: - if "draft.attn.0" in group.layer_names: - assert group.is_eagle_group - if isinstance(group.kv_cache_spec, MambaSpec): - assert not group.is_eagle_group - - -def test_draft_group_not_annotated_without_spec_decode(): - # The marker alone must not flag anything; the eagle semantics only apply - # when a speculative method is actually enabled. - config = _spec_decode_grouping_config() - config.speculative_config = None - groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True)) - - assert not any(g.is_eagle_group for g in groups) - - -def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): - # No group carries the draft marker, so consumers fall back to - # conservative behavior that silently breaks reuse for Mamba groups. - # That must at least be visible. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) - ) - - assert not any(g.is_eagle_group for g in groups) - assert "no KV cache group could be identified as the draft model's" in ( - caplog_vllm.text - ) - - -def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): - # Pure-attention models degrade gracefully under the consumers' - # conservative fallback (a one-block hit drop at most), so the warning - # stays silent to avoid noise on every unannotated EAGLE deployment. - specs = { - "target.attn.0": new_mla_spec(block_size=64), - "target.attn.1": new_mla_spec(block_size=64), - } - groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs) - - assert not any(g.is_eagle_group for g in groups) - assert "could be identified as the draft model's" not in caplog_vllm.text - - -def test_no_warning_when_draft_group_is_identified(caplog_vllm): - get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) - ) - - assert "could be identified as the draft model's" not in caplog_vllm.text - - -def _deepseek_v4_specs(model_version="deepseek_v4"): - """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers - at differing page sizes, with the MTP draft layer registered last.""" - return { - "model.layers.0.self_attn.attn": new_mla_spec(), - "model.layers.1.self_attn.attn": new_mla_spec(), - "model.layers.2.self_attn.attn": new_swa_mla_spec( - head_size=1024, model_version=model_version - ), - # The MTP block registers last, and its sliding-window size differs, so - # it lands in a group of its own. - "model.layers.3.self_attn.attn": new_swa_mla_spec( - head_size=1024, sliding_window=256, model_version=model_version - ), - } - - -@pytest.mark.parametrize( - ("method", "model_type"), - [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")], -) -def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type): - # DeepseekV4's MTP block reuses the target's decoder layer, so its spec - # carries no draft marker and only the positional rule can find it. This - # pins the pre-existing behaviour that the unified annotator must preserve. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method=method, model_type=model_type), - _deepseek_v4_specs(model_version=None), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "model.layers.3.self_attn.attn" in flagged[0].layer_names - - -def test_trailing_layer_fallback_applies_to_any_mtp_model(): - # The positional rule is sound for every MTP drafter, not just DeepseekV4: - # MTP blocks reuse the target's decoder layer (no spec marker) and always - # register after every target layer. The model_type must not gate it. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method="mtp", model_type="other"), - _deepseek_v4_specs(), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "model.layers.3.self_attn.attn" in flagged[0].layer_names - - -def _qwen3_5_hybrid_specs(with_mtp_layer: bool): - """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with - the MTP drafter's full-attn layer (spec-identical to the target's) - registered last.""" - specs = {} - idx = 0 - for _ in range(2): - for _ in range(3): - specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec( - mamba_cache_mode="align" - ) - idx += 1 - specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec() - idx += 1 - if with_mtp_layer: - specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec() - return specs - - -def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm): - # A hybrid mamba + full-attention model with an MTP drafter that is - # spec-indistinguishable from the target reaches the general multi-group - # path. The trailing-layer rule must locate the draft group there so the - # Mamba groups are not swept up by the flag-all consumer fallback. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"), - _qwen3_5_hybrid_specs(with_mtp_layer=True), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names - for group in groups: - if any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ): - assert not group.is_eagle_group - assert "could be identified as the draft model's" not in caplog_vllm.text - - -@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"]) -def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method): - # Other drafters are not covered by the trailing-layer rule, so an - # unidentifiable hybrid draft must still warn. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method=method, model_type="qwen3_5"), - _qwen3_5_hybrid_specs(with_mtp_layer=True), - ) - - assert not any(g.is_eagle_group for g in groups) - assert "could be identified as the draft model's" in caplog_vllm.text - - -def test_trailing_layer_fallback_requires_exact_partition(): - # If the groups do not partition the layers exactly (e.g. a caller that - # dropped or duplicated layers), the positional rule is meaningless and - # must not fire. - from vllm.v1.core.kv_cache_utils import _annotate_eagle_groups - - specs = _qwen3_5_hybrid_specs(with_mtp_layer=True) - config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5") - groups = get_kv_cache_groups(config, specs) - for g in groups: - g.is_eagle_group = False - # Remove one layer from its group: no longer an exact partition. - trimmed = [ - KVCacheGroupSpec( - [n for n in g.layer_names if n != "model.layers.0.linear_attn"], - g.kv_cache_spec, - ) - for g in groups - ] - _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) - - assert not any(g.is_eagle_group for g in trimmed) - - -_GQA_SPEC = FullAttentionSpec( - block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32 -) -_GQA_SWA_SPEC = SlidingWindowSpec( - block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128 -) - - -@pytest.mark.parametrize( - "specs,expected", - [ - pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"), - pytest.param( - { - "l.0": new_mla_spec(), - "l.1": new_mla_spec(), - "s.0": new_swa_mla_spec(), - "s.1": new_swa_mla_spec(), - }, - 4, - id="mla-and-swa-mla", - ), - pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"), - pytest.param( - {"l.0": new_mla_spec(), "l.1": _GQA_SPEC}, - 1, - id="mla-and-gqa-uniform-group", - ), - pytest.param( - { - "l.0": new_mla_spec(), - "l.1": new_mla_spec(), - "s.0": _GQA_SWA_SPEC, - "s.1": _GQA_SWA_SPEC, - }, - 1, - id="mla-and-gqa-swa", - ), - ], -) -def test_kv_tp_replicas(monkeypatch, specs, expected): - """Resolved per layer before scheduler flattening hides mixed groups.""" - from vllm.config import ParallelConfig - from vllm.platforms import current_platform - - monkeypatch.setattr(current_platform, "device_count", lambda: 4) - parallel_config = ParallelConfig(tensor_parallel_size=4) - vllm_config = VllmConfig( - model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config - ) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - mem = sum(s.page_size_bytes for s in specs.values()) * 10 - configs = get_kv_cache_configs(vllm_config, [specs], [mem]) - assert configs[0].kv_tp_replicas == expected - scheduler = generate_scheduler_kv_cache_config(configs) - assert scheduler.kv_tp_replicas == expected - - -@pytest.mark.parametrize( - "specs,tp_size,dcp_size,expected", - [ - pytest.param([new_mla_spec()], 8, 1, 8, id="mla"), - pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"), - pytest.param( - [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset" - ), - pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"), - pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"), - pytest.param( - [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)], - 8, - 1, - 2, - id="mla-and-gqa-partial", - ), - pytest.param( - [ - HiddenStateCacheSpec( - block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32 - ) - ], - 8, - 1, - 1, - id="hidden-state", - ), - ], -) -def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected): - """Replicas are the gcd of each layer's tp_size // max_tp_shards.""" - groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)] - assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index cf221fc22a..541e55f5b6 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -36,6 +36,10 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, + KVQuantMode, + MambaSpec, + MLAAttentionSpec, + SlidingWindowSpec, ) from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -333,3 +337,73 @@ def create_requests( ) requests.append(request) return requests + + +def new_kv_cache_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + sliding_window=None, + attention_chunk_size=None, + kv_quant_mode=KVQuantMode.NONE, +): + return FullAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + sliding_window=sliding_window, + attention_chunk_size=attention_chunk_size, + kv_quant_mode=kv_quant_mode, + ) + + +def new_sliding_window_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + sliding_window=1, +): + return SlidingWindowSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + sliding_window=sliding_window, + ) + + +def new_mamba_spec( + block_size=16, + shapes=((2, 512), (3, 32, 32)), + dtypes=(torch.float32, torch.float32), + num_speculative_blocks=2, + mamba_cache_mode="none", + page_size_padded=None, +): + return MambaSpec( + block_size=block_size, + shapes=shapes, + dtypes=dtypes, + page_size_padded=page_size_padded, + mamba_cache_mode=mamba_cache_mode, + num_speculative_blocks=num_speculative_blocks, + ) + + +def new_mla_spec(cache_dtype_str=None, block_size: int = 16): + # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576 + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + max_tp_shards=1, + head_size=576, + dtype=torch.float32, + cache_dtype_str=cache_dtype_str, + ) diff --git a/tests/v1/engine/test_init_error_messaging.py b/tests/v1/engine/test_init_error_messaging.py index 5626dec170..f9455389ef 100644 --- a/tests/v1/engine/test_init_error_messaging.py +++ b/tests/v1/engine/test_init_error_messaging.py @@ -4,7 +4,7 @@ import pytest import torch -from vllm.v1.core.kv_cache_utils import check_enough_kv_cache_memory +from vllm.v1.core.kv_cache_planning import check_enough_kv_cache_memory from vllm.v1.kv_cache_interface import FullAttentionSpec @@ -38,7 +38,7 @@ def test_kv_cache_oom_insufficient_memory(monkeypatch): config.parallel_config.decode_context_parallel_size = 1 monkeypatch.setattr( - "vllm.v1.core.kv_cache_utils.max_memory_usage_bytes", + "vllm.v1.core.kv_cache_planning.max_memory_usage_bytes", lambda c, s: 100 * 1024**3, # 100 GiB ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 07d530f97d..a6692b5cd5 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -17,10 +17,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( SchedulerOffloadConfig, ) from vllm.platforms import current_platform -from vllm.v1.core.kv_cache_utils import ( - generate_scheduler_kv_cache_config, - kv_cache_groups_tp_replicas, -) +from vllm.v1.core.kv_cache_planning import kv_cache_groups_tp_replicas +from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, diff --git a/tests/v1/kv_connector/unit/test_hidden_states_connector.py b/tests/v1/kv_connector/unit/test_hidden_states_connector.py index 735b042259..854632dce1 100644 --- a/tests/v1/kv_connector/unit/test_hidden_states_connector.py +++ b/tests/v1/kv_connector/unit/test_hidden_states_connector.py @@ -10,7 +10,7 @@ import torch from vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector import ( # noqa: E501 ExampleHiddenStatesConnector, ) -from vllm.v1.core.kv_cache_utils import get_kv_cache_groups +from vllm.v1.core.kv_cache_planning import get_kv_cache_groups from vllm.v1.kv_cache_interface import ( FullAttentionSpec, HiddenStateCacheSpec, diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py index 16e142d7cb..4f9a7c451f 100644 --- a/tests/v1/simple_kv_offload/test_worker.py +++ b/tests/v1/simple_kv_offload/test_worker.py @@ -29,11 +29,11 @@ from vllm.config import CacheConfig from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import ( SimpleCPUOffloadConnector, ) -from vllm.v1.core.kv_cache_utils import ( +from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, is_kv_cache_spec_uniform, - resolve_kv_cache_block_sizes, ) +from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, diff --git a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py index 4cc8f9ffe2..abd0830a47 100644 --- a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py +++ b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py @@ -104,7 +104,7 @@ def test_overlaid_zeroer_dedups_segments_with_max_span(): zeroed no matter which group owns it.""" from unittest.mock import MagicMock - from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups + from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import ( KVCacheGroupSpec, KVCacheLayout, diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 1247bf193a..534d813124 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -43,7 +43,10 @@ from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( ROCMAiterMLASparseBackend, ) from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs +from vllm.v1.core.kv_cache_planning import ( + estimate_max_model_len, + get_kv_cache_configs, +) from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput from vllm.v1.kv_cache_interface import ( FullAttentionSpec, diff --git a/vllm/v1/core/kv_cache_planning.py b/vllm/v1/core/kv_cache_planning.py new file mode 100644 index 0000000000..6e258a68a9 --- /dev/null +++ b/vllm/v1/core/kv_cache_planning.py @@ -0,0 +1,1637 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""KV cache planning: grouping, sizing, and memory-fit logic.""" + +import math +from collections import defaultdict +from collections.abc import Callable, Iterable, Sequence +from dataclasses import replace +from functools import partial, reduce +from typing import Any + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv, round_up +from vllm.utils.mem_utils import format_gib +from vllm.utils.torch_utils import get_dtype_size +from vllm.v1.core.kv_cache_utils import _kv_cache_device_name +from vllm.v1.hisparse.layout import ( + get_hisparse_host_pool_bytes, + get_hisparse_kv_cache_config, + get_hisparse_kv_cache_groups, +) +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + ChunkedLocalAttentionSpec, + CircularBufferSpec, + FullAttentionSpec, + HiddenStateCacheSpec, + HiSparseHotSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheLayout, + KVCacheSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, + UniformTypeKVCacheSpecs, + compute_layout_strides, + iter_layer_specs, + replace_as, +) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry + +logger = init_logger(__name__) + + +def _check_enough_kv_cache_memory( + available_memory: int, + get_needed_memory: Callable[[], int], + max_model_len: int, + estimate_max_model_len: Callable[[int], int], +): + if available_memory <= 0: + raise ValueError( + "No available memory for the cache blocks. " + "Try increasing `gpu_memory_utilization` when initializing the engine " + "(this flag also controls CPU memory reservation on the CPU " + "backend, despite its name). " + "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " + "for more details." + ) + + needed_memory = get_needed_memory() + + if needed_memory > available_memory: + estimated_max_len = estimate_max_model_len(available_memory) + estimated_msg = "" + if estimated_max_len > 0: + estimated_msg = ( + "Based on the available memory, " + f"the estimated maximum model length is {estimated_max_len}. " + ) + + raise ValueError( + f"To serve at least one request with the model's max seq len " + f"({max_model_len}), ({format_gib(needed_memory)} GiB KV " + f"cache is needed, which is larger than the available KV cache " + f"memory ({format_gib(available_memory)} GiB). {estimated_msg}" + f"Try increasing `gpu_memory_utilization` (which also controls " + f"CPU memory on the CPU backend) or decreasing `max_model_len` " + f"when initializing the engine. " + f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " + f"for more details." + ) + + +def max_memory_usage_bytes( + vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec] +) -> int: + """Get the maximum memory usage in bytes for the given KV cache specs.""" + return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs) + + +def estimate_max_model_len( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + available_memory: int, +) -> int: + """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. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + available_memory: Memory available for KV cache in bytes. + + Returns: + The estimated maximum model length that can fit in the available memory. + + """ + # Save the original max_model_len to restore after estimation + original_max_model_len = vllm_config.model_config.max_model_len + + # Define a function to check if a given model length fits in memory + def fits_in_memory(model_len: int) -> bool: + # Temporarily modify the max_model_len for this calculation + vllm_config.model_config.max_model_len = model_len + # Calculate memory needed for the given model length + memory_needed = max_memory_usage_bytes(vllm_config, kv_cache_spec.values()) + return memory_needed <= available_memory + + try: + # Binary search for the maximum model length + left, right = 1, original_max_model_len + + # If even the smallest model length doesn't fit, return 0 + if not fits_in_memory(left): + return 0 + + # Binary search for the maximum model length that fits + result = 1 + while left <= right: + mid = (left + right) // 2 + if fits_in_memory(mid): + result = mid + left = mid + 1 + else: + right = mid - 1 + return result + finally: + # Always restore the original max_model_len to avoid side effects + vllm_config.model_config.max_model_len = original_max_model_len + + +def check_enough_kv_cache_memory( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + available_memory: int, +): + """Checks whether `available_memory` is enough for the KV cache to hold at + least one request with the model's max_model_len. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + available_memory: Memory available for KV cache in bytes. + + Raises: + ValueError: If there is not enough memory available for the KV cache. + + """ + # No need to check for available memory if the kv_cache_spec is empty + if kv_cache_spec: + # Reserve the null block BlockPool permanently holds back, so the check + # plans against usable blocks, as in get_kv_cache_configs. Group a copy + # of the specs since grouping may unify them in-place. + groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec)) + check_memory = ( + available_memory - _pool_bytes_per_block(vllm_config, groups) + if groups + else available_memory + ) + _check_enough_kv_cache_memory( + check_memory, + lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), + vllm_config.model_config.max_model_len, + lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am), + ) + + +def create_kv_cache_group_specs( + kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]] +) -> list[KVCacheGroupSpec]: + """Create KVCacheGroupSpec object for each kv cache group layer. + The layers in the same group should share the same + KVCacheSpec. + + Args: + kv_cache_spec: + A mapping from each layer name to its corresponding KVCacheSpec. + grouped_layer_names: + 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: + A list of KVCacheGroupSpec objects, one for each group. + + """ + kv_cache_groups = [] + for layer_names_one_group in grouped_layer_names: + layer_specs = [ + kv_cache_spec[layer_name] for layer_name in layer_names_one_group + ] + merged_layer_spec = layer_specs[0].merge(layer_specs) + kv_cache_groups.append( + KVCacheGroupSpec(layer_names_one_group, merged_layer_spec) + ) + return kv_cache_groups + + +def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: + """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. + + Args: + kv_cache_spec: The kv cache spec of each attention layer in the model + + Returns: + True if all layers have the same type, False otherwise. + + """ + if not kv_cache_spec: + # Encoder-only models do not have KV cache, kv_cache_type can be + # regarded as uniform. + return True + try: + kv_cache_spec_values = list(kv_cache_spec.values()) + _ = kv_cache_spec_values[0].merge(kv_cache_spec_values) + except AssertionError: + return False + return True + + +def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: + """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`. + """ + if vllm_config.cache_config.num_gpu_blocks_override is not None: + num_blocks = vllm_config.cache_config.num_gpu_blocks_override + return num_blocks + + +def _pool_bytes_per_block( + vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec] +) -> int: + """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. + """ + return _get_kv_cache_bytes_per_block( + [group for group in kv_cache_groups if not group.host_resident], + vllm_config.cache_config.get_resolved_kv_cache_layout(), + ) + + +def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: + """Get the page size of the KV cache.""" + page_sizes = {layer.page_size_bytes for layer in kv_cache_specs} + assert len(page_sizes) == 1 + return page_sizes.pop() + + +def _get_kv_cache_groups_uniform_spec( + kv_cache_specs: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Generates the KV cache configuration for a model with the same KV cache + spec for all layers. + + Args: + kv_cache_specs: The kv cache spec of each attention layer in the model + + Returns: + The generated KVCacheGroupSpecs + + """ + return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())]) + + +def _get_kv_cache_groups_uniform_type( + spec: UniformTypeKVCacheSpecs, +) -> list[KVCacheGroupSpec]: + """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. + + Args: + spec: The UniformTypeKVCacheSpecs of the model + + Returns: + The generated KVCacheGroupSpecs + + """ + return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] + + +def unify_kv_cache_spec_page_size( + kv_cache_spec: dict[str, KVCacheSpec], +) -> dict[str, KVCacheSpec]: + """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). + + Args: + kv_cache_spec: The KVCacheSpec of each attention layer in the model + + Returns: + The updated KVCacheSpec with the same page_size_bytes. + + """ + page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()} + if len(page_sizes) <= 1: + # All layers have the same page size, no need to unify. + return kv_cache_spec + + max_page_size = max(page_sizes) + new_kv_cache_spec = {} + for layer_name, layer_spec in kv_cache_spec.items(): + if layer_spec.page_size_bytes == max_page_size: + new_kv_cache_spec[layer_name] = layer_spec + elif isinstance(layer_spec, MambaSpec): + # MambaSpec's page size is determined by its state shapes and does + # not scale with block_size, so pad the page instead. This is the + # same padding mechanism the platform uses to align Mamba pages + # with the main model's attention page size; it is needed here + # when another layer (e.g. from a draft model) has a larger page + # than the already-aligned Mamba page. + new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size) + assert new_spec.page_size_bytes == max_page_size + new_kv_cache_spec[layer_name] = new_spec + else: + layer_page_size = layer_spec.page_size_bytes + if max_page_size % layer_page_size == 0: + ratio = max_page_size // layer_page_size + new_block_size = layer_spec.block_size * ratio + new_spec = replace(layer_spec, block_size=new_block_size) + elif isinstance(layer_spec, AttentionSpec) and not isinstance( + layer_spec, MLAAttentionSpec + ): + new_spec = replace(layer_spec, page_size_padded=max_page_size) + else: + raise NotImplementedError( + f"Layer {layer_name}: page size is not divisible by the " + "maximum page size and cannot be padded. Padding is only " + "supported for non-MLA attention layers." + ) + assert new_spec.page_size_bytes == max_page_size + new_kv_cache_spec[layer_name] = new_spec + return new_kv_cache_spec + + +def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: + # kv_cache_spec is an empty dict for attention free models + return not kv_cache_spec + + +def _get_kv_cache_groups_uniform_page_size( + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """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. + + Args: + kv_cache_spec: The KVCacheSpec of each attention layer in the model + Returns: + The generated KVCacheGroupSpecs + + """ + # Group all layers by kv_cache_spec. + # E.g., 2 full attention layers and 3 sliding window attention layers, + # -> (full.0, full.1), (sw.0, sw.1, sw.2). + same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list) + for layer_name, layer_spec in kv_cache_spec.items(): + same_type_layers[layer_spec].append(layer_name) + + # Attempt to further merge same-type layers based on whether their KV + # cache specs can be merged, to minimize the group count. This benefits + # situations where specs share a block layout and differ only in a + # property it can reconcile (e.g. full attention layers differing only in + # sliding window / attention chunk size). + layer_buckets: list[list[str]] = [] + spec_buckets: list[list[KVCacheSpec]] = [] + for layer_spec, layer_names in same_type_layers.items(): + for names, specs in zip(layer_buckets, spec_buckets): + try: + # A raise means that the specs are incompatible. + type(specs[0]).merge([*specs, layer_spec]) + except (AssertionError, ValueError): + continue + names.extend(layer_names) + specs.append(layer_spec) + break + else: + layer_buckets.append(list(layer_names)) + spec_buckets.append([layer_spec]) + + # Split each group into smaller groups, to make the number of layers in each + # group identical. Add padding to the last group of each type if necessary. + # E.g., (full.0, full.1), (sw.0, sw.1, sw.2) + # split to 3 groups with 2 layers each: + # (full.0, full.1), (sw.0, sw.2), (sw.1, padding). + # FIXME(Chen): At the moment of writing this code (2025-06-02), all + # open-source hybrid model follows a n:1 pattern between different attention + # types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and + # full), so we can use the "1" in the n:1 pattern as the group size, which + # is the minimum number of layers among all attention types. Need a better + # strategy if we want to support more complex patterns (e.g., 20 full + 30 + # sw, where the group size should be 10). + min_num_layers = min([len(layers) for layers in layer_buckets]) + group_size = min_num_layers + max_num_layers = max([len(layers) for layers in layer_buckets]) + if max_num_layers < min_num_layers * 1.5: + # If the number of layers is not much larger than the minimum number of + # layers, use the maximum number of layers as the group size to avoid + # too many padding layers. A typical example is gpt-oss-20b + eagle, + # with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of + # (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding + # layers while accommodating speculative decoding drafters that add + # extra layers to one attention type. + group_size = max_num_layers + grouped_layers = [] + for layers in layer_buckets: + num_padding_layers = group_size - len(layers) % group_size + if num_padding_layers != group_size: + logger.warning( + "Add %d padding layers, may waste at most %.2f%% KV cache memory", # noqa + num_padding_layers, + num_padding_layers / len(layers) * 100, + ) + num_groups = cdiv(len(layers), group_size) + # In PP case, say if we have + # - stage 0: full.0, sw.0, sw.1 + # - stage 1: full.1, sw.2, sw.3 + # We should have 3 groups: (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3) + # It can't be (full.0, full.1), (sw.0, sw.1), (sw.2, sw.3) because + # the 3 groups in stage 0 will be (full.0), (sw.0, sw.1), (empty group) + # and it will be padded to (full.0, padding), (sw.0, sw.1), + # (padding, padding) to ensure the number of layers in each group is + # the same and will cause memory waste. + # To avoid this, we assign layers[i::num_groups] to the i-th group + # instead of layers[i * group_size: (i + 1) * group_size] + for i in range(num_groups): + grouped_layers.append(layers[i::num_groups]) + return create_kv_cache_group_specs(kv_cache_spec, grouped_layers) + + +def _get_per_layer_spec( + group: KVCacheGroupSpec, + layer_name: str, +) -> KVCacheSpec: + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + return spec.kv_cache_specs[layer_name] + return spec + + +def _get_layer_block_stride(spec: KVCacheSpec) -> int: + """A layer-outer block stride: the page padded to its alignment.""" + return round_up(spec.page_size_bytes, spec.get_block_stride_alignment()) + + +def _get_kv_cache_bytes_per_block( + kv_cache_groups: list[KVCacheGroupSpec], layout: KVCacheLayout +) -> int: + """Return the largest cache group's bytes per block.""" + if not layout.is_block_outermost: + # Each layer's block stride is its own padded page. + return max( + sum( + _get_layer_block_stride(_get_per_layer_spec(group, layer_name)) + for layer_name in group.layer_names + ) + for group in kv_cache_groups + ) + bytes_per_block = max( + sum( + _get_per_layer_spec(group, layer_name).page_size_bytes + for layer_name in group.layer_names + ) + for group in kv_cache_groups + ) + assert bytes_per_block > 0 + alignments = [ + group.kv_cache_spec.page_size_bytes + for group in kv_cache_groups + if isinstance(group.kv_cache_spec, HiSparseHotSpec) + ] + alignments.extend( + _get_per_layer_spec(group, layer_name).get_block_stride_alignment() + for group in kv_cache_groups + for layer_name in group.layer_names + ) + return round_up(bytes_per_block, math.lcm(*alignments)) + + +def validate_kv_cache_layout( + layout: KVCacheLayout, + kv_cache_groups: list[KVCacheGroupSpec], +) -> None: + """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. + """ + page_sizes = { + _get_per_layer_spec(group, layer_name).page_size_bytes + for group in kv_cache_groups + for layer_name in group.layer_names + } + if len(page_sizes) == 1: + # A rectangular layer dim exists; every layout can express it. + return + + # Mixed page sizes pack pages side by side within a block, which needs each page + # to be one contiguous chunk inside its block (a block-compact layout) and, with + # multiple KV cache groups, the layer dim inside the block dim. + if not layout.is_block_compact or ( + len(kv_cache_groups) > 1 and layout.is_layer_compact + ): + raise ValueError( + f"KV cache layout {layout.name} cannot express this model's " + f"mixed page sizes ({sorted(page_sizes)}); a backend should " + "declare block-outermost supported layouts (e.g. BLHNC), or " + "set VLLM_KV_CACHE_LAYOUT=BLHNC." + ) + + +def get_kv_cache_config_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, +) -> KVCacheConfig: + """Generate the KV cache configuration from the KV cache groups and spec + of each layer. + + Args: + vllm_config: The global VllmConfig + kv_cache_groups: The KV cache groups + available_memory: Memory available for KV cache in bytes + Returns: + The generated KVCacheConfig + + """ + if len(kv_cache_groups) == 0: + # Attention free models do not have KV cache. + # Return num_blocks=1 as BlockPool always needs a null_block. + return KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + ) + + if vllm_config.attention_config.hisparse_config is not None: + host_budget = get_hisparse_host_pool_bytes(vllm_config) + return get_hisparse_kv_cache_config( + vllm_config, kv_cache_groups, available_memory, host_budget + ) + + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + validate_kv_cache_layout(layout, kv_cache_groups) + bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups, layout) + interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None + + num_blocks = available_memory // bytes_per_block + num_blocks = may_override_num_blocks(vllm_config, num_blocks) + size = bytes_per_block * num_blocks + + # Groups alias from byte 0. Spec regions are laid out differently: + # + # block-outer (the same packing repeats for every block): + # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ... + # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ... + # |<--- bytes_per_block -->| + # + # layer-outer (only supported for uniform page sizes or single-group models): + # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] | + # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] | + + kv_cache_tensors = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) + if isinstance(group_spec, UniformTypeKVCacheSpecs): + for layer_name, spec in group_spec.kv_cache_specs.items(): + layers_by_spec[spec].append(layer_name) + elif group.layer_names: + layers_by_spec[group_spec].extend(group.layer_names) + + byte_offset = 0 + for spec, layer_names in layers_by_spec.items(): + block_stride = interleaved_block_stride + padded_page = _get_layer_block_stride(spec) + if block_stride is None and padded_page != spec.page_size_bytes: + assert layout.is_block_compact, f"{layout.name} cannot pad pages" + block_stride = padded_page + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, + num_blocks, + len(layer_names), + layout, + fixed_strides=(None, block_stride, None, None, None), + ) + kv_cache_tensors.append( + KVCacheTensor( + size=size, + layers=layer_names, + layer_stride=layer_stride, + block_stride=block_stride, + offset=byte_offset, + ) + ) + # Block-outer layers share each block; layer-outer ones own regions. + byte_offset += len(layer_names) * ( + spec.page_size_bytes if layout.is_block_outermost else layer_stride + ) + + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + ) + + +def _promote_local_kv_cache_specs( + kv_cache_spec: dict[str, KVCacheSpec], +) -> dict[str, KVCacheSpec]: + """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. + """ + promoted_specs = kv_cache_spec.copy() + + if is_kv_cache_spec_uniform( + promoted_specs + ) or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs): + return promoted_specs + + has_full_attention = any( + isinstance(spec, FullAttentionSpec) for spec in promoted_specs.values() + ) + has_sliding_window = any( + isinstance(spec, SlidingWindowSpec) for spec in promoted_specs.values() + ) + has_chunked_local_attention = any( + isinstance(spec, ChunkedLocalAttentionSpec) for spec in promoted_specs.values() + ) + full_block_sizes = { + spec.block_size + for spec in promoted_specs.values() + if isinstance(spec, FullAttentionSpec) + } + full_attention_block_size = ( + next(iter(full_block_sizes)) if len(full_block_sizes) == 1 else None + ) + + def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | None: + if spec.page_size_padded is None: + return None + unpadded_page_size = ( + spec.unpadded_page_size_bytes * block_size // spec.block_size + ) + return max(spec.page_size_padded, unpadded_page_size) + + promotions: dict[type[AttentionSpec], type[AttentionSpec]] = { + SlidingWindowMLASpec: MLAAttentionSpec, + SlidingWindowSpec: FullAttentionSpec, + ChunkedLocalAttentionSpec: FullAttentionSpec, + } + + if has_full_attention and (has_sliding_window or has_chunked_local_attention): + for layer_name, spec in kv_cache_spec.items(): + target_cls = next( + (promotions[c] for c in type(spec).__mro__ if c in promotions), None + ) + if target_cls is None: + continue + assert isinstance(spec, AttentionSpec) + block_size = full_attention_block_size or spec.block_size + promoted_specs[layer_name] = replace_as( + spec, + target_cls, + # Promoted specs allocate blocks for all tokens and never free + # below the window, so the trailing-edge extension is moot. + drop=("extra_retained_tokens",), + block_size=block_size, + page_size_padded=promoted_page_size_padded(spec, block_size), + ) + + if not ( + is_kv_cache_spec_uniform(promoted_specs) + or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs) + ): + raise ValueError("Failed to promote local KV cache specs to one unified type.") + + return promoted_specs + + +def _try_get_full_allocation_fallback_groups( + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Try a supported full-allocation fallback for local-attention layers.""" + if any(isinstance(spec, HiddenStateCacheSpec) for spec in kv_cache_spec.values()): + return None + if any( + isinstance(spec, (SlidingWindowMLASpec, ChunkedLocalAttentionSpec)) + for spec in kv_cache_spec.values() + ): + return None + + has_mla = any(isinstance(spec, MLAAttentionSpec) for spec in kv_cache_spec.values()) + has_regular_swa = any( + isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values() + ) + if not (has_mla and has_regular_swa): + return None + + try: + promoted_specs = _promote_local_kv_cache_specs(kv_cache_spec) + except ValueError: + return None + uniform_spec = UniformTypeKVCacheSpecs.from_specs(promoted_specs) + if uniform_spec is None: + return None + logger.warning( + "KV cache page sizes cannot be unified; treating sliding-window " + "layers as full attention for cache allocation. Sliding-window " + "attention compute is unchanged." + ) + return _get_kv_cache_groups_uniform_type(uniform_spec) + + +def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]): + """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. + + Args: + kv_cache_spec: The kv cache spec of each attention layer in the model + + """ + groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict) + for name, spec in kv_cache_spec.items(): + replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded + groups[replicated][name] = spec + for specs in groups.values(): + promoted_specs = _promote_local_kv_cache_specs(specs) + if promoted_specs == specs: + continue + logger.warning( + "Hybrid KV cache manager is disabled for this hybrid model, " + "This means we do not enable any optimizations for saving KV cache " + "memory (e.g., dropping the KV cache outside the sliding window). " + "The compute of layers like sliding window is still saved." + ) + kv_cache_spec.update(promoted_specs) + + +def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: + """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. + """ + if not values: + raise ValueError("values must be non-empty") + if any(x <= 0 for x in values): + raise ValueError(f"values must be positive, got: {list(values)!r}") + + min_d = max(1, lower_bound if lower_bound is not None else 1) + max_d = max(values) + if min_d > max_d: + return min_d + + best_d = min_d + best_pad: int | None = None + for d in range(min_d, max_d + 1): + pad = sum((d - (x % d)) % d for x in values) + if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d): + best_pad = pad + best_d = d + + return best_d + + +def _get_packed_kv_cache_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """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. + """ + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()} + if not layout.is_block_outermost or len(page_sizes) <= 1: + return None + + buckets: list[dict[str, KVCacheSpec]] = [] + for name, spec in kv_cache_spec.items(): + for bucket in buckets: + candidate = {**bucket, name: spec} + if UniformTypeKVCacheSpecs.is_uniform_type(candidate): + bucket[name] = spec + break + else: + buckets.append({name: spec}) + + bucketed = [] + for bucket in buckets: + uniform_spec = UniformTypeKVCacheSpecs.from_specs(bucket) + assert uniform_spec is not None + page_size_layers: dict[int, list[str]] = defaultdict(list) + for layer_name, layer_spec in bucket.items(): + page_size_layers[layer_spec.page_size_bytes].append(layer_name) + # Only 1:1 patterns (one layer of each page size per repeat) are + # supported; counts sharing a gcd > 1 (e.g. 2:1) could in principle + # repeat too, but such buckets are emitted whole instead. + balanced = len(set(map(len, page_size_layers.values()))) == 1 + bucketed.append((uniform_spec, page_size_layers, balanced)) + + # Balanced buckets that mix page sizes must stay whole, so the largest one + # sets a floor on the repeats per group; larger single-size buckets are + # split down toward it. No such bucket means nothing needs packing. + min_repeats_per_group = max( + ( + spec.get_max_layers_per_page_size() + for spec, page_size_layers, balanced in bucketed + if balanced and len(page_size_layers) > 1 + ), + default=0, + ) + repeats_per_group = ( + _approximate_gcd( + [ + spec.get_max_layers_per_page_size() + for spec, _, balanced in bucketed + if balanced + ], + lower_bound=min_repeats_per_group, + ) + if min_repeats_per_group + else None + ) + + def num_groups_for(spec: UniformTypeKVCacheSpecs, balanced: bool) -> int: + if balanced and repeats_per_group is not None: + return cdiv(spec.get_max_layers_per_page_size(), repeats_per_group) + return 1 + + def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: + """Page bytes of the largest of the n groups a bucket splits into.""" + return sum( + cdiv(len(names), n) * page for page, names in page_size_layers.items() + ) + + # Bytes a block must hold however the state buckets end up split: mamba, + # circular-buffer and unsplit sliding-window buckets can go down to one + # state per group, every other bucket's split is fixed by the repeat pattern. + def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: + if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)): + return True + return repeats_per_group is None and isinstance( + spec.first_spec, SlidingWindowSpec + ) + + anchor_bytes = max( + ( + widest_group_bytes( + page_size_layers, + len(spec.kv_cache_specs) + if is_state_bucket(spec) + else num_groups_for(spec, balanced), + ) + for spec, page_size_layers, balanced in bucketed + ), + default=0, + ) + + groups = [] + for spec, page_size_layers, balanced in bucketed: + num_groups = num_groups_for(spec, balanced) + # Cap a state group at the states a block already fits rather than let + # it widen the block. + if anchor_bytes and is_state_bucket(spec): + states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1) + num_groups = max( + num_groups, cdiv(len(spec.kv_cache_specs), states_per_block) + ) + if num_groups == 1: + groups.append(KVCacheGroupSpec(list(spec.kv_cache_specs), spec)) + continue + + pattern_repeats = list(zip(*page_size_layers.values())) + for i in range(num_groups): + group_layer_names = [ + name for repeat in pattern_repeats[i::num_groups] for name in repeat + ] + group_layer_specs = { + name: spec.kv_cache_specs[name] for name in group_layer_names + } + group_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs) + assert group_spec is not None + groups.append(KVCacheGroupSpec(group_layer_names, group_spec)) + + _annotate_eagle_groups( + vllm_config, + kv_cache_spec, + groups, + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), + ) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) + return groups + + +def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool: + """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. + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle(): + return False + if spec_config.method == "mtp": + return True + model_config = vllm_config.model_config + return model_config is not None and model_config.hf_config.model_type in ( + "deepseek_v4", + "deepseek_v41", + ) + + +def _groups_partition_layers_exactly( + kv_cache_spec: dict[str, KVCacheSpec], + kv_cache_groups: list[KVCacheGroupSpec], +) -> bool: + """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. + """ + names = [name for group in kv_cache_groups for name in group.layer_names] + return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec) + + +def _annotate_eagle_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + kv_cache_groups: list[KVCacheGroupSpec], + use_trailing_layer_fallback: bool = False, +) -> None: + """Flag the KV cache groups that hold drafter attention layers. + + Two detection rules, in order of preference: + + 1. Spec-driven. ``non_causal_multi_token_decode`` is 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. + 2. 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. + + Args: + vllm_config: Config supplying the speculative method, if any. + kv_cache_spec: The kv cache spec of each attention layer, in layer + registration order. Only read by rule 2. + kv_cache_groups: Groups to annotate in place. + use_trailing_layer_fallback: Enable rule 2. Callers gate this on + ``_uses_trailing_mtp_layers``. + + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle_block_drop(): + return + + for group in kv_cache_groups: + if any( + getattr(spec, "non_causal_multi_token_decode", False) + for spec in iter_layer_specs(group.kv_cache_spec) + ): + group.is_eagle_group = True + + if not use_trailing_layer_fallback: + return + if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups): + return + last_layer = next(reversed(kv_cache_spec)) + for group in kv_cache_groups: + if last_layer in group.layer_names: + group.is_eagle_group = True + break + + +def _warn_if_unannotated_eagle_mamba( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> None: + """Warn when no KV cache group could be identified as the draft model's. + + Args: + vllm_config: Config supplying the speculative method, if any. + kv_cache_groups: Groups as they will be handed to consumers. + + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle(): + return + if any(group.is_eagle_group for group in kv_cache_groups): + return + mamba_groups = [ + idx + for idx, group in enumerate(kv_cache_groups) + if any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ) + ] + if not mamba_groups: + return + logger.warning( + "Speculative decoding (method=%s) is enabled but no KV cache group " + "could be identified as the draft model's.", + spec_config.method, + ) + + +def _largest_divisor_at_most(value: int, limit: int) -> int: + for candidate in range(min(value, limit), 0, -1): + if value % candidate == 0: + return candidate + return 1 + + +def _ensure_min_page_size( + groups: list[KVCacheGroupSpec], + common_page: int, + hidden_specs: dict[str, HiddenStateCacheSpec], +) -> tuple[list[KVCacheGroupSpec], int]: + """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. + """ + max_per_token = max( + spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) + for spec in hidden_specs.values() + ) + if max_per_token <= common_page: + return groups, common_page + + scale = cdiv(max_per_token, common_page) + common_page *= scale + scaled: list[KVCacheGroupSpec] = [] + for g in groups: + s = g.kv_cache_spec + kw: dict[str, Any] = {"block_size": s.block_size * scale} + if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None: + kw["page_size_padded"] = s.page_size_padded * scale + scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw))) + return scaled, common_page + + +def get_kv_cache_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Split the layers in the model into groups with the same KV cache spec. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + + Returns: + The generated KVCacheGroups + + """ + if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: + unify_hybrid_kv_cache_specs(kv_cache_spec) + + if is_kv_cache_type_attention_free(kv_cache_spec): + # This returns an empty list to allow for the KVCacheManager to handle + # attention free models. + return [] + + if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec): + return hisparse_groups + + if is_kv_cache_spec_uniform(kv_cache_spec): + # KV cache of all layers are the same, which is true for + # most models. Allocate the same amount of memory for + # each layer. + return _get_kv_cache_groups_uniform_spec(kv_cache_spec) + elif uniform_spec := UniformTypeKVCacheSpecs.from_specs(kv_cache_spec): + # All layers need the same number of token slots (e.g., all layers are + # full attention, or all layers are sliding window attention with the + # same window size). Put all layers into one group. + return _get_kv_cache_groups_uniform_type(uniform_spec) + # Hidden-state layers use their own block table and must not be absorbed + # into a compatible attention bucket. + hidden_specs = { + k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec) + } + filtered_spec = { + k: v + for k, v in kv_cache_spec.items() + if not isinstance(v, HiddenStateCacheSpec) + } + + if packed_groups := _get_packed_kv_cache_groups(vllm_config, filtered_spec): + # Block-outermost blocks are strided by the widest group, so hidden + # groups need no page alignment. + packed_groups += [ + KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items() + ] + return packed_groups + + # Prefer preserving each layer's cache semantics. If physical pages cannot + # be unified, try a supported allocation-only fallback before failing. + try: + filtered_spec = unify_kv_cache_spec_page_size(filtered_spec) + except NotImplementedError: + fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec) + if fallback_groups is None: + raise + return fallback_groups + groups = _get_kv_cache_groups_uniform_page_size(filtered_spec) + + # Add hidden-state layers back with page aligned to the common page. + if hidden_specs: + common_page = get_uniform_page_size([g.kv_cache_spec for g in groups]) + # TP may shrink the common page below hidden-state per-token cost. + groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs) + group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups)) + for name, spec in hidden_specs.items(): + per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) + max_block_size = max(common_page // per_token, 1) + new_bs = _largest_divisor_at_most(group_block_size, max_block_size) + wasted_bytes = common_page - new_bs * per_token + logger.info( + "Using block size %d for hidden-state cache layer %s; " + "page alignment wastes %d bytes (%.2f%%) per block", + new_bs, + name, + wasted_bytes, + wasted_bytes / common_page * 100, + ) + aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) + groups.append(KVCacheGroupSpec([name], aligned)) + + _annotate_eagle_groups( + vllm_config, + kv_cache_spec, + groups, + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), + ) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) + return groups + + +def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int: + if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None: + return 1 + if spec.dcp_sharded and dcp_size > 1: + return 1 + return max(1, tp_size // spec.max_tp_shards) + + +def kv_cache_groups_tp_replicas( + groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1 +) -> int: + """Consecutive TP ranks holding identical KV for every layer.""" + specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)] + if not specs: + return 1 + return reduce( + math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size + ) + + +def _max_memory_usage_bytes_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> int: + """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``. + """ + kv_cache_groups = [group for group in kv_cache_groups if not group.host_resident] + if not kv_cache_groups: + return 0 + + bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups) + total_blocks = 0 + for group in kv_cache_groups: + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + total_blocks += spec.max_memory_usage_pages(vllm_config) + else: + total_blocks += cdiv( + spec.max_memory_usage_bytes(vllm_config), + spec.page_size_bytes, + ) + + return bytes_per_block * total_blocks + + +def _estimate_max_model_len_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, +) -> int: + """Binary search for the maximum model length that fits in available memory. + Returns 0 if even 1 token doesn't fit. + """ + original_max = vllm_config.model_config.max_model_len + hisparse_enabled = ( + vllm_config.attention_config.hisparse_config is not None + and bool(kv_cache_groups) + ) + + def fits(model_len: int) -> bool: + vllm_config.model_config.max_model_len = model_len + if hisparse_enabled: + try: + get_kv_cache_config_from_groups( + vllm_config, kv_cache_groups, available_memory + ) + except ValueError: + return False + return ( + _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups) + <= available_memory + ) + + try: + left, right = 1, original_max + if not fits(left): + return 0 + result = 1 + while left <= right: + mid = (left + right) // 2 + if fits(mid): + result = mid + left = mid + 1 + else: + right = mid - 1 + return result + finally: + vllm_config.model_config.max_model_len = original_max + + +def _auto_fit_max_model_len( + vllm_config: VllmConfig, + projected_groups_per_worker: list[list[KVCacheGroupSpec]], + available_memory: list[int], +) -> None: + """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. + + Args: + vllm_config: The global VllmConfig (will be modified in-place) + projected_groups_per_worker: KV cache groups projected to each worker. + available_memory: Memory available for KV cache in bytes for each + worker. + + """ + original_max = vllm_config.model_config.max_model_len + + if all(not groups for groups in projected_groups_per_worker): + # All workers have empty specs (attention-free model) + logger.info_once( + "Auto-fit max_model_len: attention-free model, " + "using derived max_model_len=%d", + original_max, + ) + return + + # Find the max_model_len that fits across all workers. + auto_fit_max = original_max + limiting_worker_mem = available_memory[0] + for groups, avail_mem in zip(projected_groups_per_worker, available_memory): + if not groups: + continue + worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem) + if worker_max < auto_fit_max: + auto_fit_max = worker_max + limiting_worker_mem = avail_mem + + if auto_fit_max <= 0: + raise ValueError( + "Cannot auto-fit max_model_len: not enough GPU memory available " + "to serve even a single token. Try increasing `gpu_memory_utilization`." + ) + + kv_cache_memory = ( + f"{format_gib(limiting_worker_mem)} GiB {_kv_cache_device_name(vllm_config)}" + ) + if vllm_config.attention_config.hisparse_config is not None: + host_pool_gib = format_gib(get_hisparse_host_pool_bytes(vllm_config)) + kv_cache_memory += f" + {host_pool_gib} GiB host" + + if auto_fit_max >= original_max: + # The model's full context length fits in memory + logger.info_once( + "Auto-fit max_model_len: full model context length %d fits in " + "available KV cache memory (%s)", + original_max, + kv_cache_memory, + ) + else: + # Need to reduce max_model_len to fit in memory + vllm_config.model_config.max_model_len = auto_fit_max + logger.info_once( + "Auto-fit max_model_len: reduced from %d to %d to fit in " + "available KV cache memory (%s)", + original_max, + auto_fit_max, + kv_cache_memory, + ) + + +def _project_kv_cache_groups_to_worker( + global_kv_cache_groups: list[KVCacheGroupSpec], + worker_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """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. + + Args: + global_kv_cache_groups: The global KV cache groups for the whole model. + worker_spec: The KV cache spec of each layer on this worker. + + Returns: + The projected KV cache groups containing only this worker's layers. + + """ + projected_groups: list[KVCacheGroupSpec] = [] + for group in global_kv_cache_groups: + worker_layer_names = [ + layer_name for layer_name in group.layer_names if layer_name in worker_spec + ] + if len(worker_layer_names) == len(group.layer_names): + projected_groups.append(group) + continue + group_spec = group.kv_cache_spec + if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs): + group_spec = UniformTypeKVCacheSpecs( + block_size=group_spec.block_size, + kv_cache_specs={ + layer_name: group_spec.kv_cache_specs[layer_name] + for layer_name in worker_layer_names + }, + ) + projected_groups.append( + KVCacheGroupSpec( + worker_layer_names, + group_spec, + is_eagle_group=group.is_eagle_group and bool(worker_layer_names), + ) + ) + return projected_groups + + +def get_kv_cache_configs( + vllm_config: VllmConfig, + kv_cache_specs: list[dict[str, KVCacheSpec]], + available_memory: list[int], +) -> list[KVCacheConfig]: + """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. + + Args: + vllm_config: The global VllmConfig + kv_cache_specs: List of dict[layer_name, KVCacheSpec] for each worker. + available_memory: Memory available for KV cache in bytes for each + worker. + + Returns: + The generated KVCacheConfigs for each worker. + + """ + # Merge the KV cache specs of all workers. Different PP stages may have + # different layer names, and different TP ranks of the same PP stage should + # have the same KV cache spec. + merged_kv_cache_specs: dict[str, KVCacheSpec] = {} + for kv_cache_spec_one_worker in kv_cache_specs: + for layer_name, layer_spec in kv_cache_spec_one_worker.items(): + if layer_name not in merged_kv_cache_specs: + merged_kv_cache_specs[layer_name] = layer_spec + else: + assert merged_kv_cache_specs[layer_name] == layer_spec, ( + "The KV cache specs for the same layer are different " + "across workers. This is not supported yet." + ) + + # Check if the KV cache specs are registered correctly. + # This is to prevent that some layers are initialized with unregistered specs. + KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) + + # When speculating with more than 1 speculative module (e.g. multi-layered MTP) + # tag every SlidingWindowSpec with how many extra tokens to retain in the window. + extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1) + for layer_name, layer_spec in merged_kv_cache_specs.items(): + if isinstance(layer_spec, SlidingWindowSpec): + merged_kv_cache_specs[layer_name] = replace( + layer_spec, extra_retained_tokens=extra_retained_tokens + ) + + # Get global KV cache groups. This also handles spec unification for + # hybrid models when disable_hybrid_kv_cache_manager is enabled. + # After this call, merged_kv_cache_specs may be modified in-place. + global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs) + + # If original_max_model_len was -1, automatically + # determine the maximum model length that fits in available GPU memory. + # We use per-worker projected groups to account for PP sharding. + projected_groups_per_worker = [ + _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec) + for worker_spec in kv_cache_specs + ] + + # If `num_gpu_blocks_override` is set, the cache size that will actually + # be allocated is decoupled from the profiled `available_memory`: + # `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps + # `num_blocks` to the override. Reflect that in `available_memory` here so + # auto-fit, the admission check, and the per-worker config builder all + # plan against the same effective capacity. + override = vllm_config.cache_config.num_gpu_blocks_override + if override is not None: + adjusted_memory: list[int] = [] + for groups, avail_mem in zip(projected_groups_per_worker, available_memory): + if not groups: + adjusted_memory.append(avail_mem) + continue + bytes_per_block = _pool_bytes_per_block(vllm_config, groups) + logger.info( + "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d", + avail_mem // bytes_per_block, + override, + ) + adjusted_memory.append(override * bytes_per_block) + available_memory = adjusted_memory + + # Reserve the null block BlockPool permanently holds back, so auto-fit and + # the capacity check both plan against usable blocks. Allocation below + # still uses the full memory. + check_memory = [ + avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem + for groups, avail_mem in zip(projected_groups_per_worker, available_memory) + ] + + if vllm_config.model_config.original_max_model_len == -1: + _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory) + + # Check if the available memory is enough per worker. + for groups, avail_mem in zip(projected_groups_per_worker, check_memory): + if not groups: + continue + _check_enough_kv_cache_memory( + avail_mem, + partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), + vllm_config.model_config.max_model_len, + partial(_estimate_max_model_len_from_groups, vllm_config, groups), + ) + + kv_cache_configs: list[KVCacheConfig] = [] + for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip( + projected_groups_per_worker, kv_cache_specs, available_memory + ): + assert sum(len(group.layer_names) for group in projected_groups) == len( + kv_cache_spec_one_worker + ), "Some layers are not assigned to any group." + kv_cache_configs.append( + get_kv_cache_config_from_groups( + vllm_config, projected_groups, available_memory_one_worker + ) + ) + + # Change the num_blocks of each rank to the smallest among all ranks. + # We also need to shrink the tensor size proportionally to avoid + # allocating unused memory. + min_num_blocks = min( + kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs + ) + for i, kv_cache_config in enumerate(kv_cache_configs): + if kv_cache_config.num_blocks == min_num_blocks: + continue + # Re-plan with exactly the memory the smallest rank can afford, so + # strides and offsets stay consistent with the shrunken allocation. + groups = kv_cache_config.kv_cache_groups + kv_cache_configs[i] = get_kv_cache_config_from_groups( + vllm_config, + groups, + min_num_blocks * _pool_bytes_per_block(vllm_config, groups), + ) + + for kv_cache_config in kv_cache_configs: + kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas( + kv_cache_config.kv_cache_groups, + vllm_config.parallel_config.tensor_parallel_size, + vllm_config.parallel_config.decode_context_parallel_size, + ) + + return kv_cache_configs diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 03c6c5ff62..5c51a8d131 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -6,10 +6,8 @@ import copy import hashlib import math import os -from collections import defaultdict from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, field, replace -from functools import partial, reduce from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, overload from vllm import envs @@ -17,34 +15,14 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.multimodal.utils import get_mm_features_in_window from vllm.utils.hashing import xxhash, xxhash_cbor -from vllm.utils.math_utils import cdiv, round_up -from vllm.utils.mem_utils import format_gib -from vllm.utils.torch_utils import get_dtype_size -from vllm.v1.hisparse.layout import ( - get_hisparse_host_pool_bytes, - get_hisparse_kv_cache_config, - get_hisparse_kv_cache_groups, -) +from vllm.utils.math_utils import cdiv from vllm.v1.kv_cache_interface import ( - AttentionSpec, - ChunkedLocalAttentionSpec, - CircularBufferSpec, - FullAttentionSpec, - HiddenStateCacheSpec, - HiSparseHotSpec, KVCacheConfig, KVCacheGroupSpec, - KVCacheLayout, KVCacheSpec, - KVCacheTensor, MambaSpec, - MLAAttentionSpec, - SlidingWindowMLASpec, - SlidingWindowSpec, UniformTypeKVCacheSpecs, - compute_layout_strides, iter_layer_specs, - replace_as, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.request import Request @@ -1000,198 +978,6 @@ def get_request_block_hasher( return request_block_hasher -def _check_enough_kv_cache_memory( - available_memory: int, - get_needed_memory: Callable[[], int], - max_model_len: int, - estimate_max_model_len: Callable[[int], int], -): - if available_memory <= 0: - raise ValueError( - "No available memory for the cache blocks. " - "Try increasing `gpu_memory_utilization` when initializing the engine " - "(this flag also controls CPU memory reservation on the CPU " - "backend, despite its name). " - "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " - "for more details." - ) - - needed_memory = get_needed_memory() - - if needed_memory > available_memory: - estimated_max_len = estimate_max_model_len(available_memory) - estimated_msg = "" - if estimated_max_len > 0: - estimated_msg = ( - "Based on the available memory, " - f"the estimated maximum model length is {estimated_max_len}. " - ) - - raise ValueError( - f"To serve at least one request with the model's max seq len " - f"({max_model_len}), ({format_gib(needed_memory)} GiB KV " - f"cache is needed, which is larger than the available KV cache " - f"memory ({format_gib(available_memory)} GiB). {estimated_msg}" - f"Try increasing `gpu_memory_utilization` (which also controls " - f"CPU memory on the CPU backend) or decreasing `max_model_len` " - f"when initializing the engine. " - f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " - f"for more details." - ) - - -def max_memory_usage_bytes( - vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec] -) -> int: - """Get the maximum memory usage in bytes for the given KV cache specs.""" - return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs) - - -def estimate_max_model_len( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - available_memory: int, -) -> int: - """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. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - available_memory: Memory available for KV cache in bytes. - - Returns: - The estimated maximum model length that can fit in the available memory. - - """ - # Save the original max_model_len to restore after estimation - original_max_model_len = vllm_config.model_config.max_model_len - - # Define a function to check if a given model length fits in memory - def fits_in_memory(model_len: int) -> bool: - # Temporarily modify the max_model_len for this calculation - vllm_config.model_config.max_model_len = model_len - # Calculate memory needed for the given model length - memory_needed = max_memory_usage_bytes(vllm_config, kv_cache_spec.values()) - return memory_needed <= available_memory - - try: - # Binary search for the maximum model length - left, right = 1, original_max_model_len - - # If even the smallest model length doesn't fit, return 0 - if not fits_in_memory(left): - return 0 - - # Binary search for the maximum model length that fits - result = 1 - while left <= right: - mid = (left + right) // 2 - if fits_in_memory(mid): - result = mid - left = mid + 1 - else: - right = mid - 1 - return result - finally: - # Always restore the original max_model_len to avoid side effects - vllm_config.model_config.max_model_len = original_max_model_len - - -def check_enough_kv_cache_memory( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - available_memory: int, -): - """Checks whether `available_memory` is enough for the KV cache to hold at - least one request with the model's max_model_len. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - available_memory: Memory available for KV cache in bytes. - - Raises: - ValueError: If there is not enough memory available for the KV cache. - - """ - # No need to check for available memory if the kv_cache_spec is empty - if kv_cache_spec: - # Reserve the null block BlockPool permanently holds back, so the check - # plans against usable blocks, as in get_kv_cache_configs. Group a copy - # of the specs since grouping may unify them in-place. - groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec)) - check_memory = ( - available_memory - _pool_bytes_per_block(vllm_config, groups) - if groups - else available_memory - ) - _check_enough_kv_cache_memory( - check_memory, - lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), - vllm_config.model_config.max_model_len, - lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am), - ) - - -def create_kv_cache_group_specs( - kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]] -) -> list[KVCacheGroupSpec]: - """Create KVCacheGroupSpec object for each kv cache group layer. - The layers in the same group should share the same - KVCacheSpec. - - Args: - kv_cache_spec: - A mapping from each layer name to its corresponding KVCacheSpec. - grouped_layer_names: - 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: - A list of KVCacheGroupSpec objects, one for each group. - - """ - kv_cache_groups = [] - for layer_names_one_group in grouped_layer_names: - layer_specs = [ - kv_cache_spec[layer_name] for layer_name in layer_names_one_group - ] - merged_layer_spec = layer_specs[0].merge(layer_specs) - kv_cache_groups.append( - KVCacheGroupSpec(layer_names_one_group, merged_layer_spec) - ) - return kv_cache_groups - - -def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: - """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. - - Args: - kv_cache_spec: The kv cache spec of each attention layer in the model - - Returns: - True if all layers have the same type, False otherwise. - - """ - if not kv_cache_spec: - # Encoder-only models do not have KV cache, kv_cache_type can be - # regarded as uniform. - return True - try: - kv_cache_spec_values = list(kv_cache_spec.values()) - _ = kv_cache_spec_values[0].merge(kv_cache_spec_values) - except AssertionError: - return False - return True - - def get_max_concurrency_for_kv_cache_config( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig ) -> float: @@ -1227,1052 +1013,6 @@ def get_max_concurrency_for_kv_cache_config( return min(limits) -def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: - """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`. - """ - if vllm_config.cache_config.num_gpu_blocks_override is not None: - num_blocks = vllm_config.cache_config.num_gpu_blocks_override - return num_blocks - - -def _pool_bytes_per_block( - vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec] -) -> int: - """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. - """ - return _get_kv_cache_bytes_per_block( - [group for group in kv_cache_groups if not group.host_resident], - vllm_config.cache_config.get_resolved_kv_cache_layout(), - ) - - -def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: - """Get the page size of the KV cache.""" - page_sizes = {layer.page_size_bytes for layer in kv_cache_specs} - assert len(page_sizes) == 1 - return page_sizes.pop() - - -def _get_kv_cache_groups_uniform_spec( - kv_cache_specs: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Generates the KV cache configuration for a model with the same KV cache - spec for all layers. - - Args: - kv_cache_specs: The kv cache spec of each attention layer in the model - - Returns: - The generated KVCacheGroupSpecs - - """ - return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())]) - - -def _get_kv_cache_groups_uniform_type( - spec: UniformTypeKVCacheSpecs, -) -> list[KVCacheGroupSpec]: - """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. - - Args: - spec: The UniformTypeKVCacheSpecs of the model - - Returns: - The generated KVCacheGroupSpecs - - """ - return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] - - -def unify_kv_cache_spec_page_size( - kv_cache_spec: dict[str, KVCacheSpec], -) -> dict[str, KVCacheSpec]: - """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). - - Args: - kv_cache_spec: The KVCacheSpec of each attention layer in the model - - Returns: - The updated KVCacheSpec with the same page_size_bytes. - - """ - page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()} - if len(page_sizes) <= 1: - # All layers have the same page size, no need to unify. - return kv_cache_spec - - max_page_size = max(page_sizes) - new_kv_cache_spec = {} - for layer_name, layer_spec in kv_cache_spec.items(): - if layer_spec.page_size_bytes == max_page_size: - new_kv_cache_spec[layer_name] = layer_spec - elif isinstance(layer_spec, MambaSpec): - # MambaSpec's page size is determined by its state shapes and does - # not scale with block_size, so pad the page instead. This is the - # same padding mechanism the platform uses to align Mamba pages - # with the main model's attention page size; it is needed here - # when another layer (e.g. from a draft model) has a larger page - # than the already-aligned Mamba page. - new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size) - assert new_spec.page_size_bytes == max_page_size - new_kv_cache_spec[layer_name] = new_spec - else: - layer_page_size = layer_spec.page_size_bytes - if max_page_size % layer_page_size == 0: - ratio = max_page_size // layer_page_size - new_block_size = layer_spec.block_size * ratio - new_spec = replace(layer_spec, block_size=new_block_size) - elif isinstance(layer_spec, AttentionSpec) and not isinstance( - layer_spec, MLAAttentionSpec - ): - new_spec = replace(layer_spec, page_size_padded=max_page_size) - else: - raise NotImplementedError( - f"Layer {layer_name}: page size is not divisible by the " - "maximum page size and cannot be padded. Padding is only " - "supported for non-MLA attention layers." - ) - assert new_spec.page_size_bytes == max_page_size - new_kv_cache_spec[layer_name] = new_spec - return new_kv_cache_spec - - -def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: - # kv_cache_spec is an empty dict for attention free models - return not kv_cache_spec - - -def _get_kv_cache_groups_uniform_page_size( - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """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. - - Args: - kv_cache_spec: The KVCacheSpec of each attention layer in the model - Returns: - The generated KVCacheGroupSpecs - - """ - # Group all layers by kv_cache_spec. - # E.g., 2 full attention layers and 3 sliding window attention layers, - # -> (full.0, full.1), (sw.0, sw.1, sw.2). - same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list) - for layer_name, layer_spec in kv_cache_spec.items(): - same_type_layers[layer_spec].append(layer_name) - - # Attempt to further merge same-type layers based on whether their KV - # cache specs can be merged, to minimize the group count. This benefits - # situations where specs share a block layout and differ only in a - # property it can reconcile (e.g. full attention layers differing only in - # sliding window / attention chunk size). - layer_buckets: list[list[str]] = [] - spec_buckets: list[list[KVCacheSpec]] = [] - for layer_spec, layer_names in same_type_layers.items(): - for names, specs in zip(layer_buckets, spec_buckets): - try: - # A raise means that the specs are incompatible. - type(specs[0]).merge([*specs, layer_spec]) - except (AssertionError, ValueError): - continue - names.extend(layer_names) - specs.append(layer_spec) - break - else: - layer_buckets.append(list(layer_names)) - spec_buckets.append([layer_spec]) - - # Split each group into smaller groups, to make the number of layers in each - # group identical. Add padding to the last group of each type if necessary. - # E.g., (full.0, full.1), (sw.0, sw.1, sw.2) - # split to 3 groups with 2 layers each: - # (full.0, full.1), (sw.0, sw.2), (sw.1, padding). - # FIXME(Chen): At the moment of writing this code (2025-06-02), all - # open-source hybrid model follows a n:1 pattern between different attention - # types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and - # full), so we can use the "1" in the n:1 pattern as the group size, which - # is the minimum number of layers among all attention types. Need a better - # strategy if we want to support more complex patterns (e.g., 20 full + 30 - # sw, where the group size should be 10). - min_num_layers = min([len(layers) for layers in layer_buckets]) - group_size = min_num_layers - max_num_layers = max([len(layers) for layers in layer_buckets]) - if max_num_layers < min_num_layers * 1.5: - # If the number of layers is not much larger than the minimum number of - # layers, use the maximum number of layers as the group size to avoid - # too many padding layers. A typical example is gpt-oss-20b + eagle, - # with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of - # (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding - # layers while accommodating speculative decoding drafters that add - # extra layers to one attention type. - group_size = max_num_layers - grouped_layers = [] - for layers in layer_buckets: - num_padding_layers = group_size - len(layers) % group_size - if num_padding_layers != group_size: - logger.warning( - "Add %d padding layers, may waste at most %.2f%% KV cache memory", # noqa - num_padding_layers, - num_padding_layers / len(layers) * 100, - ) - num_groups = cdiv(len(layers), group_size) - # In PP case, say if we have - # - stage 0: full.0, sw.0, sw.1 - # - stage 1: full.1, sw.2, sw.3 - # We should have 3 groups: (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3) - # It can't be (full.0, full.1), (sw.0, sw.1), (sw.2, sw.3) because - # the 3 groups in stage 0 will be (full.0), (sw.0, sw.1), (empty group) - # and it will be padded to (full.0, padding), (sw.0, sw.1), - # (padding, padding) to ensure the number of layers in each group is - # the same and will cause memory waste. - # To avoid this, we assign layers[i::num_groups] to the i-th group - # instead of layers[i * group_size: (i + 1) * group_size] - for i in range(num_groups): - grouped_layers.append(layers[i::num_groups]) - return create_kv_cache_group_specs(kv_cache_spec, grouped_layers) - - -def _get_per_layer_spec( - group: KVCacheGroupSpec, - layer_name: str, -) -> KVCacheSpec: - spec = group.kv_cache_spec - if isinstance(spec, UniformTypeKVCacheSpecs): - return spec.kv_cache_specs[layer_name] - return spec - - -def _get_layer_block_stride(spec: KVCacheSpec) -> int: - """A layer-outer block stride: the page padded to its alignment.""" - return round_up(spec.page_size_bytes, spec.get_block_stride_alignment()) - - -def _get_kv_cache_bytes_per_block( - kv_cache_groups: list[KVCacheGroupSpec], layout: KVCacheLayout -) -> int: - """Return the largest cache group's bytes per block.""" - if not layout.is_block_outermost: - # Each layer's block stride is its own padded page. - return max( - sum( - _get_layer_block_stride(_get_per_layer_spec(group, layer_name)) - for layer_name in group.layer_names - ) - for group in kv_cache_groups - ) - bytes_per_block = max( - sum( - _get_per_layer_spec(group, layer_name).page_size_bytes - for layer_name in group.layer_names - ) - for group in kv_cache_groups - ) - assert bytes_per_block > 0 - alignments = [ - group.kv_cache_spec.page_size_bytes - for group in kv_cache_groups - if isinstance(group.kv_cache_spec, HiSparseHotSpec) - ] - alignments.extend( - _get_per_layer_spec(group, layer_name).get_block_stride_alignment() - for group in kv_cache_groups - for layer_name in group.layer_names - ) - return round_up(bytes_per_block, math.lcm(*alignments)) - - -def validate_kv_cache_layout( - layout: KVCacheLayout, - kv_cache_groups: list[KVCacheGroupSpec], -) -> None: - """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. - """ - page_sizes = { - _get_per_layer_spec(group, layer_name).page_size_bytes - for group in kv_cache_groups - for layer_name in group.layer_names - } - if len(page_sizes) == 1: - # A rectangular layer dim exists; every layout can express it. - return - - # Mixed page sizes pack pages side by side within a block, which needs each page - # to be one contiguous chunk inside its block (a block-compact layout) and, with - # multiple KV cache groups, the layer dim inside the block dim. - if not layout.is_block_compact or ( - len(kv_cache_groups) > 1 and layout.is_layer_compact - ): - raise ValueError( - f"KV cache layout {layout.name} cannot express this model's " - f"mixed page sizes ({sorted(page_sizes)}); a backend should " - "declare block-outermost supported layouts (e.g. BLHNC), or " - "set VLLM_KV_CACHE_LAYOUT=BLHNC." - ) - - -def get_kv_cache_config_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], - available_memory: int, -) -> KVCacheConfig: - """Generate the KV cache configuration from the KV cache groups and spec - of each layer. - - Args: - vllm_config: The global VllmConfig - kv_cache_groups: The KV cache groups - available_memory: Memory available for KV cache in bytes - Returns: - The generated KVCacheConfig - - """ - if len(kv_cache_groups) == 0: - # Attention free models do not have KV cache. - # Return num_blocks=1 as BlockPool always needs a null_block. - return KVCacheConfig( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=kv_cache_groups, - prefix_cache_retention_interval=( - vllm_config.cache_config.prefix_cache_retention_interval - ), - ) - - if vllm_config.attention_config.hisparse_config is not None: - host_budget = get_hisparse_host_pool_bytes(vllm_config) - return get_hisparse_kv_cache_config( - vllm_config, kv_cache_groups, available_memory, host_budget - ) - - layout = vllm_config.cache_config.get_resolved_kv_cache_layout() - validate_kv_cache_layout(layout, kv_cache_groups) - bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups, layout) - interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None - - num_blocks = available_memory // bytes_per_block - num_blocks = may_override_num_blocks(vllm_config, num_blocks) - size = bytes_per_block * num_blocks - - # Groups alias from byte 0. Spec regions are laid out differently: - # - # block-outer (the same packing repeats for every block): - # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ... - # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ... - # |<--- bytes_per_block -->| - # - # layer-outer (only supported for uniform page sizes or single-group models): - # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] | - # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] | - - kv_cache_tensors = [] - for group in kv_cache_groups: - group_spec = group.kv_cache_spec - layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) - if isinstance(group_spec, UniformTypeKVCacheSpecs): - for layer_name, spec in group_spec.kv_cache_specs.items(): - layers_by_spec[spec].append(layer_name) - elif group.layer_names: - layers_by_spec[group_spec].extend(group.layer_names) - - byte_offset = 0 - for spec, layer_names in layers_by_spec.items(): - block_stride = interleaved_block_stride - padded_page = _get_layer_block_stride(spec) - if block_stride is None and padded_page != spec.page_size_bytes: - assert layout.is_block_compact, f"{layout.name} cannot pad pages" - block_stride = padded_page - layer_stride, block_stride, _, _, _ = compute_layout_strides( - spec, - num_blocks, - len(layer_names), - layout, - fixed_strides=(None, block_stride, None, None, None), - ) - kv_cache_tensors.append( - KVCacheTensor( - size=size, - layers=layer_names, - layer_stride=layer_stride, - block_stride=block_stride, - offset=byte_offset, - ) - ) - # Block-outer layers share each block; layer-outer ones own regions. - byte_offset += len(layer_names) * ( - spec.page_size_bytes if layout.is_block_outermost else layer_stride - ) - - return KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=kv_cache_tensors, - kv_cache_groups=kv_cache_groups, - prefix_cache_retention_interval=( - vllm_config.cache_config.prefix_cache_retention_interval - ), - ) - - -def _promote_local_kv_cache_specs( - kv_cache_spec: dict[str, KVCacheSpec], -) -> dict[str, KVCacheSpec]: - """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. - """ - promoted_specs = kv_cache_spec.copy() - - if is_kv_cache_spec_uniform( - promoted_specs - ) or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs): - return promoted_specs - - has_full_attention = any( - isinstance(spec, FullAttentionSpec) for spec in promoted_specs.values() - ) - has_sliding_window = any( - isinstance(spec, SlidingWindowSpec) for spec in promoted_specs.values() - ) - has_chunked_local_attention = any( - isinstance(spec, ChunkedLocalAttentionSpec) for spec in promoted_specs.values() - ) - full_block_sizes = { - spec.block_size - for spec in promoted_specs.values() - if isinstance(spec, FullAttentionSpec) - } - full_attention_block_size = ( - next(iter(full_block_sizes)) if len(full_block_sizes) == 1 else None - ) - - def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | None: - if spec.page_size_padded is None: - return None - unpadded_page_size = ( - spec.unpadded_page_size_bytes * block_size // spec.block_size - ) - return max(spec.page_size_padded, unpadded_page_size) - - promotions: dict[type[AttentionSpec], type[AttentionSpec]] = { - SlidingWindowMLASpec: MLAAttentionSpec, - SlidingWindowSpec: FullAttentionSpec, - ChunkedLocalAttentionSpec: FullAttentionSpec, - } - - if has_full_attention and (has_sliding_window or has_chunked_local_attention): - for layer_name, spec in kv_cache_spec.items(): - target_cls = next( - (promotions[c] for c in type(spec).__mro__ if c in promotions), None - ) - if target_cls is None: - continue - assert isinstance(spec, AttentionSpec) - block_size = full_attention_block_size or spec.block_size - promoted_specs[layer_name] = replace_as( - spec, - target_cls, - # Promoted specs allocate blocks for all tokens and never free - # below the window, so the trailing-edge extension is moot. - drop=("extra_retained_tokens",), - block_size=block_size, - page_size_padded=promoted_page_size_padded(spec, block_size), - ) - - if not ( - is_kv_cache_spec_uniform(promoted_specs) - or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs) - ): - raise ValueError("Failed to promote local KV cache specs to one unified type.") - - return promoted_specs - - -def _try_get_full_allocation_fallback_groups( - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec] | None: - """Try a supported full-allocation fallback for local-attention layers.""" - if any(isinstance(spec, HiddenStateCacheSpec) for spec in kv_cache_spec.values()): - return None - if any( - isinstance(spec, (SlidingWindowMLASpec, ChunkedLocalAttentionSpec)) - for spec in kv_cache_spec.values() - ): - return None - - has_mla = any(isinstance(spec, MLAAttentionSpec) for spec in kv_cache_spec.values()) - has_regular_swa = any( - isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values() - ) - if not (has_mla and has_regular_swa): - return None - - try: - promoted_specs = _promote_local_kv_cache_specs(kv_cache_spec) - except ValueError: - return None - uniform_spec = UniformTypeKVCacheSpecs.from_specs(promoted_specs) - if uniform_spec is None: - return None - logger.warning( - "KV cache page sizes cannot be unified; treating sliding-window " - "layers as full attention for cache allocation. Sliding-window " - "attention compute is unchanged." - ) - return _get_kv_cache_groups_uniform_type(uniform_spec) - - -def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]): - """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. - - Args: - kv_cache_spec: The kv cache spec of each attention layer in the model - - """ - groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict) - for name, spec in kv_cache_spec.items(): - replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded - groups[replicated][name] = spec - for specs in groups.values(): - promoted_specs = _promote_local_kv_cache_specs(specs) - if promoted_specs == specs: - continue - logger.warning( - "Hybrid KV cache manager is disabled for this hybrid model, " - "This means we do not enable any optimizations for saving KV cache " - "memory (e.g., dropping the KV cache outside the sliding window). " - "The compute of layers like sliding window is still saved." - ) - kv_cache_spec.update(promoted_specs) - - -def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: - """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. - """ - if not values: - raise ValueError("values must be non-empty") - if any(x <= 0 for x in values): - raise ValueError(f"values must be positive, got: {list(values)!r}") - - min_d = max(1, lower_bound if lower_bound is not None else 1) - max_d = max(values) - if min_d > max_d: - return min_d - - best_d = min_d - best_pad: int | None = None - for d in range(min_d, max_d + 1): - pad = sum((d - (x % d)) % d for x in values) - if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d): - best_pad = pad - best_d = d - - return best_d - - -def _get_packed_kv_cache_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec] | None: - """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. - """ - layout = vllm_config.cache_config.get_resolved_kv_cache_layout() - page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()} - if not layout.is_block_outermost or len(page_sizes) <= 1: - return None - - buckets: list[dict[str, KVCacheSpec]] = [] - for name, spec in kv_cache_spec.items(): - for bucket in buckets: - candidate = {**bucket, name: spec} - if UniformTypeKVCacheSpecs.is_uniform_type(candidate): - bucket[name] = spec - break - else: - buckets.append({name: spec}) - - bucketed = [] - for bucket in buckets: - uniform_spec = UniformTypeKVCacheSpecs.from_specs(bucket) - assert uniform_spec is not None - page_size_layers: dict[int, list[str]] = defaultdict(list) - for layer_name, layer_spec in bucket.items(): - page_size_layers[layer_spec.page_size_bytes].append(layer_name) - # Only 1:1 patterns (one layer of each page size per repeat) are - # supported; counts sharing a gcd > 1 (e.g. 2:1) could in principle - # repeat too, but such buckets are emitted whole instead. - balanced = len(set(map(len, page_size_layers.values()))) == 1 - bucketed.append((uniform_spec, page_size_layers, balanced)) - - # Balanced buckets that mix page sizes must stay whole, so the largest one - # sets a floor on the repeats per group; larger single-size buckets are - # split down toward it. No such bucket means nothing needs packing. - min_repeats_per_group = max( - ( - spec.get_max_layers_per_page_size() - for spec, page_size_layers, balanced in bucketed - if balanced and len(page_size_layers) > 1 - ), - default=0, - ) - repeats_per_group = ( - _approximate_gcd( - [ - spec.get_max_layers_per_page_size() - for spec, _, balanced in bucketed - if balanced - ], - lower_bound=min_repeats_per_group, - ) - if min_repeats_per_group - else None - ) - - def num_groups_for(spec: UniformTypeKVCacheSpecs, balanced: bool) -> int: - if balanced and repeats_per_group is not None: - return cdiv(spec.get_max_layers_per_page_size(), repeats_per_group) - return 1 - - def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: - """Page bytes of the largest of the n groups a bucket splits into.""" - return sum( - cdiv(len(names), n) * page for page, names in page_size_layers.items() - ) - - # Bytes a block must hold however the state buckets end up split: mamba, - # circular-buffer and unsplit sliding-window buckets can go down to one - # state per group, every other bucket's split is fixed by the repeat pattern. - def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: - if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)): - return True - return repeats_per_group is None and isinstance( - spec.first_spec, SlidingWindowSpec - ) - - anchor_bytes = max( - ( - widest_group_bytes( - page_size_layers, - len(spec.kv_cache_specs) - if is_state_bucket(spec) - else num_groups_for(spec, balanced), - ) - for spec, page_size_layers, balanced in bucketed - ), - default=0, - ) - - groups = [] - for spec, page_size_layers, balanced in bucketed: - num_groups = num_groups_for(spec, balanced) - # Cap a state group at the states a block already fits rather than let - # it widen the block. - if anchor_bytes and is_state_bucket(spec): - states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1) - num_groups = max( - num_groups, cdiv(len(spec.kv_cache_specs), states_per_block) - ) - if num_groups == 1: - groups.append(KVCacheGroupSpec(list(spec.kv_cache_specs), spec)) - continue - - pattern_repeats = list(zip(*page_size_layers.values())) - for i in range(num_groups): - group_layer_names = [ - name for repeat in pattern_repeats[i::num_groups] for name in repeat - ] - group_layer_specs = { - name: spec.kv_cache_specs[name] for name in group_layer_names - } - group_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs) - assert group_spec is not None - groups.append(KVCacheGroupSpec(group_layer_names, group_spec)) - - _annotate_eagle_groups( - vllm_config, - kv_cache_spec, - groups, - use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), - ) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) - return groups - - -def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool: - """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. - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): - return False - if spec_config.method == "mtp": - return True - model_config = vllm_config.model_config - return model_config is not None and model_config.hf_config.model_type in ( - "deepseek_v4", - "deepseek_v41", - ) - - -def _groups_partition_layers_exactly( - kv_cache_spec: dict[str, KVCacheSpec], - kv_cache_groups: list[KVCacheGroupSpec], -) -> bool: - """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. - """ - names = [name for group in kv_cache_groups for name in group.layer_names] - return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec) - - -def _annotate_eagle_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - kv_cache_groups: list[KVCacheGroupSpec], - use_trailing_layer_fallback: bool = False, -) -> None: - """Flag the KV cache groups that hold drafter attention layers. - - Two detection rules, in order of preference: - - 1. Spec-driven. ``non_causal_multi_token_decode`` is 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. - 2. 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. - - Args: - vllm_config: Config supplying the speculative method, if any. - kv_cache_spec: The kv cache spec of each attention layer, in layer - registration order. Only read by rule 2. - kv_cache_groups: Groups to annotate in place. - use_trailing_layer_fallback: Enable rule 2. Callers gate this on - ``_uses_trailing_mtp_layers``. - - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle_block_drop(): - return - - for group in kv_cache_groups: - if any( - getattr(spec, "non_causal_multi_token_decode", False) - for spec in iter_layer_specs(group.kv_cache_spec) - ): - group.is_eagle_group = True - - if not use_trailing_layer_fallback: - return - if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups): - return - last_layer = next(reversed(kv_cache_spec)) - for group in kv_cache_groups: - if last_layer in group.layer_names: - group.is_eagle_group = True - break - - -def _warn_if_unannotated_eagle_mamba( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], -) -> None: - """Warn when no KV cache group could be identified as the draft model's. - - Args: - vllm_config: Config supplying the speculative method, if any. - kv_cache_groups: Groups as they will be handed to consumers. - - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): - return - if any(group.is_eagle_group for group in kv_cache_groups): - return - mamba_groups = [ - idx - for idx, group in enumerate(kv_cache_groups) - if any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ) - ] - if not mamba_groups: - return - logger.warning( - "Speculative decoding (method=%s) is enabled but no KV cache group " - "could be identified as the draft model's.", - spec_config.method, - ) - - -def _largest_divisor_at_most(value: int, limit: int) -> int: - for candidate in range(min(value, limit), 0, -1): - if value % candidate == 0: - return candidate - return 1 - - -def _ensure_min_page_size( - groups: list[KVCacheGroupSpec], - common_page: int, - hidden_specs: dict[str, HiddenStateCacheSpec], -) -> tuple[list[KVCacheGroupSpec], int]: - """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. - """ - max_per_token = max( - spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) - for spec in hidden_specs.values() - ) - if max_per_token <= common_page: - return groups, common_page - - scale = cdiv(max_per_token, common_page) - common_page *= scale - scaled: list[KVCacheGroupSpec] = [] - for g in groups: - s = g.kv_cache_spec - kw: dict[str, Any] = {"block_size": s.block_size * scale} - if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None: - kw["page_size_padded"] = s.page_size_padded * scale - scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw))) - return scaled, common_page - - -def get_kv_cache_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Split the layers in the model into groups with the same KV cache spec. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - - Returns: - The generated KVCacheGroups - - """ - if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: - unify_hybrid_kv_cache_specs(kv_cache_spec) - - if is_kv_cache_type_attention_free(kv_cache_spec): - # This returns an empty list to allow for the KVCacheManager to handle - # attention free models. - return [] - - if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec): - return hisparse_groups - - if is_kv_cache_spec_uniform(kv_cache_spec): - # KV cache of all layers are the same, which is true for - # most models. Allocate the same amount of memory for - # each layer. - return _get_kv_cache_groups_uniform_spec(kv_cache_spec) - elif uniform_spec := UniformTypeKVCacheSpecs.from_specs(kv_cache_spec): - # All layers need the same number of token slots (e.g., all layers are - # full attention, or all layers are sliding window attention with the - # same window size). Put all layers into one group. - return _get_kv_cache_groups_uniform_type(uniform_spec) - # Hidden-state layers use their own block table and must not be absorbed - # into a compatible attention bucket. - hidden_specs = { - k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec) - } - filtered_spec = { - k: v - for k, v in kv_cache_spec.items() - if not isinstance(v, HiddenStateCacheSpec) - } - - if packed_groups := _get_packed_kv_cache_groups(vllm_config, filtered_spec): - # Block-outermost blocks are strided by the widest group, so hidden - # groups need no page alignment. - packed_groups += [ - KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items() - ] - return packed_groups - - # Prefer preserving each layer's cache semantics. If physical pages cannot - # be unified, try a supported allocation-only fallback before failing. - try: - filtered_spec = unify_kv_cache_spec_page_size(filtered_spec) - except NotImplementedError: - fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec) - if fallback_groups is None: - raise - return fallback_groups - groups = _get_kv_cache_groups_uniform_page_size(filtered_spec) - - # Add hidden-state layers back with page aligned to the common page. - if hidden_specs: - common_page = get_uniform_page_size([g.kv_cache_spec for g in groups]) - # TP may shrink the common page below hidden-state per-token cost. - groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs) - group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups)) - for name, spec in hidden_specs.items(): - per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) - max_block_size = max(common_page // per_token, 1) - new_bs = _largest_divisor_at_most(group_block_size, max_block_size) - wasted_bytes = common_page - new_bs * per_token - logger.info( - "Using block size %d for hidden-state cache layer %s; " - "page alignment wastes %d bytes (%.2f%%) per block", - new_bs, - name, - wasted_bytes, - wasted_bytes / common_page * 100, - ) - aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) - groups.append(KVCacheGroupSpec([name], aligned)) - - _annotate_eagle_groups( - vllm_config, - kv_cache_spec, - groups, - use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), - ) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) - return groups - - -def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int: - if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None: - return 1 - if spec.dcp_sharded and dcp_size > 1: - return 1 - return max(1, tp_size // spec.max_tp_shards) - - -def kv_cache_groups_tp_replicas( - groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1 -) -> int: - """Consecutive TP ranks holding identical KV for every layer.""" - specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)] - if not specs: - return 1 - return reduce( - math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size - ) - - def generate_scheduler_kv_cache_config( kv_cache_configs: list[KVCacheConfig], ) -> KVCacheConfig: @@ -2334,359 +1074,6 @@ def update_kv_cache_capacity( ) -def _max_memory_usage_bytes_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], -) -> int: - """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``. - """ - kv_cache_groups = [group for group in kv_cache_groups if not group.host_resident] - if not kv_cache_groups: - return 0 - - bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups) - total_blocks = 0 - for group in kv_cache_groups: - spec = group.kv_cache_spec - if isinstance(spec, UniformTypeKVCacheSpecs): - total_blocks += spec.max_memory_usage_pages(vllm_config) - else: - total_blocks += cdiv( - spec.max_memory_usage_bytes(vllm_config), - spec.page_size_bytes, - ) - - return bytes_per_block * total_blocks - - -def _estimate_max_model_len_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], - available_memory: int, -) -> int: - """Binary search for the maximum model length that fits in available memory. - Returns 0 if even 1 token doesn't fit. - """ - original_max = vllm_config.model_config.max_model_len - hisparse_enabled = ( - vllm_config.attention_config.hisparse_config is not None - and bool(kv_cache_groups) - ) - - def fits(model_len: int) -> bool: - vllm_config.model_config.max_model_len = model_len - if hisparse_enabled: - try: - get_kv_cache_config_from_groups( - vllm_config, kv_cache_groups, available_memory - ) - except ValueError: - return False - return ( - _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups) - <= available_memory - ) - - try: - left, right = 1, original_max - if not fits(left): - return 0 - result = 1 - while left <= right: - mid = (left + right) // 2 - if fits(mid): - result = mid - left = mid + 1 - else: - right = mid - 1 - return result - finally: - vllm_config.model_config.max_model_len = original_max - - -def _auto_fit_max_model_len( - vllm_config: VllmConfig, - projected_groups_per_worker: list[list[KVCacheGroupSpec]], - available_memory: list[int], -) -> None: - """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. - - Args: - vllm_config: The global VllmConfig (will be modified in-place) - projected_groups_per_worker: KV cache groups projected to each worker. - available_memory: Memory available for KV cache in bytes for each - worker. - - """ - original_max = vllm_config.model_config.max_model_len - - if all(not groups for groups in projected_groups_per_worker): - # All workers have empty specs (attention-free model) - logger.info_once( - "Auto-fit max_model_len: attention-free model, " - "using derived max_model_len=%d", - original_max, - ) - return - - # Find the max_model_len that fits across all workers. - auto_fit_max = original_max - limiting_worker_mem = available_memory[0] - for groups, avail_mem in zip(projected_groups_per_worker, available_memory): - if not groups: - continue - worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem) - if worker_max < auto_fit_max: - auto_fit_max = worker_max - limiting_worker_mem = avail_mem - - if auto_fit_max <= 0: - raise ValueError( - "Cannot auto-fit max_model_len: not enough GPU memory available " - "to serve even a single token. Try increasing `gpu_memory_utilization`." - ) - - kv_cache_memory = ( - f"{format_gib(limiting_worker_mem)} GiB {_kv_cache_device_name(vllm_config)}" - ) - if vllm_config.attention_config.hisparse_config is not None: - host_pool_gib = format_gib(get_hisparse_host_pool_bytes(vllm_config)) - kv_cache_memory += f" + {host_pool_gib} GiB host" - - if auto_fit_max >= original_max: - # The model's full context length fits in memory - logger.info_once( - "Auto-fit max_model_len: full model context length %d fits in " - "available KV cache memory (%s)", - original_max, - kv_cache_memory, - ) - else: - # Need to reduce max_model_len to fit in memory - vllm_config.model_config.max_model_len = auto_fit_max - logger.info_once( - "Auto-fit max_model_len: reduced from %d to %d to fit in " - "available KV cache memory (%s)", - original_max, - auto_fit_max, - kv_cache_memory, - ) - - -def _project_kv_cache_groups_to_worker( - global_kv_cache_groups: list[KVCacheGroupSpec], - worker_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """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. - - Args: - global_kv_cache_groups: The global KV cache groups for the whole model. - worker_spec: The KV cache spec of each layer on this worker. - - Returns: - The projected KV cache groups containing only this worker's layers. - - """ - projected_groups: list[KVCacheGroupSpec] = [] - for group in global_kv_cache_groups: - worker_layer_names = [ - layer_name for layer_name in group.layer_names if layer_name in worker_spec - ] - if len(worker_layer_names) == len(group.layer_names): - projected_groups.append(group) - continue - group_spec = group.kv_cache_spec - if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs): - group_spec = UniformTypeKVCacheSpecs( - block_size=group_spec.block_size, - kv_cache_specs={ - layer_name: group_spec.kv_cache_specs[layer_name] - for layer_name in worker_layer_names - }, - ) - projected_groups.append( - KVCacheGroupSpec( - worker_layer_names, - group_spec, - is_eagle_group=group.is_eagle_group and bool(worker_layer_names), - ) - ) - return projected_groups - - -def get_kv_cache_configs( - vllm_config: VllmConfig, - kv_cache_specs: list[dict[str, KVCacheSpec]], - available_memory: list[int], -) -> list[KVCacheConfig]: - """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. - - Args: - vllm_config: The global VllmConfig - kv_cache_specs: List of dict[layer_name, KVCacheSpec] for each worker. - available_memory: Memory available for KV cache in bytes for each - worker. - - Returns: - The generated KVCacheConfigs for each worker. - - """ - # Merge the KV cache specs of all workers. Different PP stages may have - # different layer names, and different TP ranks of the same PP stage should - # have the same KV cache spec. - merged_kv_cache_specs: dict[str, KVCacheSpec] = {} - for kv_cache_spec_one_worker in kv_cache_specs: - for layer_name, layer_spec in kv_cache_spec_one_worker.items(): - if layer_name not in merged_kv_cache_specs: - merged_kv_cache_specs[layer_name] = layer_spec - else: - assert merged_kv_cache_specs[layer_name] == layer_spec, ( - "The KV cache specs for the same layer are different " - "across workers. This is not supported yet." - ) - - # Check if the KV cache specs are registered correctly. - # This is to prevent that some layers are initialized with unregistered specs. - KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) - - # When speculating with more than 1 speculative module (e.g. multi-layered MTP) - # tag every SlidingWindowSpec with how many extra tokens to retain in the window. - extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1) - for layer_name, layer_spec in merged_kv_cache_specs.items(): - if isinstance(layer_spec, SlidingWindowSpec): - merged_kv_cache_specs[layer_name] = replace( - layer_spec, extra_retained_tokens=extra_retained_tokens - ) - - # Get global KV cache groups. This also handles spec unification for - # hybrid models when disable_hybrid_kv_cache_manager is enabled. - # After this call, merged_kv_cache_specs may be modified in-place. - global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs) - - # If original_max_model_len was -1, automatically - # determine the maximum model length that fits in available GPU memory. - # We use per-worker projected groups to account for PP sharding. - projected_groups_per_worker = [ - _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec) - for worker_spec in kv_cache_specs - ] - - # If `num_gpu_blocks_override` is set, the cache size that will actually - # be allocated is decoupled from the profiled `available_memory`: - # `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps - # `num_blocks` to the override. Reflect that in `available_memory` here so - # auto-fit, the admission check, and the per-worker config builder all - # plan against the same effective capacity. - override = vllm_config.cache_config.num_gpu_blocks_override - if override is not None: - adjusted_memory: list[int] = [] - for groups, avail_mem in zip(projected_groups_per_worker, available_memory): - if not groups: - adjusted_memory.append(avail_mem) - continue - bytes_per_block = _pool_bytes_per_block(vllm_config, groups) - logger.info( - "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d", - avail_mem // bytes_per_block, - override, - ) - adjusted_memory.append(override * bytes_per_block) - available_memory = adjusted_memory - - # Reserve the null block BlockPool permanently holds back, so auto-fit and - # the capacity check both plan against usable blocks. Allocation below - # still uses the full memory. - check_memory = [ - avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem - for groups, avail_mem in zip(projected_groups_per_worker, available_memory) - ] - - if vllm_config.model_config.original_max_model_len == -1: - _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory) - - # Check if the available memory is enough per worker. - for groups, avail_mem in zip(projected_groups_per_worker, check_memory): - if not groups: - continue - _check_enough_kv_cache_memory( - avail_mem, - partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), - vllm_config.model_config.max_model_len, - partial(_estimate_max_model_len_from_groups, vllm_config, groups), - ) - - kv_cache_configs: list[KVCacheConfig] = [] - for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip( - projected_groups_per_worker, kv_cache_specs, available_memory - ): - assert sum(len(group.layer_names) for group in projected_groups) == len( - kv_cache_spec_one_worker - ), "Some layers are not assigned to any group." - kv_cache_configs.append( - get_kv_cache_config_from_groups( - vllm_config, projected_groups, available_memory_one_worker - ) - ) - - # Change the num_blocks of each rank to the smallest among all ranks. - # We also need to shrink the tensor size proportionally to avoid - # allocating unused memory. - min_num_blocks = min( - kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs - ) - for i, kv_cache_config in enumerate(kv_cache_configs): - if kv_cache_config.num_blocks == min_num_blocks: - continue - # Re-plan with exactly the memory the smallest rank can afford, so - # strides and offsets stay consistent with the shrunken allocation. - groups = kv_cache_config.kv_cache_groups - kv_cache_configs[i] = get_kv_cache_config_from_groups( - vllm_config, - groups, - min_num_blocks * _pool_bytes_per_block(vllm_config, groups), - ) - - for kv_cache_config in kv_cache_configs: - kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas( - kv_cache_config.kv_cache_groups, - vllm_config.parallel_config.tensor_parallel_size, - vllm_config.parallel_config.decode_context_parallel_size, - ) - - return kv_cache_configs - - class BlockHashListWithBlockSize: """Convert block-hash granularity from `hash_block_size` to `target_block_size`. Used when KV cache groups have different block sizes: `hash_block_size` diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 35cc8cc186..d8c976ff69 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -50,10 +50,10 @@ from vllm.utils.hashing import get_hash_fn_by_name from vllm.utils.network_utils import make_zmq_socket from vllm.utils.system_utils import decorate_logs, set_process_title from vllm.v1.attention.backends.utils import resolve_kv_cache_layout +from vllm.v1.core.kv_cache_planning import get_kv_cache_configs from vllm.v1.core.kv_cache_utils import ( BlockHash, generate_scheduler_kv_cache_config, - get_kv_cache_configs, get_request_block_hasher, init_none_hash, resolve_cache_hit_alignment_tokens, diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py index 96d93650ad..0a72427d42 100644 --- a/vllm/v1/hisparse/layout.py +++ b/vllm/v1/hisparse/layout.py @@ -71,7 +71,7 @@ def get_hisparse_kv_cache_groups( if not mla_specs: return None - from vllm.v1.core.kv_cache_utils import get_kv_cache_groups + from vllm.v1.core.kv_cache_planning import get_kv_cache_groups mla_group_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs) assert mla_group_spec is not None @@ -323,7 +323,7 @@ def get_hisparse_kv_cache_config( available_memory: int, host_budget: int, ) -> KVCacheConfig: - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( _get_kv_cache_bytes_per_block, may_override_num_blocks, validate_kv_cache_layout, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 73be62319a..3a6409db55 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -989,7 +989,7 @@ def _extrapolate_full_graph_memory(mem_samples: list[int], total_graphs: int) -> def _init_minimal_kv_cache_for_profiling(runner: "GPUModelRunner") -> None: """Allocate the smallest KV cache that still lets every graph be captured.""" - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, get_kv_cache_groups, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e19c1b84f4..90059fd39f 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6206,7 +6206,7 @@ class GPUModelRunner( gc.collect() def _init_minimal_kv_cache_for_profiling(self) -> None: - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, get_kv_cache_groups, )