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 f597c8fa27c8..6f7f4b6967c0 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -166,6 +166,22 @@ def _mla_spec( ) +def _swa_spec(sliding_window: int = 128) -> SlidingWindowSpec: + return SlidingWindowSpec( + block_size=16, + num_kv_heads=4, + head_size=128, + dtype=torch.float32, + sliding_window=sliding_window, + ) + + +_HIDDEN_STATE_KWARGS: dict[str, Any] = { + "block_size": 16, + "num_kv_heads": 1, + "head_size": 512, + "dtype": torch.float32, +} _MAMBA_SPEC = MambaSpec( block_size=16, shapes=((16, 1),), @@ -512,6 +528,55 @@ def test_dcp_scales_uniform_type_group_alongside_mamba(spec_kind, expected): ) +@pytest.mark.parametrize( + "kv_cache_spec,window_chunks", + [ + (_full_attention_spec(), None), + (_mla_spec(), None), + (HiddenStateCacheSpec(**_HIDDEN_STATE_KWARGS), None), + (_MAMBA_SPEC, 1), + (_swa_spec(), 8), + ( + UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={ + "mla": _mla_spec(), + "hidden": HiddenStateCacheSpec(**_HIDDEN_STATE_KWARGS), + }, + ), + None, + ), + ( + UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={"swa0": _swa_spec(), "swa1": _swa_spec()}, + ), + 8, + ), + ], +) +def test_group_carries_the_chunks_the_tier_keeps_per_request( + kv_cache_spec: KVCacheSpec, window_chunks: int | None +): + """The group reports how far back its attention reaches, counted in chunks. + + A capacity estimate needs that bound: a windowed or recurrent group keeps a + fixed number of chunks however long the request grows, while an unbounded + group keeps one for every chunk of the request. A worker-side group hands + over an aggregate of its layers, which must resolve to the same bound as the + single layer it wraps. + """ + kv_cache_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)], + ) + + offloading_config = build_offloading_config(_make_vllm_config(), kv_cache_config) + + assert offloading_config.groups[0].sliding_window_size_in_chunks == window_chunks + + def test_preserves_data_parallel_config(): config = _make_vllm_config() config.parallel_config.data_parallel_index = 2 @@ -795,13 +860,7 @@ def test_parallelism_agnostic_for_single_full_attention_group(): assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())]) -_SWA_SPEC = SlidingWindowSpec( - block_size=16, - num_kv_heads=4, - head_size=128, - dtype=torch.float32, - sliding_window=128, -) +_SWA_SPEC = _swa_spec() _SWA_MLA_SPEC = SlidingWindowMLASpec( block_size=16, num_kv_heads=1, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py index b7d9d6c8efac..1a185ae0dd54 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Mapping +from dataclasses import dataclass from types import SimpleNamespace from typing import Any from unittest.mock import patch @@ -20,9 +22,15 @@ OffloadingConnector, ) from vllm.v1.kv_offload.base import ( + KV_OFFLOAD_CONFIG_INFO, + CanonicalKVCaches, + OffloadingConfigInfo, OffloadingCounterMetadata, OffloadingGaugeMetadata, OffloadingHistogramMetadata, + OffloadingManager, + OffloadingSpec, + OffloadingWorker, ) from vllm.v1.kv_offload.factory import OffloadingSpecFactory from vllm.v1.kv_offload.tiering.base import TieringOffloadingMetrics @@ -124,6 +132,10 @@ class _FakeOffloadingSpec: def build_metric_definitions(extra_config): return metric_definitions + @staticmethod + def build_info_metric_definition(extra_config): + return {} + return _FakeOffloadingSpec @@ -638,6 +650,98 @@ def test_prom_metrics_lazily_observes_labeled_metric(): assert counter_def.kwargs["labelnames"] == ["model_name", "engine", MY_LABEL] +def test_prom_metrics_declares_the_info_gauge_as_most_recent(): + """The info gauge must merge across frontends by freshest write. + + Offloading stats reach one frontend per step as complete per-engine + snapshots, so summing would report the number of participating API-server + processes instead of the 1 an info gauge is pinned to -- wrong only under + multiprocess deployment, and silently, since the facts ride the labels. + """ + prom_metrics = OffloadPromMetrics( + vllm_config=_FakeVllmConfig(), # type: ignore[arg-type] + metric_types={ + Gauge: _FakeMetric, + Counter: _FakeMetric, + Histogram: _FakeMetric, + }, + labelnames=["model_name", "engine"], + per_engine_labelvalues={0: ["model", "0"]}, + ) + + gauge_def = prom_metrics._offloading_metric_defs[KV_OFFLOAD_CONFIG_INFO] + assert gauge_def.kwargs["multiprocess_mode"] == "mostrecent" + + +def test_prom_metrics_keeps_a_gauge_declared_multiprocess_mode(): + """A gauge that declares its own multiprocess_mode keeps it. + + The value on OffloadingGaugeMetadata must stay a default. A hardcoded + "mostrecent" in _create_metric passes the test above and silently overrides + a gauge that needs another merge, for example a real cross-frontend sum. + """ + metric_definitions = { + PENDING_STORES: OffloadingGaugeMetadata( + documentation="gauge declaring a non-default multiprocess mode.", + multiprocess_mode="sum", + ), + } + with patch.object( + OffloadingSpecFactory, + "get_spec_cls", + return_value=_spec_cls_with_metric_definitions(metric_definitions), + ): + prom_metrics = OffloadPromMetrics( + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type] + metric_types={ + Gauge: _FakeMetric, + Counter: _FakeMetric, + Histogram: _FakeMetric, + }, + labelnames=["model_name", "engine"], + per_engine_labelvalues={0: ["model", "0"]}, + ) + + gauge_def = prom_metrics._offloading_metric_defs[PENDING_STORES] + assert gauge_def.kwargs["multiprocess_mode"] == "sum" + + +def test_prom_metrics_omits_multiprocess_mode_outside_a_gauge(): + """multiprocess_mode reaches a gauge only. + + prometheus_client accepts the argument on Gauge alone, so hoisting the + kwarg out of the gauge branch breaks construction of a real Counter or + Histogram. + """ + metric_definitions = { + MY_COUNTER: OffloadingCounterMetadata(documentation="a counter."), + LOOKUP_LATENCY: OffloadingHistogramMetadata( + documentation="a histogram.", + buckets=(0.1, 1), + ), + } + with patch.object( + OffloadingSpecFactory, + "get_spec_cls", + return_value=_spec_cls_with_metric_definitions(metric_definitions), + ): + prom_metrics = OffloadPromMetrics( + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type] + metric_types={ + Gauge: _FakeMetric, + Counter: _FakeMetric, + Histogram: _FakeMetric, + }, + labelnames=["model_name", "engine"], + per_engine_labelvalues={0: ["model", "0"]}, + ) + + counter_def = prom_metrics._offloading_metric_defs[MY_COUNTER] + histogram_def = prom_metrics._offloading_metric_defs[LOOKUP_LATENCY] + assert "multiprocess_mode" not in counter_def.kwargs + assert "multiprocess_mode" not in histogram_def.kwargs + + def test_prom_metrics_rejects_wrong_label_count(): metric_definitions = { MY_COUNTER: OffloadingCounterMetadata( @@ -784,3 +888,102 @@ def test_prom_metrics_rejects_undeclared_metric(): _StatsKey.DATA: {"unknown:metric": {(): 1}}, } ) + + +def test_scheduler_stats_report_the_info_metric(request_runner): + """The info metric appears with no labels, because MockOffloadingSpec + declares no tiers.""" + runner = request_runner( + block_size=4, + num_gpu_blocks=10, + async_scheduling=False, + ) + + stats = runner.connector_scheduler.get_stats() + + assert stats is not None + assert stats.data[_StatsKey.TYPES][KV_OFFLOAD_CONFIG_INFO] == _MetricType.GAUGE + assert stats.data[_StatsKey.DATA][KV_OFFLOAD_CONFIG_INFO] == {(): 1} + + +@dataclass(frozen=True) +class _CPUCacheInfo(OffloadingConfigInfo): + num_chunks: int + policy: str + + +@dataclass(frozen=True) +class _FsTierInfo(OffloadingConfigInfo): + path: str + + +class _MultiSourceOffloadingSpec(OffloadingSpec): + """Spec with a CPU cache and two file-system tiers. + + This is the shape a tiering spec reports once every source publishes its own + configuration. + """ + + def __init__(self): + self.extra_config: Mapping[str, Any] = {} + + def get_manager(self) -> OffloadingManager: + raise NotImplementedError + + def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker: + raise NotImplementedError + + @classmethod + def config_info_classes( + cls, extra_config: Mapping[str, Any] + ) -> tuple[tuple[str, type[OffloadingConfigInfo]], ...]: + return (("cpu", _CPUCacheInfo), ("fs", _FsTierInfo), ("fs", _FsTierInfo)) + + def config_info(self) -> tuple[OffloadingConfigInfo, ...]: + return ( + _CPUCacheInfo(num_chunks=512, policy="lru"), + _FsTierInfo(path="/mnt/a"), + _FsTierInfo(path="/mnt/b"), + ) + + +def test_prom_metrics_binds_the_info_labels_of_every_config_source(): + """The spec class declares the label names in the API-server process, the + spec instance reports the values in the engine process, and Prometheus binds + the two by position.""" + with patch.object( + OffloadingSpecFactory, + "get_spec_cls", + return_value=_MultiSourceOffloadingSpec, + ): + prom_metrics = OffloadPromMetrics( + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type] + metric_types={ + Gauge: _FakeMetric, + Counter: _FakeMetric, + Histogram: _FakeMetric, + }, + labelnames=["model_name", "engine"], + per_engine_labelvalues={0: ["model", "0"]}, + ) + + gauge_def = prom_metrics._offloading_metric_defs[KV_OFFLOAD_CONFIG_INFO] + assert gauge_def.kwargs["labelnames"] == [ + "model_name", + "engine", + "cpu0_num_chunks", + "cpu0_policy", + "fs1_path", + "fs2_path", + ] + + stats = OffloadingConnectorStats() + stats.set_gauge( + KV_OFFLOAD_CONFIG_INFO, 1, _MultiSourceOffloadingSpec().info_labelvalues() + ) + prom_metrics.observe(stats.data) + + labelvalues = ("512", "lru", "/mnt/a", "/mnt/b") + gauge = prom_metrics.offloading_metrics[(0, KV_OFFLOAD_CONFIG_INFO, labelvalues)] + assert gauge.set_values == [1] + assert gauge.labelvalues == ("model", "0") + labelvalues diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index e7a83cb7b507..49f0de2cea22 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -114,6 +114,7 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state): enable_kv_cache_events=False, self_describing_kv_events=False ), get_manager=lambda: manager, + info_labelvalues=lambda: (), ) config = SimpleNamespace( speculative_config=None, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 366624e1d3b1..cda8bbb6ccc9 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Mapping +from typing import Any from unittest.mock import MagicMock import pytest @@ -27,6 +29,7 @@ CanonicalKVCaches, GPULoadStoreSpec, LoadStoreSpec, + OffloadingConfigInfo, OffloadingManager, OffloadingSpec, OffloadingWorker, @@ -182,6 +185,15 @@ def get_manager(self) -> OffloadingManager: def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker: raise NotImplementedError + @classmethod + def config_info_classes( + cls, extra_config: Mapping[str, Any] + ) -> tuple[tuple[str, type[OffloadingConfigInfo]], ...]: + return () + + def config_info(self) -> tuple[OffloadingConfigInfo, ...]: + return () + # --------------------------------------------------------------------------- # Tests diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 5707003ce935..06d1db46fc43 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass from typing import Any from unittest.mock import MagicMock @@ -47,6 +47,7 @@ GPULoadStoreSpec, LoadStoreSpec, LookupResult, + OffloadingConfigInfo, OffloadingManager, OffloadingSpec, OffloadingWorker, @@ -139,6 +140,15 @@ def get_manager(self) -> OffloadingManager: def get_worker(self, _: CanonicalKVCaches) -> OffloadingWorker: return self.handler + @classmethod + def config_info_classes( + cls, extra_config: Mapping[str, Any] + ) -> tuple[tuple[str, type[OffloadingConfigInfo]], ...]: + return () + + def config_info(self) -> tuple[OffloadingConfigInfo, ...]: + return () + def complete_transfers(self): self.handler.complete_jobs(self.handler.waiting_jobs.copy()) diff --git a/tests/v1/kv_offload/cpu/test_capacity.py b/tests/v1/kv_offload/cpu/test_capacity.py new file mode 100644 index 000000000000..9f9cff259fd3 --- /dev/null +++ b/tests/v1/kv_offload/cpu/test_capacity.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the CPU offloading tier capacity estimate. + +The estimate answers one question: how many KV tokens does a full tier serve +at the best request length up to max_model_len? Both functions under test take +plain arguments, so these tests need no spec and no engine. +""" + +import pytest + +from vllm.v1.kv_offload.config import OffloadingGroupConfig +from vllm.v1.kv_offload.cpu.spec import ( + _capacity_tokens_at_max_len, + _chunks_per_request, +) + +MAX_MODEL_LEN = 36864 + +# Group shapes as (tokens_per_block, sliding_window_size_in_chunks) pairs. +FULL_16 = ((16, None),) +# One full attention group and 6 sliding window groups of 1024 tokens. +GEMMA = ((16, None),) + ((16, 4),) * 6 +# One full attention group and 9 Mamba groups of one state each. +GRANITE = ((528, None),) + ((528, 1),) * 9 +# One full attention group and one hidden state group on a smaller block. +FULL_PLUS_HIDDEN = ((16, None), (8, None)) +# One group of each kind: full attention, a 1024 token window, and Mamba. +HYBRID = ((16, None), (16, 4), (16, 1)) + + +def make_groups( + *shapes: tuple[int, int | None], +) -> tuple[OffloadingGroupConfig, ...]: + """Build a group tuple from (tokens_per_block, window_chunks) pairs. + + Args: + shapes: One pair for each group of the model. + + Returns: + The groups, each with a distinct single layer name. + """ + return tuple( + OffloadingGroupConfig( + tokens_per_block=tokens_per_block, + layer_names=(f"layer.{index}",), + group_id=index, + sliding_window_size_in_chunks=window_chunks, + ) + for index, (tokens_per_block, window_chunks) in enumerate(shapes) + ) + + +@pytest.mark.parametrize( + ("shapes", "blocks_per_chunk", "num_chunks", "expected"), + [ + # 4500 * 36864 // 144 = 1152000 + pytest.param(FULL_16, 16, 4500, 1152000, id="single-full-group"), + # 4000 * 36864 // min(144, 16) = 9216000 + pytest.param(((16, 16),), 16, 4000, 9216000, id="uniform-swa-4096"), + # 4000 * 36864 // min(144, 4) = 36864000 + pytest.param(((16, 4),), 16, 4000, 36864000, id="uniform-swa-1024"), + # 4000 * 36864 // (144 + 6 * 4) = 877714 + pytest.param(GEMMA, 16, 4000, 877714, id="full-plus-6-swa-1024"), + # 8000 * 36432 // (69 + 9 * 1) = 3736615, at the last full chunk + pytest.param(GRANITE, 1, 8000, 3736615, id="full-plus-9-mamba"), + # 8000 * 36864 // min(70, 1) = 294912000 + pytest.param(((528, 1),), 1, 8000, 294912000, id="pure-mamba"), + # 4500 * 36864 // (144 + 288) = 384000 + pytest.param(FULL_PLUS_HIDDEN, 16, 4500, 384000, id="full-plus-hidden"), + # 4000 * 36864 // (144 + 4 + 1) = 989637 + pytest.param(HYBRID, 16, 4000, 989637, id="full-plus-swa-plus-mamba"), + ], +) +def test_capacity_over_model_shapes( + shapes: tuple[tuple[int, int | None], ...], + blocks_per_chunk: int, + num_chunks: int, + expected: int, +) -> None: + """The estimate reports the reviewed capacity of each model shape. + + A window raises the capacity above the byte count of the tier, because a + windowed group holds a fixed chunk count however long the request grows. + """ + capacity = _capacity_tokens_at_max_len( + make_groups(*shapes), blocks_per_chunk, num_chunks, MAX_MODEL_LEN + ) + + assert capacity == expected + + +@pytest.mark.parametrize( + ("tokens_per_block", "blocks_per_chunk", "num_chunks"), + [ + # 4500 * 36864 // 144 = 4500 * 256 = 1152000 + pytest.param(16, 16, 4500, id="qwen3-8b-shape"), + # 8000 * 36432 // 69 = 8000 * 528 = 4224000, at the last full chunk + pytest.param(528, 1, 8000, id="chunk-size-divides-max-len-with-remainder"), + # 4500 * 36864 // 288 = 4500 * 128 = 576000 + pytest.param(8, 16, 4500, id="small-block"), + ], +) +def test_single_uncapped_group_keeps_the_exact_token_count( + tokens_per_block: int, blocks_per_chunk: int, num_chunks: int +) -> None: + """One group without a window keeps the exact count of the old metric. + + The old capacity_tokens divided the byte budget by the byte cost of one + token. The estimate must agree with it on this shape. + """ + groups = make_groups((tokens_per_block, None)) + + capacity = _capacity_tokens_at_max_len( + groups, blocks_per_chunk, num_chunks, MAX_MODEL_LEN + ) + + assert capacity == num_chunks * blocks_per_chunk * tokens_per_block + + +def test_empty_tier_reports_an_exact_zero() -> None: + """A tier of no slots serves no tokens, which is a known value.""" + capacity = _capacity_tokens_at_max_len(make_groups(*FULL_16), 16, 0, MAX_MODEL_LEN) + + assert capacity == 0 + + +@pytest.mark.parametrize( + ("shapes", "max_model_len"), + [ + # max_model_len 0 leaves no length to divide, so None + pytest.param(FULL_16, 0, id="max-model-len-not-known"), + # no group leaves the tokens_per_chunk set empty, so None + pytest.param((), MAX_MODEL_LEN, id="no-groups"), + # 16 * 0 = 0 tokens for each chunk, so None + pytest.param(((0, None),), MAX_MODEL_LEN, id="blocks-span-no-tokens"), + ], +) +def test_capacity_is_none_when_the_token_scale_is_unknown( + shapes: tuple[tuple[int, int | None], ...], max_model_len: int +) -> None: + """The estimate reports None when an input leaves the token scale open.""" + capacity = _capacity_tokens_at_max_len( + make_groups(*shapes), 16, 4500, max_model_len + ) + + assert capacity is None + + +@pytest.mark.parametrize( + ("shapes", "blocks_per_chunk", "seq_len", "expected"), + [ + # cdiv(36864, 16 * 16) = 144 + pytest.param(FULL_16, 16, 36864, 144, id="one-full-group"), + # 144 + 6 * min(144, 4) = 168 + pytest.param(GEMMA, 16, 36864, 168, id="full-plus-6-swa-1024"), + # cdiv(36432, 528) + 9 * min(69, 1) = 69 + 9 = 78 + pytest.param(GRANITE, 1, 36432, 78, id="full-plus-9-mamba"), + # min(cdiv(256, 256), 4) = 1 + pytest.param(((16, 4),), 16, 256, 1, id="request-below-the-window"), + # min(cdiv(36864, 256), 4) = 4 + pytest.param(((16, 4),), 16, 36864, 4, id="window-caps-the-request"), + ], +) +def test_chunks_per_request( + shapes: tuple[tuple[int, int | None], ...], + blocks_per_chunk: int, + seq_len: int, + expected: int, +) -> None: + """One request holds one chunk for each chunk-sized span, up to the cap.""" + chunks = _chunks_per_request(make_groups(*shapes), blocks_per_chunk, seq_len) + + assert chunks == expected + + +@pytest.mark.parametrize( + ("shapes", "blocks_per_chunk"), + [ + # 1000 * 4096 // (16 + 6 * 4) = 102400 + pytest.param(GEMMA, 16, id="full-plus-6-swa-1024"), + # 1000 * 3840 // (15 + 10) = 153600, at the 384-token floor + pytest.param(((16, None), (24, None)), 16, id="chunk-sizes-do-not-nest"), + ], +) +def test_candidate_lengths_reach_the_capacity_peak( + shapes: tuple[tuple[int, int | None], ...], blocks_per_chunk: int +) -> None: + """The estimate matches a sweep of every request length. + + The estimate measures a few candidate lengths and keeps the largest + result. The sweep shows that the candidate set holds the peak of the + capacity curve. A short max_model_len keeps the sweep fast. + """ + max_model_len = 4096 + num_chunks = 1000 + groups = make_groups(*shapes) + + peak = max( + num_chunks * seq_len // _chunks_per_request(groups, blocks_per_chunk, seq_len) + for seq_len in range(1, max_model_len + 1) + ) + + capacity = _capacity_tokens_at_max_len( + groups, blocks_per_chunk, num_chunks, max_model_len + ) + + assert capacity == peak diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index a79f480ff551..71f42b84a9a7 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -17,10 +17,7 @@ ReqContext, make_offload_key, ) -from vllm.v1.kv_offload.cpu.common import ( - CPULoadStoreSpec, - CPUOffloadingMetrics, -) +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec, CPUOffloadingMetrics from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 16539677c490..0a3a3f887287 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -2,13 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for native offloading specs and their factory.""" +from collections.abc import Mapping +from dataclasses import dataclass from typing import Any from unittest.mock import MagicMock import pytest from vllm.v1.kv_offload.base import ( + KV_OFFLOAD_CONFIG_INFO, CanonicalKVCaches, + OffloadingConfigInfo, + OffloadingGaugeMetadata, OffloadingHistogramMetadata, OffloadingManager, OffloadingSpec, @@ -21,6 +26,7 @@ OffloadingModelConfig, OffloadingParallelConfig, ) +from vllm.v1.kv_offload.cpu.common import CPUCacheOffloadingInfo from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec from vllm.v1.kv_offload.factory import OffloadingSpecFactory @@ -42,6 +48,7 @@ def _make_offloading_config( groups: tuple[OffloadingGroupConfig, ...] | None = None, tokens_per_hash: int = 16, blocks_per_chunk: int = 1, + max_model_len: int = 4096, rank: int = 0, world_size: int = 1, tp_size: int | None = None, @@ -70,7 +77,9 @@ def _make_offloading_config( enable_kv_cache_events=False, extra_config=normalized_extra_config, engine_id="test-engine", - model=OffloadingModelConfig(name="test-model", dtype="float16"), + model=OffloadingModelConfig( + name="test-model", dtype="float16", max_model_len=max_model_len + ), cache=OffloadingCacheConfig( tokens_per_hash=tokens_per_hash, blocks_per_chunk=blocks_per_chunk, @@ -102,6 +111,15 @@ def get_manager(self) -> OffloadingManager: def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker: raise NotImplementedError + @classmethod + def config_info_classes( + cls, extra_config: Mapping[str, Any] + ) -> tuple[tuple[str, type[OffloadingConfigInfo]], ...]: + return () + + def config_info(self) -> tuple[OffloadingConfigInfo, ...]: + return () + def test_pre_registered_specs_can_be_imported(): for name in OffloadingSpecFactory._registry: @@ -168,6 +186,141 @@ def test_cpu_spec_zero_worker_bytes_produces_empty_cache(): assert spec.num_chunks == 0 +@pytest.mark.parametrize("blocks_per_chunk", [1, 2, 4]) +def test_cpu_spec_tier_info_converts_slots_to_tokens(blocks_per_chunk: int): + """One uncapped group makes the capacity the slot count in KV tokens. + + A slot holds blocks_per_chunk blocks of tokens_per_block tokens each, so + dropping either factor understates the tier. + """ + alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + tokens_per_block = 8 + spec = _create_spec( + cpu_bytes_to_use=alignment * 12, + worker_kv_bytes_per_block=alignment, + blocks_per_chunk=blocks_per_chunk, + groups=(OffloadingGroupConfig(tokens_per_block, ("layer",), 0),), + ) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.num_chunks == 12 // blocks_per_chunk + assert spec.tier_info.capacity_tokens_at_max_len == ( + spec.num_chunks * blocks_per_chunk * tokens_per_block + ) + + +@pytest.mark.parametrize("world_size", [1, 2, 4]) +def test_cpu_spec_tier_info_capacity_accounts_for_tensor_parallel_copies( + world_size: int, +): + """A slot holds every worker's copy of the block, so capacity divides by TP. + + Without the num_copies factor a TP=4 tier would be reported at four times + the tokens it can hold. + """ + alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + tokens_per_block = 8 + spec = _create_spec( + cpu_bytes_to_use=alignment * 12, + worker_kv_bytes_per_block=alignment, + world_size=world_size, + groups=(OffloadingGroupConfig(tokens_per_block, ("layer",), 0),), + ) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.num_chunks == 12 // world_size + assert ( + spec.tier_info.capacity_tokens_at_max_len == spec.num_chunks * tokens_per_block + ) + + +def test_cpu_spec_tier_info_capacity_dedups_a_replicated_layout(monkeypatch): + """A replicated layout stores one copy, so capacity does not divide by TP. + + Same TP=4 sizing as the test above, which yields 3 slots; deduplicating to a + single copy yields 12. + """ + import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + + monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True) + alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + tokens_per_block = 8 + spec = _create_spec( + cpu_bytes_to_use=alignment * 12, + worker_kv_bytes_per_block=alignment, + world_size=4, + replicated_layout=True, + groups=(OffloadingGroupConfig(tokens_per_block, ("layer",), 0),), + ) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.num_chunks == 12 + assert spec.tier_info.capacity_tokens_at_max_len == 12 * tokens_per_block + + +def test_cpu_spec_tier_info_mirrors_spec_sizing(): + """The exported facts are the spec's own, not a second derivation.""" + alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + spec = _create_spec( + cpu_bytes_to_use=alignment * 5, + worker_kv_bytes_per_block=alignment, + blocks_per_chunk=1, + ) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.tier_info.num_chunks == spec.num_chunks + assert spec.tier_info.blocks_per_chunk == spec.blocks_per_chunk + assert spec.tier_info.kv_bytes_per_chunk == spec.kv_bytes_per_chunk + + +def test_cpu_spec_tier_info_zero_capacity_is_exact_not_unknown(): + """A tier sized to nothing holds zero tokens; that is known, not unknown.""" + spec = _create_spec( + worker_kv_bytes_per_block=0, + groups=(OffloadingGroupConfig(16, ("layer",), 0),), + ) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.num_chunks == 0 + assert spec.tier_info.capacity_tokens_at_max_len == 0 + + +def test_cpu_spec_tier_info_capacity_shrinks_when_a_second_group_shares_slots(): + """Two groups take two chunks for one request, so the capacity halves.""" + one_group = _create_spec(groups=(OffloadingGroupConfig(16, ("full_layer",), 0),)) + two_groups = _create_spec( + groups=( + OffloadingGroupConfig(16, ("full_layer",), 0), + OffloadingGroupConfig(16, ("swa_layer",), 1), + ), + ) + + assert isinstance(one_group, CPUOffloadingSpec) + assert isinstance(two_groups, CPUOffloadingSpec) + assert two_groups.num_chunks == one_group.num_chunks + assert two_groups.tier_info.capacity_tokens_at_max_len == ( + one_group.tier_info.capacity_tokens_at_max_len // 2 + ) + + +def test_cpu_spec_tier_info_no_token_capacity_without_a_kv_cache_group(): + """A model with no KV cache has no group whose chunk size to apply.""" + spec = _create_spec(groups=()) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.num_chunks > 0 + assert spec.tier_info.capacity_tokens_at_max_len is None + + +def test_cpu_spec_tier_info_no_token_capacity_without_max_model_len(): + """The capacity holds at max_model_len, so an unknown length gives None.""" + spec = _create_spec(max_model_len=0) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.num_chunks > 0 + assert spec.tier_info.capacity_tokens_at_max_len is None + + def test_tiering_spec_aligns_row_size(): alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT spec = _create_spec( @@ -592,3 +745,165 @@ def test_build_metric_definitions_returns_counter_at_threshold(): metrics = spec_cls.build_metric_definitions(extra_config) assert CPUOffloadingMetrics.STORES_SKIPPED in metrics + + +@dataclass(frozen=True) +class _CPUConfigInfo(OffloadingConfigInfo): + chunks: int + policy: str + + @classmethod + def help_text(cls) -> str: + return "chunks holds the capacity, policy holds the eviction policy." + + +@dataclass(frozen=True) +class _FsConfigInfo(OffloadingConfigInfo): + path: str + + @classmethod + def help_text(cls) -> str: + return "path holds the mount point." + + +class _ConfigInfoOffloadingSpec(OffloadingSpec): + """Test-only spec that declares three config sources, two with one name. + + The constructor takes the reported facts directly, because these tests + exercise the info hooks alone and need no offloading config. + """ + + def __init__(self, infos: tuple[OffloadingConfigInfo, ...]): + self._infos = infos + self.extra_config: Mapping[str, Any] = {} + + def get_manager(self) -> OffloadingManager: + raise NotImplementedError + + def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker: + raise NotImplementedError + + @classmethod + def config_info_classes( + cls, extra_config: Mapping[str, Any] + ) -> tuple[tuple[str, type[OffloadingConfigInfo]], ...]: + return (("cpu", _CPUConfigInfo), ("fs", _FsConfigInfo), ("fs", _FsConfigInfo)) + + def config_info(self) -> tuple[OffloadingConfigInfo, ...]: + return self._infos + + +def test_info_metric_prefixes_every_label_with_the_source_index(): + metadata = _ConfigInfoOffloadingSpec.build_info_metric_definition({})[ + KV_OFFLOAD_CONFIG_INFO + ] + + assert isinstance(metadata, OffloadingGaugeMetadata) + assert metadata.labelnames == ("cpu0_chunks", "cpu0_policy", "fs1_path", "fs2_path") + + +def test_info_metric_help_holds_one_title_per_source_name(): + metadata = _ConfigInfoOffloadingSpec.build_info_metric_definition({})[ + KV_OFFLOAD_CONFIG_INFO + ] + + assert metadata.documentation.count("cpu:") == 1 + assert metadata.documentation.count("fs:") == 1 + + +def test_info_labelvalues_follow_the_declared_label_order(): + spec = _ConfigInfoOffloadingSpec( + ( + _CPUConfigInfo(chunks=8, policy="lru"), + _FsConfigInfo(path="/a"), + _FsConfigInfo(path="/b"), + ) + ) + + assert spec.info_labelvalues() == ("8", "lru", "/a", "/b") + + +def test_info_labelvalues_reject_a_source_that_reports_no_facts(): + """A source owns its own placeholder, so a missing config info is a bug.""" + spec = _ConfigInfoOffloadingSpec( + (None, _FsConfigInfo(path="/a"), _FsConfigInfo(path="/b")) # type: ignore[arg-type] + ) + + with pytest.raises(AssertionError, match="got NoneType"): + spec.info_labelvalues() + + +def test_info_labelvalues_reject_a_source_count_mismatch(): + spec = _ConfigInfoOffloadingSpec((_CPUConfigInfo(chunks=8, policy="lru"),)) + + with pytest.raises(AssertionError, match="declares 3 config info"): + spec.info_labelvalues() + + +_CPU_INFO_LABELS = ( + "cpu0_num_chunks", + "cpu0_blocks_per_chunk", + "cpu0_kv_bytes_per_chunk", + "cpu0_capacity_tokens_at_max_len", +) + + +def test_cpu_spec_declares_the_info_metric_with_the_cpu_cache_labels(): + """The CPU cache is the spec's only config source, under the name 'cpu'. + + The source name and its position set every label prefix, so a change here + renames the exported labels. + """ + metadata = CPUOffloadingSpec.build_info_metric_definition({})[ + KV_OFFLOAD_CONFIG_INFO + ] + + assert isinstance(metadata, OffloadingGaugeMetadata) + assert metadata.labelnames == _CPU_INFO_LABELS + assert CPUCacheOffloadingInfo.help_text() in metadata.documentation + + +def test_cpu_spec_info_labelvalues_follow_the_declared_label_order(): + """The values bind to _CPU_INFO_LABELS by position, so the order matters. + + The label names come from the class and the values from an instance, in two + separate walks. Nothing but this pairing catches the two drifting apart. + """ + spec = _create_spec() + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.info_labelvalues() == ( + str(spec.num_chunks), + str(spec.blocks_per_chunk), + str(spec.kv_bytes_per_chunk), + str(spec.tier_info.capacity_tokens_at_max_len), + ) + + +def test_cpu_spec_info_labelvalues_render_an_unknown_capacity_as_none(): + """A label takes a string, so an unknown token capacity becomes 'None'. + + The label must stay present, because Prometheus rejects a series that drops + one of the labels its metric declares. + """ + spec = _create_spec(max_model_len=0) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.tier_info.capacity_tokens_at_max_len is None + assert spec.info_labelvalues()[-1] == "None" + + +def test_tiering_spec_inherits_the_cpu_cache_info_declaration(): + """A tiering spec offloads to the same CPU cache, so it exports it too. + + Both hooks come by inheritance, which an override in the tiering spec would + silently break. + """ + metadata = TieringOffloadingSpec.build_info_metric_definition({})[ + KV_OFFLOAD_CONFIG_INFO + ] + spec = _create_spec(spec_name="TieringOffloadingSpec") + + assert isinstance(spec, TieringOffloadingSpec) + assert metadata.labelnames == _CPU_INFO_LABELS + assert spec.info_labelvalues() == spec.tier_info.as_labelvalues() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index c4e6d0f2a2ed..2472e4de7b96 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -4,6 +4,9 @@ from typing import TYPE_CHECKING +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( + get_sliding_window_size_in_chunks, +) from vllm.utils.math_utils import round_up from vllm.v1.core.kv_cache_utils import ( resolve_dcp_kv_block_size, @@ -74,22 +77,17 @@ def build_offloading_config( ) if not selected_groups: raise ValueError("KV offloading found no eligible cache groups.") - groups = tuple( - OffloadingGroupConfig( - group_id=group_id, - tokens_per_block=resolve_dcp_kv_block_size( - group.kv_cache_spec, - parallel_config.decode_context_parallel_size, - ), - layer_names=tuple(group.layer_names), + + group_tokens_per_block = tuple( + resolve_dcp_kv_block_size( + group.kv_cache_spec, parallel_config.decode_context_parallel_size ) - for group_id, group in selected_groups + for _, group in selected_groups ) - _, tokens_per_hash = resolve_kv_cache_block_sizes(kv_cache_config, vllm_config) - for group in groups: - assert group.tokens_per_block % tokens_per_hash == 0, ( - f"tokens_per_block={group.tokens_per_block} not divisible by " + for tokens_per_block in group_tokens_per_block: + assert tokens_per_block % tokens_per_hash == 0, ( + f"tokens_per_block={tokens_per_block} not divisible by " f"tokens_per_hash={tokens_per_hash}. " f"Hybrid models (e.g. Mamba+Attention) need " f"--enable-prefix-caching to align block sizes." @@ -114,7 +112,7 @@ def build_offloading_config( elif tokens_per_chunk is not None: tokens_per_chunk_int = int(tokens_per_chunk) - unique_tokens_per_block = {group.tokens_per_block for group in groups} + unique_tokens_per_block = set(group_tokens_per_block) assert len(unique_tokens_per_block) == 1, ( "If 'block_size' is specified in kv_connector_extra_config, " @@ -134,6 +132,21 @@ def build_offloading_config( f"'blocks_per_chunk' to express the chunk size in blocks." ) + groups = tuple( + OffloadingGroupConfig( + group_id=group_id, + tokens_per_block=tokens_per_block, + layer_names=tuple(group.layer_names), + sliding_window_size_in_chunks=get_sliding_window_size_in_chunks( + next(iter(iter_layer_specs(group.kv_cache_spec))), + tokens_per_block * blocks_per_chunk, + ), + ) + for (group_id, group), tokens_per_block in zip( + selected_groups, group_tokens_per_block + ) + ) + worker_kv_bytes_per_block = 0 all_groups_selected = len(selected_groups) == len(kv_cache_config.kv_cache_groups) if ( @@ -251,6 +264,7 @@ def spec_certifiable(spec: KVCacheSpec) -> bool: model=OffloadingModelConfig( name=vllm_config.model_config.model, dtype=str(cache_dtype).removeprefix("torch."), + max_model_len=vllm_config.model_config.max_model_len, ), cache=OffloadingCacheConfig( tokens_per_hash=tokens_per_hash, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py index 58023c393561..90212edf18b8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py @@ -327,6 +327,7 @@ def __init__( spec_cls = OffloadingSpecFactory.get_spec_cls(extra_config) self._offloading_metric_metadata: dict[str, OffloadingMetricMetadata] = { **spec_cls.build_metric_definitions(extra_config), + **spec_cls.build_info_metric_definition(extra_config), **get_connector_metric_definitions(), } from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec @@ -395,6 +396,7 @@ def _create_metric( metric_cls = self._counter_cls elif isinstance(metadata, OffloadingGaugeMetadata): metric_cls = self._gauge_cls + kwargs["multiprocess_mode"] = metadata.multiprocess_mode elif isinstance(metadata, OffloadingHistogramMetadata): metric_cls = self._histogram_cls if metadata.buckets is not None: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 2fb3c5bb905d..aa9dd742dc7e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -41,6 +41,7 @@ ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.kv_offload.base import ( + KV_OFFLOAD_CONFIG_INFO, GPULoadStoreSpec, Locality, LookupResult, @@ -542,6 +543,9 @@ def __init__( spec, vllm_config, kv_cache_config ) self.manager: OffloadingManager = spec.get_manager() + # Read after get_manager(), because a tiered spec builds its tiers + # inside that call. The values stay fixed for the process lifetime. + self._info_labelvalues = spec.info_labelvalues() self._connector_stats = OffloadingConnectorStats() full_attention_groups: list[int] = [] @@ -1896,6 +1900,12 @@ def get_stats(self) -> OffloadingConnectorStats | None: else: stats.aggregate(manager_stats) + # Static per-engine tier configuration. Written on every call, because + # a Prometheus child holds the last value it received. + if stats is None: + stats = OffloadingConnectorStats() + stats.set_gauge(KV_OFFLOAD_CONFIG_INFO, 1, self._info_labelvalues) + return stats def request_finished( diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 93f8e4bda4e9..e7066f51ebf7 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -5,8 +5,8 @@ """ from abc import ABC, abstractmethod -from collections.abc import Collection, Iterable, Sequence -from dataclasses import dataclass, field +from collections.abc import Collection, Iterable, Mapping, Sequence +from dataclasses import dataclass, field, fields from enum import Enum, auto from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, NewType, TypeVar @@ -202,7 +202,12 @@ class OffloadingCounterMetadata(OffloadingMetricMetadata): @dataclass(frozen=True) class OffloadingGaugeMetadata(OffloadingMetricMetadata): - pass + # Gauge-only in prometheus_client: how MultiProcessCollector merges samples + # written by different API-server processes. Offloading stats reach one + # frontend per step as complete per-engine snapshots, so the freshest write + # is the correct value and summing would multiply it by the number of + # participating frontends. + multiprocess_mode: str = "mostrecent" @dataclass(frozen=True) @@ -210,6 +215,50 @@ class OffloadingHistogramMetadata(OffloadingMetricMetadata): buckets: tuple[float, ...] | None = None +KV_OFFLOAD_CONFIG_INFO = "vllm:kv_offload_config_info" + +_INFO_METRIC_HELP = ( + "Static configuration of the KV offload sources of this engine instance. " + "The metric appears from the first scheduler step, so an idle engine " + "exposes no series. Each engine reports its own sources, not the instance " + "total." +) + + +@dataclass(frozen=True) +class OffloadingConfigInfo: + """Static, per-engine facts that one config source adds to an info metric. + + A config source is any part of a spec that owns configuration worth + publishing, for example the CPU cache or one secondary tier. + + The label names come from the class, in the API-server process. The label + values come from an instance, in the scheduler process. Both walks read + dataclasses.fields, which keeps the declaration order, so the two sides bind + by position. + """ + + @classmethod + def as_labelnames(cls, prefix: str) -> tuple[str, ...]: + """Return one label name for each field, in declaration order. + + Args: + prefix: Prefix of every label name. It names the config source that + owns these facts, so two sources cannot collide on one field + name. + """ + return tuple(f"{prefix}_{f.name}" for f in fields(cls)) + + @classmethod + def help_text(cls) -> str: + """Return one sentence about these fields, for the info metric HELP.""" + return "" + + def as_labelvalues(self) -> tuple[str, ...]: + """Return one label value for each field, in declaration order.""" + return tuple(str(getattr(self, f.name)) for f in fields(self)) + + @dataclass(frozen=True) class OffloadingKVEventsConfig: # Global vLLM KV event publishing flag. When false, connector-specific @@ -583,6 +632,107 @@ def build_metric_definitions( """Return Prometheus metric definitions emitted by this spec.""" return {} + @classmethod + @abstractmethod + def config_info_classes( + cls, extra_config: Mapping[str, Any] + ) -> tuple[tuple[str, type[OffloadingConfigInfo]], ...]: + """Return the name and the OffloadingConfigInfo class of every config + source of this spec. + + Runs in the API-server process, on the class, with no instance. The + position of an entry gives the source index, and the base builds each + label-name prefix from the source name and that index, for example + "fs1". + + config_info() must return one entry for each entry here, in the same + order. Return an empty tuple to emit the info metric with no labels. + + The two halves run in different processes. info_labelvalues() catches + a disagreement about the source count at the first scheduler step, not + at startup. + + Args: + extra_config: The kv_connector_extra_config of this engine. A spec + may read its own source list from it. + """ + raise NotImplementedError + + @abstractmethod + def config_info(self) -> tuple["OffloadingConfigInfo", ...]: + """Return the static facts of every config source, in declared order. + + Runs in the scheduler process, after get_manager(), because a tiering + spec builds its tiers inside that call. Entry i holds an instance of + config_info_classes()[i][1]. A source that cannot read one of its own + facts must put its own placeholder in that field. + """ + raise NotImplementedError + + @classmethod + def build_info_metric_definition( + cls, extra_config: Mapping[str, Any] + ) -> dict[str, "OffloadingMetricMetadata"]: + """Return the info metric definition built from config_info_classes(). + + The label names come from the classes alone, so this runs in the + API-server process. OffloadingConnectorScheduler renders the matching + values in the scheduler process. + + Args: + extra_config: The kv_connector_extra_config of this engine, passed + to config_info_classes(). + """ + declared = cls.config_info_classes(extra_config) + labelnames: tuple[str, ...] = () + for idx, (name, info_cls) in enumerate(declared): + labelnames += info_cls.as_labelnames(f"{name}{idx}") + + # One title per unique source name: two sources of one name share a + # class, and one HELP string serves the whole metric family. + # dict.fromkeys drops the duplicate and keeps the declared order, which + # a set loses. + documentation = [_INFO_METRIC_HELP] + for name, info_cls in dict.fromkeys(declared): + help_text = info_cls.help_text() + if help_text: + documentation.append(f"{name}:\n{help_text}") + + return { + KV_OFFLOAD_CONFIG_INFO: OffloadingGaugeMetadata( + documentation="\n\n".join(documentation), + labelnames=labelnames, + ) + } + + def info_labelvalues(self) -> tuple[str, ...]: + """Return one label value for each label of the info metric. + + build_info_metric_definition declares the label names from the same + config sources, in the same order, so the values bind by position. + + Raises: + AssertionError: config_info() and config_info_classes() disagree + about the source count, or a source reports the wrong + OffloadingConfigInfo class. + """ + declared = type(self).config_info_classes(self.extra_config) + infos = self.config_info() + if len(infos) != len(declared): + raise AssertionError( + f"{type(self).__name__} declares {len(declared)} config info " + f"source(s) but reports {len(infos)}" + ) + + labelvalues: tuple[str, ...] = () + for (_, info_cls), info in zip(declared, infos): + assert isinstance(info, info_cls), ( + f"{type(self).__name__}.config_info() must report a " + f"{info_cls.__name__}, got {type(info).__name__}" + ) + labelvalues += info.as_labelvalues() + return labelvalues + def __init__(self, config: OffloadingConfig): self.config = config self.extra_config = config.extra_config diff --git a/vllm/v1/kv_offload/config.py b/vllm/v1/kv_offload/config.py index 74b619dc0f8a..6e531dc7da50 100644 --- a/vllm/v1/kv_offload/config.py +++ b/vllm/v1/kv_offload/config.py @@ -16,6 +16,9 @@ class OffloadingGroupConfig: layer_names: tuple[str, ...] # Original KVCacheConfig group index. group_id: int + # Chunks of this group the tier keeps for one request, or None when + # attention reaches back without a bound (a full layer). + sliding_window_size_in_chunks: int | None = None @dataclass(frozen=True) @@ -24,6 +27,8 @@ class OffloadingModelConfig: name: str # KV cache data type (e.g. "float16"). dtype: str + # Longest request the engine accepts, in tokens. 0 when not known. + max_model_len: int = 0 @dataclass(frozen=True) diff --git a/vllm/v1/kv_offload/cpu/common.py b/vllm/v1/kv_offload/cpu/common.py index 61dce9365ea8..c3edc807b89d 100644 --- a/vllm/v1/kv_offload/cpu/common.py +++ b/vllm/v1/kv_offload/cpu/common.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass + import numpy as np -from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec +from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec, OffloadingConfigInfo class CPUOffloadingMetrics: @@ -13,6 +15,35 @@ class CPUOffloadingMetrics: CPU_CACHE_READ_USAGE_PERC = "vllm:kv_offload_cpu_cache_read_usage_perc" +@dataclass(frozen=True) +class CPUCacheOffloadingInfo(OffloadingConfigInfo): + """Static, per-engine facts about the CPU offload tier. + + One config source of the KV offload info metric. The base renders the + label names from these fields and the label values from an instance. + """ + + # Chunk count, not GPU blocks; see blocks_per_chunk. + num_chunks: int + # GPU blocks per chunk; the CPU-slot to GPU-block conversion factor. + blocks_per_chunk: int + # Page-aligned bytes per chunk. With num_chunks this is the tier's exact + # size in bytes, the only capacity valid for every model shape. + kv_bytes_per_chunk: int + # Upper bound on the KV tokens the tier holds, over the request lengths up + # to max_model_len. None when max_model_len is not known. + # See _capacity_tokens_at_max_len. + capacity_tokens_at_max_len: int | None + + @classmethod + def help_text(cls) -> str: + return ( + "The size of the CPU cache, in chunks and in bytes, and an upper bound " + "on the KV tokens it holds. The bound covers the request lengths up to " + "max_model_len. It is 'None' when max_model_len is not known." + ) + + class CPULoadStoreSpec(BlockIDsLoadStoreSpec): """Spec for loading/storing KV chunks to/from CPU memory. diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 67b469b6a394..0eba80e3187b 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -19,10 +19,7 @@ ReqContext, RequestOffloadingContext, ) -from vllm.v1.kv_offload.cpu.common import ( - CPULoadStoreSpec, - CPUOffloadingMetrics, -) +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec, CPUOffloadingMetrics from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 00b36e863a07..3a579f5ba06d 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -1,14 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Mapping from typing import Any import torch from typing_extensions import override from vllm.platforms import current_platform -from vllm.utils.math_utils import round_up +from vllm.utils.math_utils import cdiv, round_down, round_up from vllm.v1.kv_offload.base import ( CanonicalKVCaches, + OffloadingConfigInfo, OffloadingCounterMetadata, OffloadingGaugeMetadata, OffloadingHistogramMetadata, @@ -17,8 +19,8 @@ OffloadingSpec, OffloadingWorker, ) -from vllm.v1.kv_offload.config import OffloadingConfig -from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics +from vllm.v1.kv_offload.config import OffloadingConfig, OffloadingGroupConfig +from vllm.v1.kv_offload.cpu.common import CPUCacheOffloadingInfo, CPUOffloadingMetrics from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion @@ -41,6 +43,61 @@ def _all_workers_barrier() -> None: group.barrier() +def _chunks_per_request( + groups: tuple[OffloadingGroupConfig, ...], + blocks_per_chunk: int, + seq_len: int, +) -> int: + """Chunks one request of seq_len tokens holds, summed over every group. + + A group without a window holds one chunk for each chunk-sized span of the + request. A group with a window holds no more than the window, because the + tier lets the older chunks of that group age out. + """ + total = 0 + for group in groups: + chunk_count = cdiv(seq_len, blocks_per_chunk * group.tokens_per_block) + window_chunks = group.sliding_window_size_in_chunks + if window_chunks is not None: + chunk_count = min(chunk_count, window_chunks) + total += chunk_count + return total + + +def _capacity_tokens_at_max_len( + groups: tuple[OffloadingGroupConfig, ...], + blocks_per_chunk: int, + num_chunks: int, + max_model_len: int, +) -> int | None: + """Largest token count the tier serves at a length up to max_model_len. + + A tier of num_chunks slots holds num_chunks / _chunks_per_request(seq_len) + requests, so it serves num_chunks * seq_len // _chunks_per_request(seq_len) + tokens. _chunks_per_request is a step function of seq_len, so that ratio + rises between two steps and drops at each step. Every peak therefore sits at + the last token of a chunk. Measure one such length for each group chunk size, + add max_model_len, and keep the largest result. + + Returns: + The token count, or None when max_model_len is 0 and the caller + therefore did not know the longest request. + """ + tokens_per_chunk = {blocks_per_chunk * group.tokens_per_block for group in groups} + if max_model_len <= 0 or not tokens_per_chunk or min(tokens_per_chunk) <= 0: + return None + + seq_len_candidates = { + round_down(max_model_len, tokens) for tokens in tokens_per_chunk + } + seq_len_candidates.add(max_model_len) + return max( + num_chunks * seq_len // _chunks_per_request(groups, blocks_per_chunk, seq_len) + for seq_len in seq_len_candidates + if seq_len > 0 + ) + + class CPUOffloadingSpec(OffloadingSpec): BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT @@ -91,6 +148,18 @@ def build_metric_definitions( ) return definitions + @classmethod + @override + def config_info_classes( + cls, extra_config: Mapping[str, Any] + ) -> tuple[tuple[str, type[OffloadingConfigInfo]], ...]: + """Declare the CPU cache as the only config source.""" + return (("cpu", CPUCacheOffloadingInfo),) + + @override + def config_info(self) -> tuple[OffloadingConfigInfo, ...]: + return (self.tier_info,) + def __init__(self, config: OffloadingConfig): super().__init__(config) @@ -126,6 +195,8 @@ def __init__(self, config: OffloadingConfig): # or |--- C0 (single copy) ---| *** maybe-pad *** | self.kv_bytes_per_chunk = aligned_kv_bytes_per_chunk + self.tier_info = self._build_tier_info(config) + # scheduler-side self._manager: OffloadingManager | None = None @@ -137,6 +208,20 @@ def __init__(self, config: OffloadingConfig): "cache_policy_module_path" ) + def _build_tier_info(self, config: OffloadingConfig) -> CPUCacheOffloadingInfo: + """Resolve the tier's static facts, including its token capacity.""" + return CPUCacheOffloadingInfo( + num_chunks=self.num_chunks, + blocks_per_chunk=self.blocks_per_chunk, + kv_bytes_per_chunk=self.kv_bytes_per_chunk, + capacity_tokens_at_max_len=_capacity_tokens_at_max_len( + config.groups, + self.blocks_per_chunk, + self.num_chunks, + config.model.max_model_len, + ), + ) + @override def get_manager(self) -> OffloadingManager: if not self._manager: