diff --git a/.buildkite/scripts/install-kv-offload.sh b/.buildkite/scripts/install-kv-offload.sh new file mode 100644 index 0000000000..46d756b621 --- /dev/null +++ b/.buildkite/scripts/install-kv-offload.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +set -euo pipefail + +if [ -z "${CUDA_VERSION:-}" ]; then + exit 0 +fi + +# TODO: Use the PyPI wheel once it is published. +uv pip install --system \ + "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.0" + +# Keep only the NIXL wheel matching the CI image's CUDA runtime. +NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p') +uv pip uninstall --system nixl-cu12 nixl-cu13 +uv pip install --system --no-deps "nixl-cu${CUDA_VERSION%%.*}==${NIXL_VERSION}" diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index b45b26b2f2..ed82c7d71d 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -130,6 +130,7 @@ steps: source_file_dependencies: *v1-core-kv-metrics-deps commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - bash /vllm-workspace/.buildkite/scripts/install-kv-offload.sh - pytest -v -s v1/kv_offload - pytest -v -s v1/simple_kv_offload mirror: diff --git a/tests/v1/kv_offload/tiering/test_factory.py b/tests/v1/kv_offload/tiering/test_factory.py index f5a576d254..c391ac2016 100644 --- a/tests/v1/kv_offload/tiering/test_factory.py +++ b/tests/v1/kv_offload/tiering/test_factory.py @@ -10,6 +10,7 @@ These tests verify: 3. Error paths — missing tier_type, unknown tier_type, duplicate registration. """ +from importlib.util import find_spec from unittest.mock import MagicMock import pytest @@ -42,8 +43,11 @@ def _make_mock_args(): def test_pre_registered_tiers_can_be_imported(): - """CI sentinel: example/fs/obj paths must import and yield SecondaryTierManager.""" + """CI sentinel: registered tiers import and yield SecondaryTierManager.""" for tier_type in SecondaryTierFactory._registry: + # KVCR is an optional external dependency and may not be installed. + if tier_type == "kvcr" and find_spec("kvcr") is None: + continue cls = SecondaryTierFactory._registry[tier_type]() assert issubclass(cls, SecondaryTierManager) diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py new file mode 100644 index 0000000000..38a2ae044c --- /dev/null +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -0,0 +1,544 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Collection, Iterable, Mapping +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +pytest.importorskip("kvcr") + +from kvcr import ROUTER_HINT_KEY, KVCRBindings +from kvcr.config import G3Options, KVCRBackendConfigs, KVCRConfig, KVCRGuardConfig +from kvcr.policy import FIFOPolicy, G3FIFOPolicy, G3LRUPolicy, LRUPolicy +from kvcr.types import ( + BlockKey, + CacheTier, + InventoryEvent, + MemDescriptor, + OpEntryResult, + OpEntryStatus, + OpHandle, + QueryStatus, +) + +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, +) +from vllm.v1.kv_offload.base import ( + LookupResult, + Medium, + OffloadKey, + ReqContext, + make_offload_key, +) +from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob +from vllm.v1.kv_offload.tiering.kvcr import manager as kvcr_manager +from vllm.v1.kv_offload.tiering.kvcr.manager import KVCRSecondaryTierManager + + +def _op_entries( + entries: Mapping[BlockKey, bool], +) -> dict[BlockKey, OpEntryResult]: + return { + key: OpEntryResult(OpEntryStatus.SUCCESS if success else OpEntryStatus.FAILED) + for key, success in entries.items() + } + + +class _StubControlChannel: + def __init__(self, endpoint: str) -> None: + self.endpoint = endpoint + + def send(self, endpoint: str, message: bytes) -> bool: + return True + + def recv(self) -> list[bytes]: + return [] + + def close(self) -> None: + pass + + +class RecordingKVCR: + def __init__(self) -> None: + self.config: KVCRConfig | None = None + self.guard_config: KVCRGuardConfig | None = None + self.backend_configs = KVCRBackendConfigs() + self.constructor_bindings: KVCRBindings | None = None + self.nixl_agent_name = "recording" + self.framework_control: _StubControlChannel | None = None + self.inventory_sink = None + self.query_status = QueryStatus.MISS + self.stats: OffloadingConnectorStats | None = None + self.submit_hint_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + self.discard_hint_calls: list[str] = [] + self.deliver_calls: list[ + tuple[OpHandle, dict[BlockKey, list[MemDescriptor]], str | None] + ] = [] + self.deposit_calls: list[ + tuple[OpHandle, dict[BlockKey, list[MemDescriptor]]] + ] = [] + self.completed: list[tuple[OpHandle, dict[BlockKey, OpEntryResult]]] = [] + self._next_op_handle = 1 + + def submit_hint(self, *args: Any, **kwargs: Any) -> None: + self.submit_hint_calls.append((args, kwargs)) + + def discard_hint(self, request_id: str) -> None: + self.discard_hint_calls.append(request_id) + + def query( + self, + keys: Collection[BlockKey], + request_id: str | None = None, + ) -> list[tuple[QueryStatus, CacheTier | None]]: + tier = { + QueryStatus.HIT: CacheTier.LOCAL_G2, + QueryStatus.FETCHING: CacheTier.LOCAL_G2, + QueryStatus.FETCHABLE: CacheTier.REMOTE_G2, + QueryStatus.MISS: None, + }[self.query_status] + return [(self.query_status, tier) for _ in keys] + + def deliver( + self, + blocks: Mapping[BlockKey, list[MemDescriptor]], + request_id: str | None = None, + ) -> OpHandle: + op_handle = self._next_op_handle + self._next_op_handle += 1 + self.deliver_calls.append((op_handle, dict(blocks), request_id)) + self.complete(op_handle, {key: True for key in blocks}) + return op_handle + + def deposit( + self, + blocks: Mapping[BlockKey, list[MemDescriptor]], + ) -> OpHandle: + op_handle = self._next_op_handle + self._next_op_handle += 1 + self.deposit_calls.append((op_handle, dict(blocks))) + self.complete(op_handle, {key: True for key in blocks}) + return op_handle + + def complete( + self, + op_handle: OpHandle, + entries: Mapping[BlockKey, bool], + ) -> None: + self.completed.append((op_handle, _op_entries(entries))) + + def poll_completed( + self, + ) -> Iterable[tuple[OpHandle, dict[BlockKey, OpEntryResult]]]: + completed = self.completed + self.completed = [] + return completed + + def get_stats(self) -> OffloadingConnectorStats | None: + stats = self.stats + self.stats = None + return stats + + def close(self) -> None: + pass + + +class _ExternalPolicy(FIFOPolicy): + pass + + +def _job( + job_id: int, + req_context: ReqContext, + key: OffloadKey | None = None, + chunk_id: int = 0, +) -> TransferJob: + return TransferJob( + job_id=job_id, + keys=[key if key is not None else OffloadKey(b"k0")], + chunk_ids=np.array([chunk_id], dtype=np.int64), + is_promotion=True, + req_context=req_context, + ) + + +def _make_tier( + monkeypatch, + kvcr: RecordingKVCR, + *, + enable_telemetry: bool = False, + secondary_g2_slots: int = 0, + kvcr_service_socket_path: str | None = None, + compatibility_digest: str | None = None, + enable_kv_cache_events: bool = False, + self_describing_kv_events: bool = False, + policy: str | None = None, + g3: dict[str, object] | None = None, + control_ports: list[int] | None = None, + data_parallel_rank_local: int | None = None, +) -> KVCRSecondaryTierManager: + def make_control(_bind_host, bind_port, advertise_host): + return _StubControlChannel(f"tcp://{advertise_host}:{int(bind_port)}") + + def make_kvcr(config, bindings, backend_configs, guard_config): + kvcr.config = config + kvcr.guard_config = guard_config + kvcr.nixl_agent_name = config.nixl_agent_name + kvcr.backend_configs = backend_configs + kvcr.constructor_bindings = bindings + kvcr.framework_control = bindings.framework_control + kvcr.inventory_sink = bindings.inventory_sink + return kvcr + + monkeypatch.setattr(kvcr_manager, "KVCR", make_kvcr) + monkeypatch.setattr(kvcr_manager, "ZmqPeerControlChannel", make_control) + return KVCRSecondaryTierManager( + offloading_spec=SimpleNamespace( + config=SimpleNamespace( + parallel=SimpleNamespace( + data_parallel_rank_local=data_parallel_rank_local, + ) + ), + kv_events_config=SimpleNamespace( + enable_kv_cache_events=enable_kv_cache_events, + self_describing_kv_events=self_describing_kv_events, + ), + ), + primary_kv_view=memoryview(np.zeros((4, 16), dtype=np.int8)), + tier_type="kvcr", + router_capabilities=["router_hint"], + control_host="127.0.0.1", + control_ports=control_ports if control_ports is not None else [7777], + control_advertise_host="127.0.0.1", + enable_telemetry=enable_telemetry, + secondary_g2_slots=secondary_g2_slots, + kvcr_service_socket_path=kvcr_service_socket_path, + compatibility_digest=compatibility_digest, + policy=policy, + g3=g3, + local_dram_backend="UCX", + remote_fw_dram_backend="UCX", + ) + + +def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch): + """Keep the control endpoint and guard pool aligned to the local DP rank.""" + kvcr = RecordingKVCR() + tier = _make_tier( + monkeypatch, + kvcr, + kvcr_service_socket_path="/tmp/kvcr.sock", + compatibility_digest="Opaque-Digest", + secondary_g2_slots=1, + control_ports=[7001, 7002], + data_parallel_rank_local=1, + ) + + assert kvcr.framework_control is not None + assert kvcr.framework_control.endpoint == "tcp://127.0.0.1:7002" + assert kvcr.guard_config == KVCRGuardConfig( + kvcr_service_socket_path="/tmp/kvcr.sock", + guard_index=1, + compatibility_digest="Opaque-Digest", + ) + assert kvcr.config is not None + assert kvcr.config.pool_layouts == [("", tier._primary_row_stride)] + assert kvcr.backend_configs.local_dram is None + + +def test_kvcr_tier_converts_g3_paths(monkeypatch, tmp_path): + """Convert user-provided G3 paths to KVCR's typed configuration.""" + kvcr = RecordingKVCR() + path = tmp_path / "g3.data" + + _make_tier( + monkeypatch, + kvcr, + g3={"paths": [str(path)], "capacity_bytes_per_file": 64}, + ) + + assert kvcr.backend_configs.g3 == G3Options( + paths=(path,), + capacity_bytes_per_file=64, + ) + + +def test_kvcr_tier_adapts_request_and_load(monkeypatch): + """Check hint forwarding, key conversion, load descriptors, and cleanup.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + router_hint = {"opaque": True} + ctx = ReqContext( + req_id="req", + kv_transfer_params={ROUTER_HINT_KEY: router_hint, "unrelated": object()}, + ) + key = make_offload_key((123).to_bytes(8, "big"), 0) + same_hash_other_group = make_offload_key((123).to_bytes(8, "big"), 7) + other_key = make_offload_key((124).to_bytes(8, "big"), 0) + + tier.on_new_request(ctx) + + assert kvcr.submit_hint_calls == [((), {"request_id": "req", "hints": router_hint})] + + bindings = kvcr.constructor_bindings + assert bindings is not None + assert bindings.key_adapter is not None + monkeypatch.setenv("VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES", "0") + decode = bindings.key_adapter.decode + hash_123 = (123).to_bytes(8, "big") + assert decode(BlockKey(bytes(key))) == hash_123 + assert decode(BlockKey(bytes(same_hash_other_group))) == hash_123 + assert decode(BlockKey(bytes(other_key))) == (124).to_bytes(8, "big") + + tier.submit_load(_job(7, ctx, key=key, chunk_id=2)) + + # Check load destinations and job completion independently of hint parsing. + assert len(kvcr.submit_hint_calls) == 1 + _, blocks, request_id = kvcr.deliver_calls[0] + assert request_id == "req" + assert list(blocks) == [key] + (descriptor,) = blocks[key] + assert descriptor.end_point_name == kvcr.nixl_agent_name + assert descriptor.addr == tier._primary_base_addr + 2 * 16 + assert descriptor.size == 16 + assert list(tier.get_finished_jobs()) == [JobResult(7, True)] + + # Here we verify that request cleanup discards the request-scoped hint in + # KVCR. + tier.on_request_finished(ctx) + assert kvcr.discard_hint_calls == ["req"] + + +def test_kvcr_tier_allows_request_without_router_hint(monkeypatch): + """Keep router hints optional for requests from non-hint-aware routers.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + + tier.on_new_request(ReqContext(req_id="req")) + + assert kvcr.submit_hint_calls == [] + + +@pytest.mark.parametrize( + ("status", "expected"), + [ + (QueryStatus.MISS, LookupResult.MISS), + (QueryStatus.HIT, LookupResult.HIT), + (QueryStatus.FETCHABLE, LookupResult.HIT), + (QueryStatus.FETCHING, LookupResult.RETRY), + ], +) +def test_kvcr_tier_maps_query_status(monkeypatch, status, expected): + """Map KVCR cache states to the scheduler's lookup contract.""" + kvcr = RecordingKVCR() + kvcr.query_status = status + tier = _make_tier(monkeypatch, kvcr) + + assert tier.lookup(OffloadKey(b"k0"), ReqContext(req_id="req")) is expected + + +def test_kvcr_tier_serves_primary_pin_request(monkeypatch): + """Hold primary-tier hits until KVCR releases the corresponding pin.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + keys = (BlockKey(b"k0"), BlockKey(b"k1"), BlockKey(b"k2")) + hit_keys = (keys[0], keys[2]) + chunk_ids = {keys[0]: 1, keys[2]: 5} + lifecycle: list[str] = [] + + class Parent: + def on_new_request(self, req_context): + lifecycle.append("new") + return SimpleNamespace() + + def lookup(self, key, req_context): + return LookupResult.HIT if key in hit_keys else LookupResult.MISS + + def create_store_job(self, requested_keys, req_context): + return TransferJob( + job_id=11, + keys=requested_keys, + chunk_ids=np.array([chunk_ids[key] for key in requested_keys]), + is_promotion=True, + req_context=req_context, + ) + + def on_request_finished(self, req_context): + lifecycle.append("finished") + + bindings = kvcr.constructor_bindings + assert bindings is not None + request = bindings.request_pin(keys) + + tier.serve_external_requests(Parent()) + + [(queued_request, result)] = bindings.poll_pin_results() + assert queued_request == request + assert result is not None + pin_handle, descriptors = result + assert descriptors[keys[1]] is None + block_descriptors = descriptors[keys[2]] + assert block_descriptors is not None + (descriptor,) = block_descriptors + assert descriptor.addr == (tier._primary_base_addr + 5 * tier._primary_row_stride) + assert lifecycle == ["new", "finished"] + + polls = 0 + + def release_pin(): + nonlocal polls + polls += 1 + bindings.release_pin(pin_handle) + return [] + + monkeypatch.setattr(kvcr, "poll_completed", release_pin) + tier.drain_jobs() + + assert polls == 1 + assert list(tier.get_finished_jobs()) == [JobResult(11, True)] + + +def test_kvcr_telemetry_is_opt_in_and_namespaced_at_vllm_boundary(monkeypatch): + """Keep telemetry opt-in and namespace metrics only at the vLLM boundary.""" + assert KVCRSecondaryTierManager.build_metric_definitions({}) == {} + definitions = KVCRSecondaryTierManager.build_metric_definitions( + {"enable_telemetry": True} + ) + assert "vllm:kvcr_duration_seconds" in definitions + assert "vllm:kvcr_transfer_blocks" in definitions + + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr, enable_telemetry=True) + bindings = kvcr.constructor_bindings + assert bindings is not None + assert bindings.stats_factory is not None + stats = bindings.stats_factory() + stats.increase_counter( + "kvcr_transfer_blocks", + 2, + ("remote_deliver",), + ) + kvcr.stats = stats + + returned = tier.get_stats() + + assert returned is stats + assert returned.reduce() == {"vllm:kvcr_transfer_blocks:('remote_deliver',)": 2} + + +@pytest.mark.parametrize( + ("policy", "expected_type"), + [ + ("fifo", FIFOPolicy), + ("lru", LRUPolicy), + ("g3_fifo", G3FIFOPolicy), + ("g3_lru", G3LRUPolicy), + (f"{__name__}._ExternalPolicy", _ExternalPolicy), + ], +) +def test_kvcr_tier_passes_selected_policy(monkeypatch, policy, expected_type): + """Resolve every built-in and fully qualified external policy.""" + kvcr = RecordingKVCR() + _make_tier(monkeypatch, kvcr, policy=policy) + + bindings = kvcr.constructor_bindings + assert bindings is not None + assert type(bindings.policy) is expected_type + + +@pytest.mark.parametrize( + ("socket_path", "digest"), + [("/tmp/kvcr.sock", None), (None, "Opaque-Digest")], +) +def test_kvcr_tier_requires_complete_service_config(monkeypatch, socket_path, digest): + """Reject partial guard configuration before connecting to the service.""" + with pytest.raises(ValueError, match="configured together"): + _make_tier( + monkeypatch, + RecordingKVCR(), + kvcr_service_socket_path=socket_path, + compatibility_digest=digest, + ) + + +def test_kvcr_tier_stores_and_emits_inventory(monkeypatch): + """Cover store descriptors and inventory translation at the KVCR boundary.""" + kvcr = RecordingKVCR() + tier = _make_tier( + monkeypatch, + kvcr, + secondary_g2_slots=2, + enable_kv_cache_events=True, + self_describing_kv_events=True, + ) + local_dram = kvcr.backend_configs.local_dram + assert local_dram is not None + assert [(name, size) for name, _, size in local_dram.pools] == [ + ("", 2 * tier._primary_row_stride) + ] + + key = OffloadKey(b"k0") + tier.submit_store(_job(11, ReqContext(req_id="req"), key=key, chunk_id=2)) + + _, blocks = kvcr.deposit_calls[0] + assert blocks[key][0].addr == tier._primary_base_addr + 2 * 16 + assert list(tier.get_finished_jobs()) == [JobResult(11, True)] + + assert kvcr.inventory_sink is not None + kvcr.inventory_sink( + InventoryEvent((BlockKey(bytes(key)),), CacheTier.LOCAL_G2, False) + ) + kvcr.inventory_sink(InventoryEvent((BlockKey(bytes(key)),), CacheTier.G3, False)) + events = list(tier.take_events()) + assert [ + (event.keys, event.medium, event.ownership, event.removed) for event in events + ] == [ + ([key], Medium.CPU, "kvcr", False), + ([key], Medium.STORAGE, "kvcr", False), + ] + assert all(event.removal_expected for event in events) + tier.shutdown() + + +def test_kvcr_tier_requires_self_describing_inventory_events(monkeypatch): + """Require tier-aware events whenever KVCR owns local cache inventory.""" + with pytest.raises(ValueError, match="self_describing_kv_events"): + _make_tier( + monkeypatch, + RecordingKVCR(), + secondary_g2_slots=1, + enable_kv_cache_events=True, + ) + + +def test_kvcr_tier_accumulates_block_results(monkeypatch): + """Wait for every block result and preserve partial success while draining.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + keys = [OffloadKey(b"k0"), OffloadKey(b"k1")] + tier.submit_load( + TransferJob( + job_id=13, + keys=keys, + chunk_ids=np.array([0, 1], dtype=np.int64), + is_promotion=True, + req_context=ReqContext(req_id="req"), + ) + ) + + op_handle, _, _ = kvcr.deliver_calls[0] + kvcr.completed.clear() + kvcr.complete(op_handle, {keys[0]: True}) + assert list(tier.get_finished_jobs()) == [] + + kvcr.complete(op_handle, {keys[1]: False}) + tier.drain_jobs() + + [result] = tier.get_finished_jobs() + assert result.job_id == 13 + assert not result.success + assert result.successful_keys == {keys[0]} diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index 17f8430204..3310167b87 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -123,3 +123,9 @@ SecondaryTierFactory.register_tier( "vllm.v1.kv_offload.tiering.obj.manager", "ObjectStoreSecondaryTierManager", ) + +SecondaryTierFactory.register_tier( + "kvcr", + "vllm.v1.kv_offload.tiering.kvcr.manager", + "KVCRSecondaryTierManager", +) diff --git a/vllm/v1/kv_offload/tiering/kvcr/__init__.py b/vllm/v1/kv_offload/tiering/kvcr/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/v1/kv_offload/tiering/kvcr/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py new file mode 100644 index 0000000000..00186279a4 --- /dev/null +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -0,0 +1,635 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM secondary-tier adapter for KVCR.""" + +import ctypes +import mmap +import socket +import time +import uuid +from collections.abc import Collection, Iterable, Mapping +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +from kvcr import ( + DURATION_METRIC, + KVCR, + ROUTER_HINT_CAPABILITIES, + ROUTER_HINT_KEY, + STATE_METRIC, + TRANSFER_BLOCKS_METRIC, + TRANSFER_BYTES_METRIC, + KVCRBindings, +) +from kvcr.config import ( + FrameworkDramInput, + G3Options, + KeyAdapter, + KVCRBackendConfigs, + KVCRConfig, + KVCRGuardConfig, + LocalDramOptions, + RemoteFWDramOptions, +) +from kvcr.control_channels import ZmqPeerControlChannel +from kvcr.policy import ( + FIFOPolicy, + G3FIFOPolicy, + G3LRUPolicy, + KVCachePolicy, + LRUPolicy, +) +from kvcr.types import ( + BlockKey, + CacheTier, + InventoryEvent, + MemDescriptor, + OpHandle, + PinRequestId, + PinResult, + QueryStatus, +) +from typing_extensions import override + +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, +) +from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname +from vllm.v1.core.kv_cache_utils import ( + BlockHash, + maybe_convert_block_hash, +) +from vllm.v1.kv_offload.base import ( + LookupResult, + Medium, + OffloadingCounterMetadata, + OffloadingEvent, + OffloadingGaugeMetadata, + OffloadingHistogramMetadata, + OffloadingMetricMetadata, + OffloadKey, + ReqContext, + RequestOffloadingContext, + get_offload_block_hash, +) +from vllm.v1.kv_offload.tiering.base import ( + JobResult, + ParentManager, + SecondaryTierManager, + TransferJob, +) + +if TYPE_CHECKING: + from vllm.v1.kv_offload.base import OffloadingSpec + + +_REQUIRED_ROUTER_CAPABILITIES = ROUTER_HINT_CAPABILITIES + +logger = init_logger(__name__) + +_BUILTIN_POLICIES: dict[str, type[KVCachePolicy]] = { + "fifo": FIFOPolicy, + "lru": LRUPolicy, + "g3_fifo": G3FIFOPolicy, + "g3_lru": G3LRUPolicy, +} + + +def _resolve_policy(name: str | None) -> KVCachePolicy | None: + if name is None: + return None + policy_type = _BUILTIN_POLICIES.get(name) + if policy_type is None: + if "." not in name: + raise ValueError( + f"Unknown KVCR policy: {name!r}. " + f"Supported: {list(_BUILTIN_POLICIES)}; external policies " + "must use a fully qualified module.Class path." + ) + policy_type = resolve_obj_by_qualname(name) + if not isinstance(policy_type, type) or not issubclass( + policy_type, KVCachePolicy + ): + raise TypeError(f"{name} is not a KVCachePolicy") + return policy_type() + + +def _vllm_metric_name(name: str) -> str: + return f"vllm:{name}" + + +def _kvcr_metric_definitions() -> dict[str, OffloadingMetricMetadata]: + return { + _vllm_metric_name(DURATION_METRIC): OffloadingHistogramMetadata( + documentation="KVCR operation and stage duration in seconds.", + labelnames=("scope", "result"), + ), + _vllm_metric_name(TRANSFER_BYTES_METRIC): OffloadingCounterMetadata( + documentation="Bytes transferred by KVCR operations.", + labelnames=("operation",), + ), + _vllm_metric_name(TRANSFER_BLOCKS_METRIC): OffloadingCounterMetadata( + documentation="KV blocks transferred by successful KVCR operations.", + labelnames=("operation",), + ), + _vllm_metric_name(STATE_METRIC): OffloadingGaugeMetadata( + documentation="Current KVCR metadata and operation counts.", + labelnames=("resource",), + ), + } + + +class _FrameworkPinAdapter: + """Adapt KVCR framework pin callbacks to primary-tier store jobs.""" + + def __init__(self, tier: "KVCRSecondaryTierManager") -> None: + self._tier = tier + self._next_request_id = 0 + self._next_pin_request_id = 0 + self._pending_pins: dict[PinRequestId, tuple[BlockKey, ...]] = {} + self._completed_pins: list[tuple[PinRequestId, PinResult]] = [] + self._pin_jobs: dict[str, int] = {} + self._pin_job_results: list[JobResult] = [] + + def request_pin(self, keys: Collection[BlockKey]) -> PinRequestId: + # The parent callback is valid only from serve_external_requests(). + request = PinRequestId(self._next_pin_request_id) + self._next_pin_request_id += 1 + self._pending_pins[request] = tuple(keys) + return request + + def process_pending(self, parent: ParentManager) -> None: + pending = self._pending_pins + self._pending_pins = {} + self._completed_pins.extend( + (request, self._resolve_pin(parent, keys)) + for request, keys in pending.items() + ) + + def poll_pin_results(self) -> list[tuple[PinRequestId, PinResult]]: + completed = self._completed_pins + self._completed_pins = [] + return completed + + def cancel_pin_request(self, request: PinRequestId) -> None: + self._pending_pins.pop(request, None) + + def _resolve_pin( + self, parent: ParentManager, keys: Collection[BlockKey] + ) -> tuple[str, dict[BlockKey, list[MemDescriptor] | None]] | None: + offload_keys = tuple(OffloadKey(bytes(key)) for key in keys) + request_id = f"kvcr-source:{self._next_request_id}" + self._next_request_id += 1 + req_context = ReqContext(req_id=request_id) + started = False + job: TransferJob | None = None + pin_handle: str | None = None + try: + parent.on_new_request(req_context) + started = True + hit_keys = tuple( + key + for key in offload_keys + if parent.lookup(key, req_context) is LookupResult.HIT + ) + if not hit_keys: + return None + + job = parent.create_store_job(hit_keys, req_context) + job_keys = tuple(job.keys) + chunk_ids = tuple(int(chunk_id) for chunk_id in job.chunk_ids) + if len(job_keys) != len(chunk_ids) or set(job_keys) != set(hit_keys): + return None + + descriptors: dict[BlockKey, list[MemDescriptor] | None] = { + BlockKey(bytes(key)): None for key in offload_keys + } + descriptors.update( + ( + BlockKey(bytes(key)), + [self._tier._make_descriptor(chunk_id)], + ) + for key, chunk_id in zip(job_keys, chunk_ids) + ) + pin_handle = f"kvcr-job-{job.job_id}" + self._pin_jobs[pin_handle] = job.job_id + return pin_handle, descriptors + except Exception: + logger.warning("KVCR source acquisition failed", exc_info=True) + return None + finally: + if started: + try: + parent.on_request_finished(req_context) + except Exception: + logger.warning("KVCR source request cleanup failed", exc_info=True) + if job is not None and pin_handle is None: + self._pin_job_results.append( + JobResult(job_id=job.job_id, success=False) + ) + + def release_pin(self, pin_handle: str) -> bool: + job_id = self._pin_jobs.pop(pin_handle, None) + if job_id is not None: + self._pin_job_results.append(JobResult(job_id=job_id, success=True)) + return True + + def has_active_pins(self) -> bool: + return bool(self._pin_jobs) + + def take_pin_job_results(self) -> list[JobResult]: + results = self._pin_job_results + self._pin_job_results = [] + return results + + +class _VllmKeyAdapter: + """Adapt vLLM's local key format to KVCR.""" + + def encode(self, framework_key: object) -> BlockKey: + if not isinstance(framework_key, bytes): + raise TypeError("vLLM offload keys must be bytes") + return BlockKey(framework_key) + + def decode(self, key: BlockKey) -> int | bytes: + block_hash = BlockHash(get_offload_block_hash(OffloadKey(key))) + return maybe_convert_block_hash(block_hash) + + +# Job ID, remaining blocks, aggregate success, and successful load keys. +_JobState = tuple[int, int, bool, set[OffloadKey] | None] + + +class KVCRSecondaryTierManager(SecondaryTierManager): + """Secondary tier wrapper around the KVCR KV P2P API.""" + + @classmethod + @override + def build_metric_definitions( + cls, extra_config: dict[str, Any] + ) -> dict[str, OffloadingMetricMetadata]: + if not extra_config.get("enable_telemetry", False): + return {} + return _kvcr_metric_definitions() + + def __init__( + self, + offloading_spec: "OffloadingSpec", + primary_kv_view: memoryview, + tier_type: str, + router_capabilities: Iterable[str] | None = None, + control_host: str = "0.0.0.0", + control_ports: list[int] | None = None, + control_advertise_host: str | None = None, + eager_ctrl_connect: bool = True, + opportunistic_query: bool = True, + enable_telemetry: bool = False, + operation_timeout_ms: int = 1000, + abandon_timeout_ms: int = 5000, + metadata_retry_interval_ms: int = 100, + secondary_g2_slots: int = 0, + kvcr_service_socket_path: str | None = None, + compatibility_digest: str | None = None, + policy: str | None = None, + g3: dict[str, Any] | None = None, + local_dram_backend: str = "UCX", + remote_fw_dram_backend: str = "UCX", + ) -> None: + super().__init__(offloading_spec, primary_kv_view, tier_type) + selected_policy = _resolve_policy(policy) + if (kvcr_service_socket_path is None) != (compatibility_digest is None): + raise ValueError( + "kvcr_service_socket_path and compatibility_digest must be " + "configured together" + ) + missing = _REQUIRED_ROUTER_CAPABILITIES - set(router_capabilities or ()) + if missing: + raise ValueError( + f"KVCR secondary tier requires router capabilities: {sorted(missing)}" + ) + events_enabled = offloading_spec.kv_events_config.enable_kv_cache_events + self_describing_events = ( + offloading_spec.kv_events_config.self_describing_kv_events + ) + if ( + events_enabled + and ( + secondary_g2_slots > 0 + or kvcr_service_socket_path is not None + or g3 is not None + ) + and not self_describing_events + ): + raise ValueError( + "KVCR local G2/G3 inventory requires self_describing_kv_events" + ) + self._key_adapter: KeyAdapter = _VllmKeyAdapter() + advertise_host = control_advertise_host or socket.gethostname() + dp_local_rank = offloading_spec.config.parallel.data_parallel_rank_local + if dp_local_rank is None: + if control_ports is None or len(control_ports) != 1: + raise ValueError( + "control_ports must contain exactly one port when " + "data_parallel_rank_local is unset" + ) + dp_local_rank = 0 + if control_ports is None or not 0 <= dp_local_rank < len(control_ports): + raise ValueError( + "control_ports must contain a port for " + f"local data-parallel rank {dp_local_rank}" + ) + if any(not 1 <= port <= 65535 for port in control_ports): + raise ValueError("control_ports must be between 1 and 65535") + control = ZmqPeerControlChannel( + control_host, + control_ports[dp_local_rank], + advertise_host, + ) + # Give colocated workers distinct transport listen ports. + with socket.socket() as _s: + _s.bind(("", 0)) + _nixl_listen_port = _s.getsockname()[1] + self._primary_base_addr = ctypes.addressof( + ctypes.c_char.from_buffer(primary_kv_view) + ) + if primary_kv_view.strides is None: + raise ValueError("primary KV memoryview must expose strides") + self._primary_row_stride = int(primary_kv_view.strides[0]) + if secondary_g2_slots < 0: + raise ValueError("secondary_g2_slots must be non-negative") + if ( + secondary_g2_slots or kvcr_service_socket_path is not None + ) and not events_enabled: + logger.warning( + "KVCR local DRAM is enabled, but KV cache events are disabled; " + "local DRAM inventory will not be published" + ) + local_mapping: mmap.mmap | None = None + local_dram: LocalDramOptions | None = None + if secondary_g2_slots: + if kvcr_service_socket_path is not None: + logger.warning( + "secondary_g2_slots is ignored when " + "kvcr_service_socket_path is configured" + ) + else: + local_mapping = mmap.mmap( + -1, secondary_g2_slots * self._primary_row_stride + ) + local_dram = LocalDramOptions( + pools=[ + ( + "", + ctypes.addressof(ctypes.c_char.from_buffer(local_mapping)), + len(local_mapping), + ) + ], + backend=local_dram_backend, + ) + guard_config = ( + KVCRGuardConfig( + kvcr_service_socket_path=kvcr_service_socket_path, + guard_index=dp_local_rank, + compatibility_digest=compatibility_digest, + ) + if kvcr_service_socket_path is not None and compatibility_digest is not None + else None + ) + g3_config = None + if g3 is not None: + g3_values = dict(g3) + g3_values["paths"] = tuple(Path(path) for path in g3_values["paths"]) + g3_config = G3Options(**g3_values) + nixl_agent_name = f"KVCR-{uuid.uuid4()}" + self._framework_pin_adapter = _FrameworkPinAdapter(self) + self._inventory_events: list[OffloadingEvent] = [] + + try: + self._kvcr = KVCR( + KVCRConfig( + pool_layouts=[("", self._primary_row_stride)], + enable_telemetry=enable_telemetry, + operation_timeout_ms=operation_timeout_ms, + abandon_timeout_ms=abandon_timeout_ms, + nixl_agent_name=nixl_agent_name, + nixl_listen_port=_nixl_listen_port, + ), + KVCRBindings( + request_pin=self._framework_pin_adapter.request_pin, + poll_pin_results=self._framework_pin_adapter.poll_pin_results, + release_pin=self._framework_pin_adapter.release_pin, + cancel_pin_request=(self._framework_pin_adapter.cancel_pin_request), + framework_control=control, + key_adapter=self._key_adapter, + inventory_sink=(self._record_inventory if events_enabled else None), + stats_factory=( + OffloadingConnectorStats if enable_telemetry else None + ), + policy=selected_policy, + ), + KVCRBackendConfigs( + framework_dram=FrameworkDramInput( + self._primary_base_addr, primary_kv_view.nbytes + ), + local_dram=local_dram, + g3=g3_config, + remote_fw_dram=RemoteFWDramOptions( + eager_ctrl_connect=eager_ctrl_connect, + opportunistic_query=opportunistic_query, + metadata_retry_interval_ms=metadata_retry_interval_ms, + backend=remote_fw_dram_backend, + ), + ), + guard_config, + ) + except BaseException: + control.close() + if local_mapping is not None: + local_mapping.close() + raise + self._local_dram_mmap = local_mapping + self._finished_jobs: list[JobResult] = [] + self._jobs_by_op: dict[OpHandle, _JobState] = {} + + @override + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + block_key = self._key_adapter.encode(key) + status, _ = self._kvcr.query((block_key,), req_context.req_id)[0] + if status is QueryStatus.FETCHING: + return LookupResult.RETRY + if status in (QueryStatus.HIT, QueryStatus.FETCHABLE): + return LookupResult.HIT + return LookupResult.MISS + + @override + def submit_load(self, job_metadata: TransferJob) -> None: + blocks = { + self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + for key, chunk_id in zip( + job_metadata.keys, job_metadata.chunk_ids, strict=True + ) + } + if not blocks: + self._finished_jobs.append( + JobResult(job_id=job_metadata.job_id, success=True) + ) + return + + op_handle = self._kvcr.deliver( + blocks, request_id=job_metadata.req_context.req_id + ) + self._jobs_by_op[op_handle] = ( + job_metadata.job_id, + len(blocks), + True, + set(), + ) + + @override + def submit_store(self, job_metadata: TransferJob) -> None: + blocks = { + self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + for key, chunk_id in zip( + job_metadata.keys, job_metadata.chunk_ids, strict=True + ) + } + if not blocks: + self._finished_jobs.append( + JobResult(job_id=job_metadata.job_id, success=True) + ) + return + op_handle = self._kvcr.deposit(blocks) + self._jobs_by_op[op_handle] = ( + job_metadata.job_id, + len(blocks), + True, + None, + ) + + @override + def get_finished_jobs(self) -> Iterable[JobResult]: + results = self._finished_jobs + self._finished_jobs = [] + results.extend(self._poll_finished_jobs()) + return results + + @override + def get_stats(self) -> OffloadingConnectorStats | None: + stats = cast(OffloadingConnectorStats | None, self._kvcr.get_stats()) + if stats is not None: + # Add the framework namespace only at the vLLM boundary. + for metrics in stats.data.values(): + for name in tuple(metrics): + metrics[_vllm_metric_name(name)] = metrics.pop(name) + return stats + + @override + def serve_external_requests(self, parent: ParentManager) -> None: + self._framework_pin_adapter.process_pending(parent) + + def _poll_finished_jobs(self) -> list[JobResult]: + results: list[JobResult] = [] + for op_handle, entries in self._kvcr.poll_completed(): + job_state = self._jobs_by_op.get(op_handle) + if job_state is None: + continue + job_id, remaining, success, successful_keys = job_state + remaining -= len(entries) + success = success and all(entry.success for entry in entries.values()) + if successful_keys is not None: + successful_keys.update( + OffloadKey(bytes(key)) + for key, entry in entries.items() + if entry.success + ) + if remaining > 0: + self._jobs_by_op[op_handle] = ( + job_id, + remaining, + success, + successful_keys, + ) + continue + self._jobs_by_op.pop(op_handle, None) + results.append( + JobResult( + job_id=job_id, + success=success, + successful_keys=successful_keys if not success else None, + ) + ) + results.extend(self._framework_pin_adapter.take_pin_job_results()) + return results + + @override + def has_pending_work(self) -> bool: + # Keep scheduler polling: control messages may arrive without a load job. + return True + + @override + def take_events(self) -> Iterable[OffloadingEvent]: + events = self._inventory_events + self._inventory_events = [] + return events + + @override + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + params = getattr(req_context, "kv_transfer_params", None) + if isinstance(params, Mapping): + hint = params.get(ROUTER_HINT_KEY) + if hint is not None: + self._kvcr.submit_hint( + request_id=req_context.req_id, + hints=hint, + ) + return RequestOffloadingContext() + + @override + def on_request_finished(self, req_context: ReqContext) -> None: + self._kvcr.discard_hint(req_context.req_id) + + @override + def drain_jobs(self) -> None: + while self._jobs_by_op or self._framework_pin_adapter.has_active_pins(): + drained = self._poll_finished_jobs() + if drained: + self._finished_jobs.extend(drained) + else: + time.sleep(0.001) + + @override + def shutdown(self) -> None: + self.drain_jobs() + # TODO: Keep registered buffers alive if KVCR.close() fails. + self._kvcr.close() + if self._local_dram_mmap is not None: + self._local_dram_mmap.close() + self._local_dram_mmap = None + + def _record_inventory(self, event: InventoryEvent) -> None: + if event.tier is CacheTier.LOCAL_G2: + medium = Medium.CPU + elif event.tier is CacheTier.G3: + medium = Medium.STORAGE + else: + raise ValueError(f"unsupported KVCR inventory tier: {event.tier}") + self._inventory_events.append( + OffloadingEvent( + keys=[OffloadKey(bytes(key)) for key in event.keys], + medium=medium, + removed=event.removed, + ownership="kvcr", + removal_expected=True, + ) + ) + + def _make_descriptor(self, chunk_id: int) -> MemDescriptor: + return MemDescriptor( + end_point_name=self._kvcr.config.nixl_agent_name, + mem_type="DRAM", + addr=self._primary_base_addr + chunk_id * self._primary_row_stride, + size=self._primary_row_stride, + device_Id=0, + info="", + )