From b5445874080ce85c575f2bcb847852235f1af7bc Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Sat, 17 Oct 2026 07:16:47 +0800 Subject: [PATCH] [Core][Feat] Pluggable KVCacheConfigBuilder for platform/model-specific KV cache planning (#53558) Signed-off-by: wangxiyuan Signed-off-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 Co-authored-by: Yifan Qiao --- tests/v1/core/test_kv_cache_config_builder.py | 294 ++++++++++++++++++ vllm/config/model.py | 4 + vllm/model_executor/models/interfaces.py | 42 +++ vllm/model_executor/models/registry.py | 3 + vllm/platforms/interface.py | 22 ++ vllm/v1/core/kv_cache_config_builder.py | 103 ++++++ vllm/v1/core/kv_cache_planning.py | 214 +++++++++++-- vllm/v1/worker/gpu/cudagraph_utils.py | 22 +- vllm/v1/worker/gpu_model_runner.py | 20 +- 9 files changed, 675 insertions(+), 49 deletions(-) create mode 100644 tests/v1/core/test_kv_cache_config_builder.py create mode 100644 vllm/v1/core/kv_cache_config_builder.py diff --git a/tests/v1/core/test_kv_cache_config_builder.py b/tests/v1/core/test_kv_cache_config_builder.py new file mode 100644 index 0000000000..4f4c8d0262 --- /dev/null +++ b/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) diff --git a/vllm/config/model.py b/vllm/config/model.py index 7a883eac54..6dc2d16d52 100644 --- a/vllm/config/model.py +++ b/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 diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 5f51ece0b3..818c161050 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/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.""" diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 23e67c9f50..50280061e2 100644 --- a/vllm/model_executor/models/registry.py +++ b/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", () diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 3034d7bc1a..65f165ad3a 100644 --- a/vllm/platforms/interface.py +++ b/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 diff --git a/vllm/v1/core/kv_cache_config_builder.py b/vllm/v1/core/kv_cache_config_builder.py new file mode 100644 index 0000000000..09332946cc --- /dev/null +++ b/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() diff --git a/vllm/v1/core/kv_cache_planning.py b/vllm/v1/core/kv_cache_planning.py index 6e258a68a9..10016409af 100644 --- a/vllm/v1/core/kv_cache_planning.py +++ b/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 + ) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index a1955b09c7..4a1dd0fbfb 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index f18e219d8e..3f1c3fe4c0 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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