[Core][Refactor] Move KV cache planning logic to kv_cache_planning.py (#61183)

Signed-off-by: Dakai An <dakaian108@gmail.com>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

.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

.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

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)

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 (

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

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

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,
+    )

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
     )
 

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,

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,

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,

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,

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,

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,

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,

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,
     )

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,
         )