[Core][Feat] Pluggable KVCacheConfigBuilder for platform/model-specific KV cache planning (#53558)

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

tests/v1/core/test_kv_cache_config_builder.py

@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for KVCacheConfigBuilder resolution."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm.platforms import Platform
+from vllm.v1.core.kv_cache_config_builder import (
+    KVCacheConfigBuilder,
+    get_kv_cache_config_builder,
+)
+from vllm.v1.core.kv_cache_planning import (
+    DefaultKVCacheConfigBuilder,
+    get_kv_cache_groups_from_builder,
+)
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
+
+
+def _make_vllm_config(supports_grouping: bool = False) -> MagicMock:
+    """Create a minimal mock VllmConfig for builder resolution tests.
+
+    ``supports_grouping`` selects whether the resolved model class declares
+    ``SupportsKVCacheGrouping`` (exposing a ``get_kv_cache_groups`` classmethod).
+    """
+    cfg = MagicMock()
+    cfg.model_config.supports_kv_cache_grouping = supports_grouping
+    cfg.attention_config.hisparse_config = None
+    if supports_grouping:
+        cfg.model_config.registry.resolve_model_cls.return_value = (
+            _GroupedModel,
+            "arch",
+        )
+    return cfg
+
+
+CUSTOM_PATH = "tests.v1.core.test_kv_cache_config_builder.CustomBuilder"
+DEFAULT_PATH = "vllm.v1.core.kv_cache_planning.DefaultKVCacheConfigBuilder"
+
+
+class CustomBuilder(DefaultKVCacheConfigBuilder):
+    """A test builder subclass."""
+
+    pass
+
+
+class _GroupedModel:
+    """A test model class with a ``get_kv_cache_groups`` classmethod."""
+
+    @classmethod
+    def get_kv_cache_groups(cls, vllm_config, kv_cache_spec):
+        return []
+
+
+class _PartialGroupedModel:
+    """A test model class whose grouping omits layers."""
+
+    @classmethod
+    def get_kv_cache_groups(cls, vllm_config, kv_cache_spec):
+        from vllm.v1.kv_cache_interface import KVCacheGroupSpec
+
+        return [KVCacheGroupSpec(["a"], kv_cache_spec["a"])]
+
+
+class TestPlatformHookResolution:
+    """The platform hook selects only the placement builder."""
+
+    def test_default_hook_returns_default_builder(self):
+        cfg = _make_vllm_config()
+        assert Platform.get_kv_cache_config_builder_cls(cfg) == DEFAULT_PATH
+
+
+class TestBuilderResolution:
+    @patch("vllm.platforms.current_platform")
+    @pytest.mark.parametrize(
+        ("builder_cls_path", "expected_cls"),
+        [
+            (None, DefaultKVCacheConfigBuilder),
+            (CUSTOM_PATH, CustomBuilder),
+        ],
+    )
+    def test_resolves_hook_selected_builder(
+        self, mock_platform, builder_cls_path, expected_cls
+    ):
+        mock_platform.get_kv_cache_config_builder_cls.return_value = (
+            builder_cls_path or DEFAULT_PATH
+        )
+        cfg = _make_vllm_config()
+        assert type(get_kv_cache_config_builder(cfg)) is expected_cls
+
+    @patch("vllm.platforms.current_platform")
+    def test_resolves_fresh_instance_per_call(self, mock_platform):
+        mock_platform.get_kv_cache_config_builder_cls.return_value = CUSTOM_PATH
+        cfg = _make_vllm_config()
+        assert get_kv_cache_config_builder(cfg) is not get_kv_cache_config_builder(cfg)
+
+    def test_public_interface_declares_two_placement_hooks(self):
+        # get_kv_cache_configs lives in kv_cache_planning and is Core-owned;
+        # the builder only declares the placement hooks.
+        assert KVCacheConfigBuilder.__abstractmethods__ == {
+            "get_kv_cache_groups",
+            "get_kv_cache_config_from_groups",
+        }
+        assert issubclass(DefaultKVCacheConfigBuilder, KVCacheConfigBuilder)
+        assert (
+            "pool_bytes_per_block"
+            not in DefaultKVCacheConfigBuilder.__abstractmethods__
+        )
+        assert (
+            "max_memory_usage_bytes"
+            not in DefaultKVCacheConfigBuilder.__abstractmethods__
+        )
+
+
+class TestModelGroupingPriority:
+    """A model class with a get_kv_cache_groups classmethod owns grouping."""
+
+    def test_model_classmethod_owns_grouping(self):
+        cfg = _make_vllm_config(supports_grouping=True)
+        assert DefaultKVCacheConfigBuilder().get_kv_cache_groups(cfg, {}) == []
+
+    def test_model_grouping_rejects_hisparse(self):
+        cfg = _make_vllm_config(supports_grouping=True)
+        cfg.attention_config.hisparse_config = object()
+        with pytest.raises(ValueError, match="HiSparse"):
+            DefaultKVCacheConfigBuilder().get_kv_cache_groups(cfg, {})
+
+
+class TestCoreChecksGrouping:
+    """Core checks the grouping whether a model or a platform produced it."""
+
+    _SPEC = FullAttentionSpec(
+        block_size=16, num_kv_heads=1, head_size=128, dtype="float16"
+    )
+
+    def test_model_grouping_is_checked(self):
+        cfg = _make_vllm_config(supports_grouping=True)
+        cfg.model_config.registry.resolve_model_cls.return_value = (
+            _PartialGroupedModel,
+            "arch",
+        )
+        with pytest.raises(AssertionError, match="miss layers"):
+            get_kv_cache_groups_from_builder(
+                DefaultKVCacheConfigBuilder(), cfg, {"a": self._SPEC, "b": self._SPEC}
+            )
+
+    @pytest.mark.parametrize(
+        ("groups", "match"),
+        [
+            ([KVCacheGroupSpec(["a"], _SPEC)], "miss layers"),
+            (
+                [KVCacheGroupSpec(["a", "b"], _SPEC), KVCacheGroupSpec(["b"], _SPEC)],
+                "overlaps",
+            ),
+        ],
+    )
+    def test_platform_grouping_is_checked(self, groups, match):
+        class FixedGroupsBuilder(DefaultKVCacheConfigBuilder):
+            def get_kv_cache_groups(self, vllm_config, kv_cache_spec):
+                return groups
+
+        with pytest.raises(AssertionError, match=match):
+            get_kv_cache_groups_from_builder(
+                FixedGroupsBuilder(),
+                _make_vllm_config(),
+                {"a": self._SPEC, "b": self._SPEC},
+            )
+
+
+class TestPlatformCustomPriority:
+    """A vendor platform can override the hook to customize priority."""
+
+    def test_platform_builder_wins_over_model_declaration(self):
+        class PlatformFirstPlatform(Platform):
+            @classmethod
+            def get_kv_cache_config_builder_cls(cls, vllm_config):
+                return CUSTOM_PATH
+
+        cfg = _make_vllm_config()
+        assert PlatformFirstPlatform.get_kv_cache_config_builder_cls(cfg) == CUSTOM_PATH
+        with patch("vllm.platforms.current_platform", PlatformFirstPlatform):
+            assert isinstance(get_kv_cache_config_builder(cfg), CustomBuilder)
+
+
+class TestProjectionPreservesGroupMetadata:
+    """PP projection must carry host_resident/role/enable_kv_transfer."""
+
+    def test_partial_projection_preserves_metadata(self):
+        from vllm.v1.core.kv_cache_planning import _project_kv_cache_groups_to_worker
+        from vllm.v1.kv_cache_interface import (
+            FullAttentionSpec,
+            KVCacheGroupRole,
+        )
+
+        spec = FullAttentionSpec(
+            block_size=16, num_kv_heads=1, head_size=128, dtype="float16"
+        )
+        group = MagicMock()
+        group.layer_names = ["layer1", "layer2", "layer3"]
+        group.kv_cache_spec = spec
+        group.is_eagle_group = False
+        group.host_resident = True
+        group.enable_kv_transfer = False
+        group.role = KVCacheGroupRole.HISPARSE_SOURCE
+
+        projected = _project_kv_cache_groups_to_worker(
+            [group], {"layer1": spec, "layer2": spec}
+        )
+        assert len(projected) == 1
+        p = projected[0]
+        assert p.host_resident is True
+        assert p.enable_kv_transfer is False
+        assert p.role is KVCacheGroupRole.HISPARSE_SOURCE
+        assert p.layer_names == ["layer1", "layer2"]
+
+
+class TestBytesPerBlockProbe:
+    """The single-block probe must not be polluted by num_gpu_blocks_override."""
+
+    def test_probe_suppresses_override(self):
+        import torch
+
+        from vllm.config import CacheConfig, ModelConfig, VllmConfig
+        from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
+        from vllm.v1.core.kv_cache_planning import (
+            DefaultKVCacheConfigBuilder,
+            _pool_bytes_per_block,
+            _pool_bytes_per_block_from_builder,
+        )
+        from vllm.v1.kv_cache_interface import (
+            FullAttentionSpec,
+            KVCacheGroupSpec,
+        )
+
+        spec = FullAttentionSpec(
+            block_size=16, num_kv_heads=1, head_size=128, dtype=torch.float16
+        )
+        vc = VllmConfig(
+            model_config=ModelConfig(max_model_len=100),
+            cache_config=CacheConfig(block_size=16),
+        )
+        resolve_kv_cache_layout(vc, [["BLHNC"]], [spec])
+        groups = [
+            KVCacheGroupSpec(["a"], spec),
+            KVCacheGroupSpec(["b"], spec),
+        ]
+
+        pure_c = _pool_bytes_per_block(vc, groups)
+        vc.cache_config.num_gpu_blocks_override = 500
+        probed_c = _pool_bytes_per_block_from_builder(
+            DefaultKVCacheConfigBuilder(), vc, groups
+        )
+        assert probed_c == pure_c
+        assert vc.cache_config.num_gpu_blocks_override == 500
+
+    @pytest.mark.parametrize("declares_pool_bytes", [False, True])
+    def test_several_device_allocations_need_declared_pool_bytes(
+        self, declares_pool_bytes
+    ):
+        """A placement with one allocation per group cannot be sized from one
+        backing; the builder must declare its per-block bytes."""
+        from vllm.v1.core.kv_cache_planning import _pool_bytes_per_block_from_builder
+        from vllm.v1.kv_cache_interface import (
+            KVCacheConfig,
+            KVCacheGroupSpec,
+            KVCacheTensor,
+        )
+
+        class PerGroupAllocationBuilder(DefaultKVCacheConfigBuilder):
+            def get_kv_cache_config_from_groups(self, vllm_config, groups, num_blocks):
+                tensors = [
+                    KVCacheTensor(
+                        size=page * num_blocks,
+                        layers=group.layer_names,
+                        layer_stride=page * num_blocks,
+                        block_stride=page,
+                    )
+                    for group, page in zip(groups, (64, 32))
+                ]
+                return KVCacheConfig(num_blocks, tensors, groups)
+
+            def _get_pool_bytes_per_block(self, vllm_config, groups):
+                return 64 + 32 if declares_pool_bytes else None
+
+        vc = _make_vllm_config()
+        vc.cache_config.num_gpu_blocks_override = None
+        groups = [KVCacheGroupSpec([name], MagicMock()) for name in "ab"]
+        builder = PerGroupAllocationBuilder()
+        if declares_pool_bytes:
+            assert _pool_bytes_per_block_from_builder(builder, vc, groups) == 96
+        else:
+            with pytest.raises(ValueError, match="_get_pool_bytes_per_block"):
+                _pool_bytes_per_block_from_builder(builder, vc, groups)

vllm/config/model.py

@@ -2034,6 +2034,10 @@ class ModelConfig:
     def has_noops(self) -> bool:
         return self._model_info.has_noops
 
+    @property
+    def supports_kv_cache_grouping(self) -> bool:
+        return self._model_info.supports_kv_cache_grouping
+
     @property
     def has_inner_state(self):
         return self._model_info.has_inner_state

vllm/model_executor/models/interfaces.py

@@ -64,6 +64,7 @@ if TYPE_CHECKING:
     from vllm.tasks import ScoreType
     from vllm.tokenizers import TokenizerLike
     from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+    from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec
     from vllm.v1.worker.encoder_cudagraph_defs import (
         EncoderCudaGraphCaptureInputs,
         EncoderCudaGraphConfig,
@@ -669,6 +670,47 @@ def supports_score_template(
     return getattr(model, "supports_score_template", False)
 
 
+@runtime_checkable
+class SupportsKVCacheGrouping(Protocol):
+    """The interface required for models that own their KV cache grouping.
+
+    .. note::
+        Experimental. The classmethod contract may change as model-owned
+        grouping stabilizes.
+
+    Implementing it lets a model replace the default KV cache grouping (which
+    layers share a block table) while Core keeps owning placement, byte
+    accounting, and admission.
+    """
+
+    @classmethod
+    def get_kv_cache_groups(
+        cls,
+        vllm_config: "VllmConfig",
+        kv_cache_spec: dict[str, "KVCacheSpec"],
+    ) -> list["KVCacheGroupSpec"]:
+        """Split the model's layers into logical KV cache groups."""
+        ...
+
+
+@overload
+def supports_kv_cache_grouping(
+    model: type[object],
+) -> TypeIs[type[SupportsKVCacheGrouping]]: ...
+
+
+@overload
+def supports_kv_cache_grouping(
+    model: object,
+) -> TypeIs[SupportsKVCacheGrouping]: ...
+
+
+def supports_kv_cache_grouping(
+    model: type[object] | object,
+) -> TypeIs[type[SupportsKVCacheGrouping]] | TypeIs[SupportsKVCacheGrouping]:
+    return isinstance(model, SupportsKVCacheGrouping)
+
+
 @runtime_checkable
 class SupportsLoRA(Protocol):
     """The interface required for all models that support LoRA."""

vllm/model_executor/models/registry.py

@@ -50,6 +50,7 @@ from .interfaces import (
     is_hybrid,
     requires_raw_input_tokens,
     supports_hc_sp,
+    supports_kv_cache_grouping,
     supports_multimodal,
     supports_multimodal_encoder_tp_data,
     supports_multimodal_raw_input_only,
@@ -882,6 +883,7 @@ class _ModelInfo:
     supports_replayssm: bool
     supports_transcription: bool
     supports_transcription_only: bool
+    supports_kv_cache_grouping: bool
     supported_video_pruning_methods: tuple[str, ...]
     supports_mm_device_do_normalize: bool
 
@@ -913,6 +915,7 @@ class _ModelInfo:
             supports_transcription_only=(
                 supports_transcription(model) and model.supports_transcription_only
             ),
+            supports_kv_cache_grouping=supports_kv_cache_grouping(model),
             has_noops=has_noops(model),
             supported_video_pruning_methods=getattr(
                 model, "supported_video_pruning_methods", ()

vllm/platforms/interface.py

@@ -1467,6 +1467,28 @@ class Platform:
                         "Unknown vLLM environment variable detected: %s", env
                     )
 
+    @classmethod
+    def get_kv_cache_config_builder_cls(cls, vllm_config: "VllmConfig") -> str:
+        """Resolve the fully-qualified class path of the platform's KV cache
+        config builder.
+
+        .. note::
+            Experimental, like the ``KVCacheConfigBuilder`` interface it selects.
+
+        The builder supplies grouping (``get_kv_cache_groups``) and tensor
+        placement (``get_kv_cache_config_from_groups``). Vendors may override
+        this to return a platform-specific builder. The default builder honors a
+        model implementing
+        :class:`~vllm.model_executor.models.interfaces.SupportsKVCacheGrouping`;
+        a builder that overrides ``get_kv_cache_groups`` owns grouping and
+        bypasses that model hook unless it calls ``super()``.
+
+        Returns:
+            Fully-qualified class path of a ``KVCacheConfigBuilder`` subclass.
+
+        """
+        return "vllm.v1.core.kv_cache_planning.DefaultKVCacheConfigBuilder"
+
 
 class UnspecifiedPlatform(Platform):
     _enum = PlatformEnum.UNSPECIFIED

vllm/v1/core/kv_cache_config_builder.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Interface and resolution for pluggable KV cache config builders."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.kv_cache_interface import (
+        KVCacheConfig,
+        KVCacheGroupSpec,
+        KVCacheSpec,
+    )
+
+
+class KVCacheConfigBuilder(ABC):
+    """Interface for platform KV cache placement.
+
+    .. note::
+        This interface is experimental and subject to change. The exact hooks
+        and the ``num_blocks`` contract may evolve before it stabilizes.
+
+    Subclasses normally inherit from ``DefaultKVCacheConfigBuilder`` and
+    override only the hooks they need. The cross-worker planning flow
+    (:func:`vllm.v1.core.kv_cache_planning.get_kv_cache_configs`) is Core-owned;
+    this builder supplies the placement decisions: how layers are split into
+    groups and how a group list is materialized into tensors for a given block
+    count. Grouping coordination lives in ``DefaultKVCacheConfigBuilder`` (a
+    model implementing :class:`~vllm.model_executor.models.interfaces.
+    SupportsKVCacheGrouping` owns its grouping there); a platform that replaces
+    the builder fully owns its own grouping.
+    """
+
+    @abstractmethod
+    def get_kv_cache_groups(
+        self,
+        vllm_config: "VllmConfig",
+        kv_cache_spec: dict[str, "KVCacheSpec"],
+    ) -> list["KVCacheGroupSpec"]:
+        """Split the given layers into KV cache groups.
+
+        Planning passes all workers' layers; profiling passes one worker's.
+        Core rejects a result that does not put each layer in exactly one group.
+
+        The default builder honors a model's
+        :class:`~vllm.model_executor.models.interfaces.SupportsKVCacheGrouping`
+        declaration here before falling back to the default grouping; a
+        platform-specific builder decides its own grouping.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_kv_cache_config_from_groups(
+        self,
+        vllm_config: "VllmConfig",
+        kv_cache_groups: list["KVCacheGroupSpec"],
+        num_blocks: int,
+    ) -> "KVCacheConfig":
+        """Materialize the groups for exactly ``num_blocks`` blocks.
+
+        Each device backing's size is ``bytes_per_block * num_blocks``. All
+        device tensors share one backing allocation unless the builder
+        overrides ``_get_pool_bytes_per_block``. Host-resident tensors, such as
+        HiSparse source storage, are excluded from the device byte accounting
+        and sized independently.
+        """
+        raise NotImplementedError
+
+    def _get_pool_bytes_per_block(
+        self,
+        vllm_config: "VllmConfig",
+        kv_cache_groups: list["KVCacheGroupSpec"],
+    ) -> int | None:
+        """Return the device bytes one block ID consumes, or ``None``.
+
+        Private and transitional; expected to go away with worker-driven
+        allocation (RFC #42449). With ``None`` (the default), Core derives the
+        value from a one-block placement with one device backing. Override it
+        only for placements with several device allocations.
+        """
+        return None
+
+
+def get_kv_cache_config_builder(
+    vllm_config: "VllmConfig",
+) -> KVCacheConfigBuilder:
+    """Resolve the builder selected by the current platform.
+
+    Resolution priority is owned by the platform hook
+    (:meth:`vllm.platforms.interface.Platform.get_kv_cache_config_builder_cls`).
+    Builders are stateless (every method takes ``vllm_config``), so a fresh
+    instance is returned per call and engines with different configs in the
+    same process always get the correct builder.
+    """
+    from vllm.platforms import current_platform
+
+    builder_cls = resolve_obj_by_qualname(
+        current_platform.get_kv_cache_config_builder_cls(vllm_config)
+    )
+    return builder_cls()

vllm/v1/core/kv_cache_planning.py

@@ -14,6 +14,7 @@ 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_config_builder import KVCacheConfigBuilder
 from vllm.v1.core.kv_cache_utils import _kv_cache_device_name
 from vllm.v1.hisparse.layout import (
     get_hisparse_host_pool_bytes,
@@ -261,6 +262,48 @@ def _pool_bytes_per_block(
     )
 
 
+def _pool_bytes_per_block_from_builder(
+    builder: KVCacheConfigBuilder,
+    vllm_config: VllmConfig,
+    kv_cache_groups: list[KVCacheGroupSpec],
+) -> int:
+    """Bytes consumed by one block in the worker's device KV cache pool.
+
+    Core places one block with the builder and reads the size of its single
+    device backing. A placement with several device allocations declares the
+    value through ``KVCacheConfigBuilder._get_pool_bytes_per_block``. HiSparse
+    keeps the pure computation, since probing its placement would check the host
+    pool against ``max_model_len`` before auto-fit shortens it.
+    ``num_gpu_blocks_override`` is suppressed so it cannot inflate the probe.
+    """
+    if not kv_cache_groups:
+        return 0
+    if vllm_config.attention_config.hisparse_config is not None:
+        return _pool_bytes_per_block(vllm_config, kv_cache_groups)
+    pool_bytes = builder._get_pool_bytes_per_block(vllm_config, kv_cache_groups)
+    if pool_bytes is not None:
+        return pool_bytes
+    saved_override = vllm_config.cache_config.num_gpu_blocks_override
+    vllm_config.cache_config.num_gpu_blocks_override = None
+    try:
+        probe = builder.get_kv_cache_config_from_groups(
+            vllm_config, kv_cache_groups, num_blocks=1
+        )
+    finally:
+        vllm_config.cache_config.num_gpu_blocks_override = saved_override
+    backing_sizes = {
+        tensor.size for tensor in probe.kv_cache_tensors if not tensor.host_resident
+    }
+    if len(backing_sizes) != 1:
+        raise ValueError(
+            "KV cache tensors must share one device backing allocation, but got "
+            f"sizes {sorted(backing_sizes)}. Override "
+            "KVCacheConfigBuilder._get_pool_bytes_per_block() for placements "
+            "with several device allocations."
+        )
+    return backing_sizes.pop()
+
+
 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}
@@ -1264,6 +1307,36 @@ def get_kv_cache_groups(
     return groups
 
 
+def get_kv_cache_groups_from_builder(
+    builder: KVCacheConfigBuilder,
+    vllm_config: VllmConfig,
+    kv_cache_spec: dict[str, KVCacheSpec],
+) -> list[KVCacheGroupSpec]:
+    """Group the layers with the builder and check the result.
+
+    Core checks every grouping, whether it comes from the default grouping, a
+    model's ``SupportsKVCacheGrouping`` or a platform builder.
+    """
+    groups = builder.get_kv_cache_groups(vllm_config, kv_cache_spec)
+    _check_exact_partition(groups, kv_cache_spec)
+    return groups
+
+
+def _check_exact_partition(
+    groups: list[KVCacheGroupSpec],
+    kv_cache_spec: dict[str, KVCacheSpec],
+) -> None:
+    """Ensure the groups partition every input layer exactly once."""
+    grouped_layers: set[str] = set()
+    for group in groups:
+        assert not set(group.layer_names) & grouped_layers, (
+            f"KV cache group overlaps: {set(group.layer_names) & grouped_layers}."
+        )
+        grouped_layers |= set(group.layer_names)
+    missing = set(kv_cache_spec) - grouped_layers
+    assert not missing, f"KV cache groups miss layers: {sorted(missing)}."
+
+
 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
@@ -1285,6 +1358,7 @@ def kv_cache_groups_tp_replicas(
 
 
 def _max_memory_usage_bytes_from_groups(
+    builder: KVCacheConfigBuilder,
     vllm_config: VllmConfig,
     kv_cache_groups: list[KVCacheGroupSpec],
 ) -> int:
@@ -1301,7 +1375,9 @@ def _max_memory_usage_bytes_from_groups(
     if not kv_cache_groups:
         return 0
 
-    bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups)
+    bytes_per_block = _pool_bytes_per_block_from_builder(
+        builder, vllm_config, kv_cache_groups
+    )
     total_blocks = 0
     for group in kv_cache_groups:
         spec = group.kv_cache_spec
@@ -1317,6 +1393,7 @@ def _max_memory_usage_bytes_from_groups(
 
 
 def _estimate_max_model_len_from_groups(
+    builder: KVCacheConfigBuilder,
     vllm_config: VllmConfig,
     kv_cache_groups: list[KVCacheGroupSpec],
     available_memory: int,
@@ -1334,13 +1411,16 @@ def _estimate_max_model_len_from_groups(
         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
+                num_blocks = available_memory // _pool_bytes_per_block_from_builder(
+                    builder, vllm_config, kv_cache_groups
+                )
+                builder.get_kv_cache_config_from_groups(
+                    vllm_config, kv_cache_groups, num_blocks
                 )
             except ValueError:
                 return False
         return (
-            _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups)
+            _max_memory_usage_bytes_from_groups(builder, vllm_config, kv_cache_groups)
             <= available_memory
         )
 
@@ -1362,6 +1442,7 @@ def _estimate_max_model_len_from_groups(
 
 
 def _auto_fit_max_model_len(
+    builder: KVCacheConfigBuilder,
     vllm_config: VllmConfig,
     projected_groups_per_worker: list[list[KVCacheGroupSpec]],
     available_memory: list[int],
@@ -1372,6 +1453,7 @@ def _auto_fit_max_model_len(
     workers.
 
     Args:
+        builder: The builder whose hooks drive block accounting.
         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
@@ -1395,7 +1477,9 @@ def _auto_fit_max_model_len(
     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)
+        worker_max = _estimate_max_model_len_from_groups(
+            builder, vllm_config, groups, avail_mem
+        )
         if worker_max < auto_fit_max:
             auto_fit_max = worker_max
             limiting_worker_mem = avail_mem
@@ -1473,6 +1557,9 @@ def _project_kv_cache_groups_to_worker(
                 worker_layer_names,
                 group_spec,
                 is_eagle_group=group.is_eagle_group and bool(worker_layer_names),
+                host_resident=group.host_resident,
+                enable_kv_transfer=group.enable_kv_transfer,
+                role=group.role,
             )
         )
     return projected_groups
@@ -1512,6 +1599,10 @@ def get_kv_cache_configs(
         The generated KVCacheConfigs for each worker.
 
     """
+    from vllm.v1.core.kv_cache_config_builder import get_kv_cache_config_builder
+
+    builder = get_kv_cache_config_builder(vllm_config)
+
     # 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.
@@ -1542,7 +1633,9 @@ def get_kv_cache_configs(
     # 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)
+    global_kv_cache_groups = get_kv_cache_groups_from_builder(
+        builder, 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.
@@ -1553,11 +1646,10 @@ def get_kv_cache_configs(
     ]
 
     # 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.
+    # be allocated is decoupled from the profiled `available_memory`. 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] = []
@@ -1565,7 +1657,9 @@ def get_kv_cache_configs(
             if not groups:
                 adjusted_memory.append(avail_mem)
                 continue
-            bytes_per_block = _pool_bytes_per_block(vllm_config, groups)
+            bytes_per_block = _pool_bytes_per_block_from_builder(
+                builder, vllm_config, groups
+            )
             logger.info(
                 "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d",
                 avail_mem // bytes_per_block,
@@ -1578,12 +1672,18 @@ def get_kv_cache_configs(
     # 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
+        (
+            avail_mem - _pool_bytes_per_block_from_builder(builder, 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)
+        _auto_fit_max_model_len(
+            builder, 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):
@@ -1591,9 +1691,9 @@ def get_kv_cache_configs(
             continue
         _check_enough_kv_cache_memory(
             avail_mem,
-            partial(_max_memory_usage_bytes_from_groups, vllm_config, groups),
+            partial(_max_memory_usage_bytes_from_groups, builder, vllm_config, groups),
             vllm_config.model_config.max_model_len,
-            partial(_estimate_max_model_len_from_groups, vllm_config, groups),
+            partial(_estimate_max_model_len_from_groups, builder, vllm_config, groups),
         )
 
     kv_cache_configs: list[KVCacheConfig] = []
@@ -1603,9 +1703,15 @@ def get_kv_cache_configs(
         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."
+        bytes_per_block = _pool_bytes_per_block_from_builder(
+            builder, vllm_config, projected_groups
+        )
+        num_blocks = (
+            available_memory_one_worker // bytes_per_block if projected_groups else 1
+        )
         kv_cache_configs.append(
-            get_kv_cache_config_from_groups(
-                vllm_config, projected_groups, available_memory_one_worker
+            builder.get_kv_cache_config_from_groups(
+                vllm_config, projected_groups, num_blocks
             )
         )
 
@@ -1618,13 +1724,12 @@ def get_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(
+        # Re-plan with exactly the smallest rank's block count, so strides and
+        # offsets stay consistent with the shrunken allocation.
+        kv_cache_configs[i] = builder.get_kv_cache_config_from_groups(
             vllm_config,
-            groups,
-            min_num_blocks * _pool_bytes_per_block(vllm_config, groups),
+            kv_cache_config.kv_cache_groups,
+            min_num_blocks,
         )
 
     for kv_cache_config in kv_cache_configs:
@@ -1635,3 +1740,64 @@ def get_kv_cache_configs(
         )
 
     return kv_cache_configs
+
+
+class DefaultKVCacheConfigBuilder(KVCacheConfigBuilder):
+    """Default grouping and placement, backed by the module-level functions here.
+
+    .. note::
+        This class is part of the experimental ``KVCacheConfigBuilder``
+        interface and is subject to change.
+
+    Subclasses override only the hooks they need. The planning flow and its
+    invariants (``num_gpu_blocks_override``, null-block reservation, auto-fit,
+    admission checks, cross-rank convergence) stay in the Core-owned
+    :func:`get_kv_cache_configs`.
+    """
+
+    def get_kv_cache_groups(
+        self,
+        vllm_config: VllmConfig,
+        kv_cache_spec: dict[str, KVCacheSpec],
+    ) -> list[KVCacheGroupSpec]:
+        """See ``KVCacheConfigBuilder.get_kv_cache_groups``.
+
+        A model implementing ``SupportsKVCacheGrouping`` owns the grouping
+        decision; otherwise the default grouping is used.
+        """
+        if not vllm_config.model_config.supports_kv_cache_grouping:
+            return get_kv_cache_groups(vllm_config, kv_cache_spec)
+        if vllm_config.attention_config.hisparse_config is not None:
+            raise ValueError(
+                "A model-declared get_kv_cache_groups cannot be combined "
+                "with HiSparse, which owns its own KV cache grouping. "
+                "Remove one of them."
+            )
+        model_cls, _ = vllm_config.model_config.registry.resolve_model_cls(
+            vllm_config.model_config.architectures,
+            vllm_config.model_config,
+        )
+        return model_cls.get_kv_cache_groups(vllm_config, kv_cache_spec)
+
+    def get_kv_cache_config_from_groups(
+        self,
+        vllm_config: VllmConfig,
+        kv_cache_groups: list[KVCacheGroupSpec],
+        num_blocks: int,
+    ) -> KVCacheConfig:
+        """See ``KVCacheConfigBuilder.get_kv_cache_config_from_groups``."""
+        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
+                ),
+            )
+        bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups)
+        return get_kv_cache_config_from_groups(
+            vllm_config, kv_cache_groups, num_blocks * bytes_per_block
+        )

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -1036,26 +1036,22 @@ 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_planning import (
-        get_kv_cache_config_from_groups,
-        get_kv_cache_groups,
-    )
+    from vllm.v1.core.kv_cache_config_builder import get_kv_cache_config_builder
+    from vllm.v1.core.kv_cache_planning import get_kv_cache_groups_from_builder
 
     kv_cache_spec = runner.get_kv_cache_spec()
-    kv_cache_groups = get_kv_cache_groups(runner.vllm_config, kv_cache_spec)
+    builder = get_kv_cache_config_builder(runner.vllm_config)
+    kv_cache_groups = get_kv_cache_groups_from_builder(
+        builder, runner.vllm_config, kv_cache_spec
+    )
     # At least one block per sequence is required to capture the graphs.
     min_blocks = (
         min(runner.max_num_reqs, runner.compilation_config.max_cudagraph_capture_size)
         or 1
     )
-    saved_override = runner.cache_config.num_gpu_blocks_override
-    runner.cache_config.num_gpu_blocks_override = min_blocks
-    try:
-        minimal_config = get_kv_cache_config_from_groups(
-            runner.vllm_config, kv_cache_groups, available_memory=0
-        )
-    finally:
-        runner.cache_config.num_gpu_blocks_override = saved_override
+    minimal_config = builder.get_kv_cache_config_from_groups(
+        runner.vllm_config, kv_cache_groups, num_blocks=min_blocks
+    )
 
     runner.initialize_kv_cache(minimal_config, is_profiling=True)
     runner.cache_config.num_gpu_blocks = minimal_config.num_blocks

vllm/v1/worker/gpu_model_runner.py

@@ -6174,27 +6174,23 @@ class GPUModelRunner(
         gc.collect()
 
     def _init_minimal_kv_cache_for_profiling(self) -> None:
-        from vllm.v1.core.kv_cache_planning import (
-            get_kv_cache_config_from_groups,
-            get_kv_cache_groups,
-        )
+        from vllm.v1.core.kv_cache_config_builder import get_kv_cache_config_builder
+        from vllm.v1.core.kv_cache_planning import get_kv_cache_groups_from_builder
 
         kv_cache_spec = self.get_kv_cache_spec()
         KVCacheSpecRegistry.check_kv_cache_spec_registry(kv_cache_spec)
-        kv_cache_groups = get_kv_cache_groups(self.vllm_config, kv_cache_spec)
+        builder = get_kv_cache_config_builder(self.vllm_config)
+        kv_cache_groups = get_kv_cache_groups_from_builder(
+            builder, self.vllm_config, kv_cache_spec
+        )
         # the minimum number of blocks required is 1 block *per sequence*
         min_blocks = (
             min(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size)
             or 1
         )
-
-        # Temporarily change num_gpu_blocks_override to allocate a minimal KV cache
-        saved_override = self.cache_config.num_gpu_blocks_override
-        self.cache_config.num_gpu_blocks_override = min_blocks
-        minimal_config = get_kv_cache_config_from_groups(
-            self.vllm_config, kv_cache_groups, available_memory=0
+        minimal_config = builder.get_kv_cache_config_from_groups(
+            self.vllm_config, kv_cache_groups, num_blocks=min_blocks
         )
-        self.cache_config.num_gpu_blocks_override = saved_override
 
         self.initialize_kv_cache(minimal_config, is_profiling=True)
         self.cache_config.num_gpu_blocks = minimal_config.num_blocks