From a9f7b2d41c5e92f0d60ebdd3dbe04e627019c179 Mon Sep 17 00:00:00 2001 From: Change72 Date: Mon, 22 Jun 2026 00:27:46 -0700 Subject: [PATCH] [feature][kv_offload] Self-describing KV events for OffloadingConnector (#43468) Signed-off-by: Change72 Co-authored-by: Claude --- docs/features/kv_offloading_usage.md | 1 + .../unit/offloading_connector/test_events.py | 355 ++++++++++++++++++ .../offloading_connector/test_scheduler.py | 29 -- .../unit/offloading_connector/utils.py | 17 +- .../kv_connector/v1/offloading/events.py | 286 ++++++++++++++ .../kv_connector/v1/offloading/scheduler.py | 51 ++- vllm/v1/kv_offload/base.py | 19 + vllm/v1/kv_offload/cpu/spec.py | 7 +- vllm/v1/kv_offload/tiering/spec.py | 17 +- 9 files changed, 720 insertions(+), 62 deletions(-) create mode 100644 tests/v1/kv_connector/unit/offloading_connector/test_events.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 93da2ed036..8ef5d6c63d 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -74,6 +74,7 @@ vllm serve \ | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. | | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). | | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. | +| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). | ## Secondary Tiers diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py new file mode 100644 index 0000000000..9e5f564bba --- /dev/null +++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock + +import pytest +import torch + +from tests.v1.kv_connector.unit.utils import create_vllm_config +from vllm.config import KVEventsConfig, KVTransferConfig +from vllm.distributed.kv_events import BlockRemoved, BlockStored +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import ( + OffloadingEventGroupSpec, + OffloadingEventsTracker, +) +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( + GroupOffloadConfig, +) +from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheSpecKind, +) +from vllm.v1.kv_offload.base import ( + OffloadingEvent, + OffloadingKVEventsConfig, + OffloadKey, + make_offload_key, +) +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec +from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec + +_CPU_MEDIUM = CPULoadStoreSpec.medium() +_FULL_ATTENTION_EVENT_SPEC = OffloadingEventGroupSpec( + kv_cache_spec_kind=KVCacheSpecKind.FULL_ATTENTION.value, + kv_cache_spec_sliding_window=None, +) + + +def _tracker( + *, + enable_kv_cache_events: bool = True, + self_describing_kv_events: bool = True, +) -> OffloadingEventsTracker: + return OffloadingEventsTracker( + OffloadingKVEventsConfig( + enable_kv_cache_events=enable_kv_cache_events, + self_describing_kv_events=self_describing_kv_events, + ) + ) + + +def _hash(i: int) -> BlockHash: + return BlockHash(str(i).encode()) + + +def _wire_hash(block_hash: BlockHash): + return maybe_convert_block_hash(block_hash) + + +def _request(*, block_hashes: list[BlockHash], token_count: int): + req = MagicMock() + req.block_hashes = block_hashes + req.all_token_ids = list(range(1, token_count + 1)) + req.lora_request = None + return req + + +def _group_config( + *, + group_idx: int = 0, + block_size: int = 4, + block_size_factor: int = 1, + sliding_window_size_in_blocks: int | None = None, +) -> GroupOffloadConfig: + return GroupOffloadConfig( + group_idx=group_idx, + gpu_block_size=block_size, + offloaded_block_size=block_size * block_size_factor, + hash_block_size_factor=block_size_factor, + sliding_window_size_in_blocks=sliding_window_size_in_blocks, + kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC, + ) + + +def _record_chunks( + tracker: OffloadingEventsTracker, + req, + group_config: GroupOffloadConfig, + num_chunks: int, +) -> list[OffloadKey]: + keys: list[OffloadKey] = [] + hbf = group_config.hash_block_size_factor + for chunk_idx in range(num_chunks): + tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1] + assert tail_hash is not None + key = make_offload_key(tail_hash, group_config.group_idx) + tracker.record_store(req, group_config, chunk_idx, key) + keys.append(key) + return keys + + +def _stored_event(keys: list[OffloadKey]) -> OffloadingEvent: + return OffloadingEvent(keys=keys, medium=_CPU_MEDIUM, removed=False) + + +def _removed_event(keys: list[OffloadKey]) -> OffloadingEvent: + return OffloadingEvent(keys=keys, medium=_CPU_MEDIUM, removed=True) + + +def test_take_events_publishes_routable_block_stored(): + block_size = 4 + tracker = _tracker() + group_config = _group_config(block_size=block_size) + req = _request( + block_hashes=[_hash(i) for i in range(6)], + token_count=block_size * 6, + ) + keys = _record_chunks(tracker, req, group_config, num_chunks=6) + + batch1 = list(tracker.take_events([_stored_event(keys[:3])])) + assert len(batch1) == 3 + + for i, event in enumerate(batch1): + assert isinstance(event, BlockStored) + assert event.medium == _CPU_MEDIUM + assert event.block_hashes == [_wire_hash(_hash(i))] + assert event.block_size == block_size + assert event.token_ids == list( + range(i * block_size + 1, (i + 1) * block_size + 1) + ) + if i == 0: + assert event.parent_block_hash is None + else: + assert event.parent_block_hash == _wire_hash(_hash(i - 1)) + assert event.lora_id is None + assert event.lora_name is None + assert event.extra_keys is None + assert event.group_idx == 0 + assert event.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value + assert event.kv_cache_spec_sliding_window is None + + batch2 = list(tracker.take_events([_stored_event(keys[3:])])) + assert len(batch2) == 3 + assert batch2[0].parent_block_hash == batch1[-1].block_hashes[-1] + + assert len(tracker._pending_event_metadata) == 6 + + +def test_take_events_factor_gt_1_chunk_store_and_remove(): + block_size = 4 + block_size_factor = 3 + tracker = _tracker() + group_config = _group_config( + block_size=block_size, block_size_factor=block_size_factor + ) + req = _request( + block_hashes=[_hash(i) for i in range(6)], + token_count=block_size * block_size_factor * 2, + ) + keys = _record_chunks(tracker, req, group_config, num_chunks=2) + + stored = list(tracker.take_events([_stored_event(keys)])) + assert len(stored) == 2 + + expected_hashes = [] + for chunk_idx, event in enumerate(stored): + assert isinstance(event, BlockStored) + expected_chunk_hashes = [ + _wire_hash(_hash(i)) + for i in range( + chunk_idx * block_size_factor, + (chunk_idx + 1) * block_size_factor, + ) + ] + assert event.block_hashes == expected_chunk_hashes + assert event.block_size == block_size + assert len(event.token_ids) == block_size * block_size_factor + if chunk_idx == 0: + assert event.parent_block_hash is None + else: + assert event.parent_block_hash == _wire_hash(_hash(block_size_factor - 1)) + expected_hashes.extend(expected_chunk_hashes) + + assert len(tracker._pending_event_metadata) == 2 + + removed = list(tracker.take_events([_removed_event(keys)])) + assert len(removed) == 1 + assert isinstance(removed[0], BlockRemoved) + assert removed[0].block_hashes == expected_hashes + assert removed[0].medium == _CPU_MEDIUM + assert removed[0].group_idx == 0 + assert not tracker._pending_event_metadata + + +def test_take_events_factor_gt_1_store_is_order_independent(): + block_size_factor = 3 + tracker = _tracker() + group_config = _group_config(block_size_factor=block_size_factor) + req = _request( + block_hashes=[_hash(i) for i in range(6)], + token_count=4 * block_size_factor * 2, + ) + keys = _record_chunks(tracker, req, group_config, num_chunks=2) + unknown_key = make_offload_key(_hash(12345), 0) + + events = list(tracker.take_events([_stored_event([keys[1], unknown_key, keys[0]])])) + + assert len(events) == 3 + chunk1, placeholder, chunk0 = events + assert [len(event.block_hashes) for event in events] == [3, 1, 3] + assert placeholder.block_size == 0 + assert placeholder.token_ids == [] + assert chunk0.parent_block_hash is None + assert chunk1.parent_block_hash == chunk0.block_hashes[-1] + + +def test_take_events_opt_out_keeps_placeholders(): + tracker = _tracker(self_describing_kv_events=False) + group_config = _group_config() + req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12) + keys = _record_chunks(tracker, req, group_config, num_chunks=3) + + assert not tracker.self_describing_enabled + assert not tracker._pending_event_metadata + + events = list( + tracker.take_events( + [ + _stored_event(keys), + _removed_event(keys), + ] + ) + ) + assert len(events) == 4 + for event in events[:3]: + assert isinstance(event, BlockStored) + assert event.block_size == 0 + assert event.token_ids == [] + assert event.parent_block_hash is None + assert isinstance(events[3], BlockRemoved) + assert len(events[3].block_hashes) == 3 + + +def test_record_store_skips_sliding_window_group(): + tracker = _tracker() + group_config = _group_config(sliding_window_size_in_blocks=2) + req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12) + keys = _record_chunks(tracker, req, group_config, num_chunks=3) + + assert not tracker._pending_event_metadata + + events = list(tracker.take_events([_stored_event(keys[:1])])) + assert len(events) == 1 + assert isinstance(events[0], BlockStored) + assert events[0].block_size == 0 + + +def test_take_events_groups_removed_hashes_by_kv_group(): + tracker = _tracker() + group0_config = _group_config(group_idx=0, block_size_factor=2) + group1_config = _group_config(group_idx=1, block_size_factor=2) + req0 = _request(block_hashes=[_hash(0), _hash(1)], token_count=8) + req1 = _request(block_hashes=[_hash(10), _hash(11)], token_count=8) + key0 = _record_chunks(tracker, req0, group0_config, num_chunks=1)[0] + key1 = _record_chunks(tracker, req1, group1_config, num_chunks=1)[0] + + removed = list(tracker.take_events([_removed_event([key0, key1])])) + + assert len(removed) == 2 + by_group = {event.group_idx: event.block_hashes for event in removed} + assert by_group == { + 0: [_wire_hash(_hash(0)), _wire_hash(_hash(1))], + 1: [_wire_hash(_hash(10)), _wire_hash(_hash(11))], + } + + +def test_take_events_supports_restore_after_eviction(): + block_size = 4 + tracker = _tracker() + group_config = _group_config(block_size=block_size) + req = _request(block_hashes=[_hash(0)], token_count=block_size) + key = _record_chunks(tracker, req, group_config, num_chunks=1)[0] + + first_store = list(tracker.take_events([_stored_event([key])])) + assert len(first_store) == 1 + assert isinstance(first_store[0], BlockStored) + assert first_store[0].token_ids == [1, 2, 3, 4] + + removed = list(tracker.take_events([_removed_event([key])])) + assert len(removed) == 1 + assert isinstance(removed[0], BlockRemoved) + assert not tracker._pending_event_metadata + + req.all_token_ids = [5, 6, 7, 8] + tracker.record_store(req, group_config, offload_block_idx=0, offload_key=key) + + second_store = list(tracker.take_events([_stored_event([key])])) + assert len(second_store) == 1 + assert isinstance(second_store[0], BlockStored) + assert second_store[0].token_ids == [5, 6, 7, 8] + + +def test_reset_cache_clears_side_table(): + tracker = _tracker() + group_config = _group_config() + req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12) + _record_chunks(tracker, req, group_config, num_chunks=3) + + assert tracker._pending_event_metadata + + tracker.reset() + + assert not tracker._pending_event_metadata + + +def test_tiering_rejects_self_describing_kv_events(): + vllm_config = create_vllm_config( + block_size=4, + max_num_batched_tokens=16, + disable_hybrid_kv_cache_manager=False, + ) + vllm_config.kv_transfer_config = KVTransferConfig( + kv_connector="OffloadingConnector", + kv_role="kv_both", + kv_connector_extra_config={ + "spec_name": "TieringOffloadingSpec", + "cpu_bytes_to_use": 1 << 20, + "self_describing_kv_events": True, + "secondary_tiers": [{"type": "example"}], + }, + ) + vllm_config.kv_events_config = KVEventsConfig( + enable_kv_cache_events=True, + publisher="null", + ) + kv_cache_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=4, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ) + + with pytest.raises(ValueError, match="TieringOffloadingSpec"): + TieringOffloadingSpec(vllm_config, kv_cache_config) 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 7b8f6119f5..e4e5c50ecd 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable from types import SimpleNamespace from unittest.mock import MagicMock @@ -12,19 +11,16 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import ( to_keys, ) from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID -from vllm.distributed.kv_events import BlockRemoved, BlockStored from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( OffloadingConnectorScheduler, RequestOffloadState, ) -from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, SlidingWindowSpec, ) from vllm.v1.kv_offload.base import ( - OffloadingEvent, OffloadingManager, OffloadPolicy, ReqContext, @@ -146,31 +142,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5)) - # test take_events - def to_hashes(int_hashes: list[int]) -> list[BlockHash]: - return [BlockHash(str(i).encode()) for i in int_hashes] - - def take_events() -> Iterable[OffloadingEvent]: - yield OffloadingEvent(keys=to_keys([1, 2, 3]), medium="A", removed=False) - yield OffloadingEvent(keys=to_keys([4, 5, 6]), medium="B", removed=True) - - runner.manager.take_events.side_effect = take_events - events = list(runner.scheduler_connector.take_events()) - assert len(events) == 2 - event = events[0] - assert isinstance(event, BlockStored) - assert event.block_hashes == to_hashes([1, 2, 3]) - assert event.block_size == 0 - assert event.medium == "A" - assert event.token_ids == [] - assert event.parent_block_hash is None - assert event.lora_id is None - assert event.lora_name is None - event = events[1] - assert isinstance(event, BlockRemoved) - assert event.block_hashes == to_hashes([4, 5, 6]) - assert event.medium == "B" - @pytest.mark.parametrize("async_scheduling", [True, False]) def test_request_preemption(request_runner, async_scheduling: bool): diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 4464531914..95980690d6 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -14,7 +14,12 @@ from tests.v1.kv_connector.unit.utils import ( create_vllm_config, ) from vllm import SamplingParams -from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config +from vllm.config import ( + KVEventsConfig, + KVTransferConfig, + VllmConfig, + set_current_vllm_config, +) from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( OffloadingConnectorMetadata, @@ -198,6 +203,9 @@ class RequestRunner: "spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils", # noqa: E501 # Preserve legacy behavior for tests; new opt-in tests override. "offload_prompt_only": False, + # Exercise the self-describing KV events path by default; + # opt-out tests override this to cover the legacy placeholders. + "self_describing_kv_events": True, } if block_size_factor > 1: extra_config["block_size"] = block_size * block_size_factor @@ -209,6 +217,13 @@ class RequestRunner: kv_role="kv_both", kv_connector_extra_config=extra_config, ) + vllm_config.kv_events_config = KVEventsConfig( + # Enable so the offloading events tracker is active, but use the + # null publisher: these tests drain take_events directly and a + # real ZMQ publisher would bind a port per test. + enable_kv_cache_events=True, + publisher="null", + ) if kv_cache_groups is None: kv_cache_groups = [ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py new file mode 100644 index 0000000000..410f84c50d --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Self-describing KV cache events for the offloading connector. + +The OffloadingManager identifies an offloaded chunk only by its OffloadKey, +so its raw events carry no token ids, parent hash, or block size. +:class:`OffloadingEventsTracker` snapshots each chunk's full ``BlockStored`` +payload while the ``Request`` is alive and publishes stores as block-granular +payloads: a chunk event may carry multiple constituent per-block hashes, and +evictions fan out to the same hashes. Chunks overlapping a non-chunk-aligned +shared prefix re-announce the shared hashes once per chunk; consumers are +expected to deduplicate (reference-count) repeated store/remove announcements +of the same hash. Opt-in via +``kv_connector_extra_config["self_describing_kv_events"]``; inert unless +KV cache events are enabled. See the PR description for the full design. +""" + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, NamedTuple + +from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent +from vllm.logger import init_logger +from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash +from vllm.v1.kv_cache_interface import ( + KVCacheGroupSpec, + get_kv_cache_spec_kind, + get_kv_cache_spec_sliding_window, +) +from vllm.v1.kv_offload.base import ( + OffloadingEvent, + OffloadingKVEventsConfig, + OffloadKey, + get_offload_block_hash, + get_offload_group_idx, +) +from vllm.v1.request import Request + +if TYPE_CHECKING: + from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( + GroupOffloadConfig, + ) + +logger = init_logger(__name__) + + +class OffloadingEventGroupSpec(NamedTuple): + kv_cache_spec_kind: str | None + kv_cache_spec_sliding_window: int | None + + +def get_offloading_event_group_spec( + kv_cache_group: KVCacheGroupSpec, +) -> OffloadingEventGroupSpec: + kv_cache_spec = kv_cache_group.kv_cache_spec + return OffloadingEventGroupSpec( + kv_cache_spec_kind=get_kv_cache_spec_kind(kv_cache_spec).value, + kv_cache_spec_sliding_window=get_kv_cache_spec_sliding_window(kv_cache_spec), + ) + + +@dataclass(slots=True) +class _OffloadEventMetadata: + """BlockStored payload snapshot for one OffloadKey, captured at store + time and kept until the matching eviction event. ``medium`` is forwarded + from the OffloadingEvent.""" + + # The chunk's constituent block hashes; the last one is the OffloadKey. + block_hashes: tuple[BlockHash, ...] + parent_block_hash: BlockHash | None + token_ids: tuple[int, ...] + block_size: int + lora_id: int | None + lora_name: str | None + # Deferred: needs the same incremental curr_mm_idx handling as GPU events. + extra_keys: tuple[tuple[Any, ...] | None, ...] | None + group_idx: int + kv_cache_spec: OffloadingEventGroupSpec + + +class OffloadingEventsTracker: + """Tracks offloaded chunks' KV event payloads from store to eviction. + + The scheduler calls :meth:`record_store` from ``_build_store_jobs`` + while the ``Request`` is available, and routes the manager's raw + :class:`OffloadingEvent` stream through :meth:`take_events`. All state + is bounded by the CPU pool capacity and cleared by :meth:`reset`. + """ + + def __init__(self, config: OffloadingKVEventsConfig): + self.config = config + self.self_describing_enabled = ( + config.enable_kv_cache_events and config.self_describing_kv_events + ) + + # OffloadKey -> payload snapshot, kept until the eviction event so + # BlockRemoved can fan out. Bounded: one entry per offloaded chunk. + self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {} + + def record_store( + self, + req: Request, + group_config: "GroupOffloadConfig", + offload_block_idx: int, + offload_key: OffloadKey, + ) -> None: + """Snapshot the KV cache event payload for one offloaded chunk. + + No-op when self-describing event capture is disabled or for + sliding-window / SSM groups, which keep the legacy placeholder payload. + """ + if not self.self_describing_enabled: + return + if group_config.sliding_window_size_in_blocks is not None: + return + meta = self._build_event_metadata(req, group_config, offload_block_idx) + self._pending_event_metadata[offload_key] = meta + + def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]: + """Translate raw OffloadingEvents into self-describing KV events. + + Complete metadata is available only for full-attention groups when + the tracker is enabled. Other shapes retain the legacy placeholder + payload so consumers can ignore them. + + Yields: + ``BlockStored`` or ``BlockRemoved`` events corresponding to + the underlying :class:`OffloadingEvent` stream. + """ + for event in events: + if event.removed: + yield from self._take_removed_event(event) + else: + yield from self._take_stored_event(event) + + def reset(self) -> None: + """Drop all tracked state; pending payloads are stale after a + manager cache reset.""" + self._pending_event_metadata.clear() + + def _build_event_metadata( + self, + req: Request, + group_config: "GroupOffloadConfig", + offload_block_idx: int, + ) -> _OffloadEventMetadata: + """Build the payload snapshot for one offloaded chunk: its + constituent per-block hashes, the whole chunk's tokens, and the + per-block ``block_size``.""" + hbf = group_config.hash_block_size_factor + assert hbf > 0 + assert offload_block_idx >= 0 + # per-block token count (= the GPU/hash block size) + sub_block_size = group_config.offloaded_block_size // hbf + # chunk c covers hash-blocks [c*hbf, (c+1)*hbf); its tail block's hash + # is the chunk's OffloadKey. + first_hash_idx = offload_block_idx * hbf + last_hash_idx = first_hash_idx + hbf + assert first_hash_idx >= 0 + assert last_hash_idx <= len(req.block_hashes) + chunk_hashes: list[BlockHash] = [] + for block_hash in req.block_hashes[first_hash_idx:last_hash_idx]: + assert block_hash is not None + chunk_hashes.append(block_hash) + assert len(chunk_hashes) == hbf + + if group_config.sliding_window_size_in_blocks is not None: + # record_store filters these out before calling this helper. + raise AssertionError("self-describing events only support full attention") + + parent_block_hash: BlockHash | None + if first_hash_idx == 0: + parent_block_hash = None + else: + parent_block_hash = req.block_hashes[first_hash_idx - 1] + assert parent_block_hash is not None + + tok_start = offload_block_idx * group_config.offloaded_block_size + tok_end = tok_start + group_config.offloaded_block_size + assert tok_end <= len(req.all_token_ids) + token_ids = tuple(req.all_token_ids[tok_start:tok_end]) + + lora_id: int | None = None + lora_name: str | None = None + if req.lora_request is not None: + lora_id = req.lora_request.adapter_id + lora_name = req.lora_request.name + + return _OffloadEventMetadata( + block_hashes=tuple(chunk_hashes), + parent_block_hash=parent_block_hash, + token_ids=token_ids, + block_size=sub_block_size, + lora_id=lora_id, + lora_name=lora_name, + extra_keys=None, + group_idx=group_config.group_idx, + kv_cache_spec=group_config.kv_event_group_spec, + ) + + def _placeholder_stored(self, key: OffloadKey, medium: str) -> BlockStored: + return BlockStored( + block_hashes=[ + maybe_convert_block_hash(BlockHash(get_offload_block_hash(key))) + ], + parent_block_hash=None, + token_ids=[], + lora_id=None, + block_size=0, + medium=medium, + lora_name=None, + group_idx=get_offload_group_idx(key), + ) + + def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]: + # Metadata is read, NOT popped: the entry must survive until the + # eviction event so BlockRemoved can fan out to the same hashes. + # Events are self-contained (own parent), so key order is free. + for key in event.keys: + meta = self._pending_event_metadata.get(key) + if meta is None: + if self.self_describing_enabled: + # Expected for unsupported shapes; warn once only. + logger.warning_once( + "OffloadingEventsTracker: no event metadata for " + "offload key during BlockStored emission; emitting a " + "placeholder payload. Expected for non-full-attention " + "groups; otherwise indicates a missing populate path." + ) + yield self._placeholder_stored(key, event.medium) + continue + + yield BlockStored( + block_hashes=list( + maybe_convert_block_hash(h) for h in meta.block_hashes + ), + parent_block_hash=( + maybe_convert_block_hash(meta.parent_block_hash) + if meta.parent_block_hash is not None + else None + ), + token_ids=list(meta.token_ids), + block_size=meta.block_size, + lora_id=meta.lora_id, + medium=event.medium, + lora_name=meta.lora_name, + extra_keys=( + list(meta.extra_keys) if meta.extra_keys is not None else None + ), + group_idx=meta.group_idx, + kv_cache_spec_kind=meta.kv_cache_spec.kv_cache_spec_kind, + kv_cache_spec_sliding_window=( + meta.kv_cache_spec.kv_cache_spec_sliding_window + ), + ) + + def _take_removed_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]: + # Keep group_idx unambiguous if a manager batch spans groups. + by_group: dict[int, list] = {} + for key in event.keys: + meta = self._pending_event_metadata.pop(key, None) + if meta is not None: + group_idx = meta.group_idx + by_group.setdefault(group_idx, []).extend( + maybe_convert_block_hash(h) for h in meta.block_hashes + ) + else: + if self.self_describing_enabled: + logger.warning_once( + "OffloadingEventsTracker: no event metadata for " + "offload key during BlockRemoved emission; emitting a " + "placeholder removal. Expected if the matching store " + "used the legacy placeholder payload; otherwise " + "indicates missing store metadata." + ) + group_idx = get_offload_group_idx(key) + by_group.setdefault(group_idx, []).append( + maybe_convert_block_hash(BlockHash(get_offload_block_hash(key))) + ) + + for group_idx, hashes in by_group.items(): + yield BlockRemoved( + block_hashes=hashes, + medium=event.medium, + group_idx=group_idx, + ) 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 5527772788..5884186cc9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, field from itertools import islice from typing import Any, NamedTuple -from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent +from vllm.distributed.kv_events import KVCacheEvent from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( @@ -14,6 +14,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( ReqId, TransferJob, ) +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import ( + OffloadingEventGroupSpec, + OffloadingEventsTracker, + get_offloading_event_group_spec, +) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, _TransferMetricName, @@ -36,7 +41,6 @@ from vllm.v1.kv_offload.base import ( OffloadPolicy, ReqContext, RequestOffloadingContext, - get_offload_block_hash, make_offload_key, ) from vllm.v1.outputs import KVConnectorOutput @@ -69,6 +73,9 @@ class GroupOffloadConfig(NamedTuple): gpu_block_size: int offloaded_block_size: int hash_block_size_factor: int + # KV cache spec metadata propagated onto emitted BlockStored events so + # KV-aware consumers can classify and filter the group. + kv_event_group_spec: OffloadingEventGroupSpec # None below means full attention sliding_window_size_in_blocks: int | None # Number of this group's offloaded blocks per full-attention alignment @@ -200,6 +207,9 @@ class SchedulerOffloadConfig(NamedTuple): alignment_block_count=_alignment_block_count( gpu_block_size * spec.block_size_factor, sw ), + kv_event_group_spec=get_offloading_event_group_spec( + spec.kv_cache_config.kv_cache_groups[idx] + ), is_eagle_group=idx in eagle_groups, ) for idx, gpu_block_size in enumerate(spec.gpu_block_size) @@ -361,6 +371,8 @@ class OffloadingConnectorScheduler: # be freed before a request finishes). self._block_id_to_pending_jobs: dict[int, set[int]] = {} + self._events_tracker = OffloadingEventsTracker(spec.kv_events_config) + def _generate_job_id(self) -> int: job_id = self._job_counter self._job_counter += 1 @@ -934,6 +946,11 @@ class OffloadingConnectorScheduler: continue offloaded_block_idx = start_block_idx + idx + + self._events_tracker.record_store( + req, group_config, offloaded_block_idx, offload_key + ) + gpu_block_idx = offloaded_block_idx * block_size_factor for i in range(block_size_factor): block_id = block_ids[gpu_block_idx + i] @@ -1184,25 +1201,17 @@ class OffloadingConnectorScheduler: return False, None def take_events(self) -> Iterable[KVCacheEvent]: - """Take the KV cache events from the connector. + """Drain pending KV cache events. - Returns: - A list of KV cache events. + Complete metadata is available only when self-describing KV events + are enabled, and only for full-attention groups. Other shapes retain + the previous placeholder payload so consumers can ignore them. + + Yields: + ``BlockStored`` or ``BlockRemoved`` events corresponding to + the underlying :class:`OffloadingEvent` stream. """ - for event in self.manager.take_events(): - block_hashes = [get_offload_block_hash(key) for key in event.keys] - if event.removed: - yield BlockRemoved(block_hashes=block_hashes, medium=event.medium) - else: - yield BlockStored( - block_hashes=block_hashes, - parent_block_hash=None, - token_ids=[], - lora_id=None, - block_size=0, - medium=event.medium, - lora_name=None, - ) + yield from self._events_tracker.take_events(self.manager.take_events()) def reset_cache(self) -> None: """Reset the offloading manager cache, evicting all stored blocks.""" @@ -1238,6 +1247,10 @@ class OffloadingConnectorScheduler: self._jobs.clear() self._block_id_to_pending_jobs.clear() + # The manager pool is empty; pending event payloads and announced + # reference counts are stale. + self._events_tracker.reset() + # Note: _current_batch_jobs_to_flush is intentionally NOT cleared. # The load flush IDs collected above must be delivered to workers. if self._blocks_being_loaded is not None: diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 904003bdcb..d410f42701 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -147,6 +147,16 @@ class OffloadingHistogramMetadata(OffloadingMetricMetadata): buckets: tuple[float, ...] | None = None +@dataclass(frozen=True) +class OffloadingKVEventsConfig: + # Global vLLM KV event publishing flag. When false, connector-specific + # event capture must stay inert because take_events() is not drained. + enable_kv_cache_events: bool + # OffloadingConnector opt-in for self-describing BlockStored payloads. + # Effective only when enable_kv_cache_events is true. + self_describing_kv_events: bool + + class OffloadingManager(ABC): @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: @@ -443,6 +453,15 @@ class OffloadingSpec(ABC): kv_transfer_config = vllm_config.kv_transfer_config assert kv_transfer_config is not None self.extra_config = kv_transfer_config.kv_connector_extra_config + kv_events_config = vllm_config.kv_events_config + self.kv_events_config = OffloadingKVEventsConfig( + enable_kv_cache_events=( + kv_events_config is not None and kv_events_config.enable_kv_cache_events + ), + self_describing_kv_events=bool( + self.extra_config.get("self_describing_kv_events", False) + ), + ) # When True, only prompt (prefill) blocks are offloaded; decode-phase # blocks (KV generated after the prompt) are skipped. Useful when prior diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 9b1dff24a8..7d3ba9c753 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -114,11 +114,6 @@ class CPUOffloadingSpec(OffloadingSpec): @override def get_manager(self) -> OffloadingManager: if not self._manager: - kv_events_config = self.vllm_config.kv_events_config - enable_events = ( - kv_events_config is not None and kv_events_config.enable_kv_cache_events - ) - # store_threshold: how many times a block must appear in lookup() # before it is eligible for CPU offloading. Values < 2 disable # filtering (a threshold of 1 equals no filter; 0 is the default). @@ -130,7 +125,7 @@ class CPUOffloadingSpec(OffloadingSpec): self._manager = CPUOffloadingManager( num_blocks=self.num_blocks, cache_policy=self.eviction_policy, # type: ignore[arg-type] - enable_events=enable_events, + enable_events=self.kv_events_config.enable_kv_cache_events, store_threshold=store_threshold, max_tracker_size=max_tracker_size, ) diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index f223d81aa5..e9dd68c44f 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -69,6 +69,14 @@ class TieringOffloadingSpec(CPUOffloadingSpec): super().__init__(vllm_config, kv_cache_config) # Redeclare for mypy: parent sets this but `--follow-imports skip` hides it self._manager: OffloadingManager | None = None + if self.kv_events_config.self_describing_kv_events: + raise ValueError( + "self_describing_kv_events is not supported by " + "TieringOffloadingSpec. Tier promotions can emit primary-tier " + "store events that do not correspond to GPU store jobs, so the " + "current self-describing side table cannot describe them " + "correctly." + ) # Parse secondary tier configurations self.secondary_tier_configs = self.extra_config.get("secondary_tiers", []) @@ -91,11 +99,6 @@ class TieringOffloadingSpec(CPUOffloadingSpec): TieringOffloadingManager instance """ if not self._manager: - kv_events_config = self.vllm_config.kv_events_config - enable_events = ( - kv_events_config is not None and kv_events_config.enable_kv_cache_events - ) - # Create scheduler-side SharedOffloadRegion (rank=None) so the # primary tier can eagerly create a memoryview over _base. scheduler_mmap = SharedOffloadRegion( @@ -111,7 +114,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec): primary_tier = CPUPrimaryTierOffloadingManager( num_blocks=self.num_blocks, cache_policy=self.eviction_policy, # type: ignore[arg-type] - enable_events=enable_events, + enable_events=self.kv_events_config.enable_kv_cache_events, mmap_region=scheduler_mmap, ) @@ -143,7 +146,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec): tiering_manager = TieringOffloadingManager( primary_tier=primary_tier, secondary_tiers=secondary_tiers, - enable_events=enable_events, + enable_events=self.kv_events_config.enable_kv_cache_events, ) if int(self.extra_config.get("store_threshold", 0)) >= 2: raise ValueError(