From 71a6be79da9002133d2513eccb595acbfbbdf2e3 Mon Sep 17 00:00:00 2001 From: felix h Date: Mon, 7 Sep 2026 18:32:59 +0200 Subject: [PATCH] [KV Offload] Avoid quadratic event polling Use ordered deques so transfer completion polling scales with completed prefixes instead of the total event backlog. Co-authored-by: OpenAI Codex Signed-off-by: felix h --- .../simple_kv_offload/test_worker_events.py | 64 +++++++++++++++++++ vllm/v1/simple_kv_offload/copy_backend.py | 3 +- vllm/v1/simple_kv_offload/disk_backend.py | 3 +- vllm/v1/simple_kv_offload/worker.py | 41 +++++++----- 4 files changed, 92 insertions(+), 19 deletions(-) create mode 100644 tests/v1/simple_kv_offload/test_worker_events.py diff --git a/tests/v1/simple_kv_offload/test_worker_events.py b/tests/v1/simple_kv_offload/test_worker_events.py new file mode 100644 index 000000000000..2a64af8e1190 --- /dev/null +++ b/tests/v1/simple_kv_offload/test_worker_events.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host-side event bookkeeping tests for SimpleCPUOffloadConnector.""" + +from typing import cast + +import pytest +import torch + +from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadMetadata +from vllm.v1.simple_kv_offload.worker import SimpleCPUOffloadWorker + + +class _FakeEvent: + def __init__(self, complete: bool) -> None: + self.complete = complete + self.query_count = 0 + + def query(self) -> bool: + self.query_count += 1 + return self.complete + + +@pytest.mark.skip_global_cleanup +def test_get_finished_reports_only_completed_prefixes(): + worker = SimpleCPUOffloadWorker( + vllm_config=None, kv_cache_config=None, cpu_capacity_bytes=0 + ) + load_events = [_FakeEvent(True), _FakeEvent(False), _FakeEvent(True)] + store_events = [_FakeEvent(True), _FakeEvent(False), _FakeEvent(True)] + worker._load_events.extend( + (i, cast(torch.Event, event)) for i, event in enumerate(load_events) + ) + worker._store_events.extend( + (i, cast(torch.Event, event)) for i, event in enumerate(store_events) + ) + + load_event_to_reqs = {i: [f"req-{i}"] for i in range(len(load_events))} + for event_idx in range(len(load_events)): + worker.bind_connector_metadata( + SimpleCPUOffloadMetadata( + load_event=event_idx, + load_event_to_reqs=load_event_to_reqs, + store_event=event_idx, + ) + ) + + assert worker.get_finished(set()) == (None, {"req-0"}) + worker_meta = worker.build_connector_worker_meta() + assert worker_meta is not None + assert worker_meta.completed_store_events == {0: 1} + assert [event.query_count for event in load_events] == [1, 1, 0] + assert [event.query_count for event in store_events] == [1, 1, 0] + + load_events[1].complete = True + store_events[1].complete = True + assert worker.get_finished(set()) == (None, {"req-1", "req-2"}) + worker_meta = worker.build_connector_worker_meta() + assert worker_meta is not None + assert worker_meta.completed_store_events == {1: 1, 2: 1} + assert worker.build_connector_worker_meta() is None + assert worker.get_finished(set()) == (None, None) + assert [event.query_count for event in load_events] == [1, 2, 1] + assert [event.query_count for event in store_events] == [1, 2, 1] diff --git a/vllm/v1/simple_kv_offload/copy_backend.py b/vllm/v1/simple_kv_offload/copy_backend.py index 58de7a7e9eff..da5608583dfc 100644 --- a/vllm/v1/simple_kv_offload/copy_backend.py +++ b/vllm/v1/simple_kv_offload/copy_backend.py @@ -6,6 +6,7 @@ import queue import threading +from collections import deque import torch @@ -74,7 +75,7 @@ def launch_copy( dst_blocks: list[int], is_store: bool, event_idx: int, - events_list: list[tuple[int, torch.Event]], + events_list: list[tuple[int, torch.Event]] | deque[tuple[int, torch.Event]], wait_event: torch.Event | None = None, ) -> None: params = self._store_params if is_store else self._load_params diff --git a/vllm/v1/simple_kv_offload/disk_backend.py b/vllm/v1/simple_kv_offload/disk_backend.py index f39be7d05f56..fca91ab1c70f 100644 --- a/vllm/v1/simple_kv_offload/disk_backend.py +++ b/vllm/v1/simple_kv_offload/disk_backend.py @@ -13,6 +13,7 @@ import os import queue import threading +from collections import deque import torch @@ -189,7 +190,7 @@ def launch_copy( dst_blocks: list[int], is_store: bool, event_idx: int, - events_list: list[tuple[int, torch.Event]], + events_list: list[tuple[int, torch.Event]] | deque[tuple[int, torch.Event]], wait_event: torch.Event | None = None, ) -> None: q = self._store_queue if is_store else self._load_queue diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py index d49de261e575..640312b2e95a 100644 --- a/vllm/v1/simple_kv_offload/worker.py +++ b/vllm/v1/simple_kv_offload/worker.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Worker-side handler for SimpleCPUOffloadConnector.""" +from collections import deque from typing import TYPE_CHECKING import torch @@ -57,9 +58,9 @@ def __init__( self._backend: DmaCopyBackend | DiskBackend | None = None - # Ordered (event_idx, Event). Events pre-allocated on main thread. - self._load_events: list[tuple[int, torch.Event]] = [] - self._store_events: list[tuple[int, torch.Event]] = [] + # Ordered (event_idx, Event) queues populated by the backend threads. + self._load_events: deque[tuple[int, torch.Event]] = deque() + self._store_events: deque[tuple[int, torch.Event]] = deque() # High-water marks: highest event_idx completed per stream. # When the event list is empty, the hwm covers all prior events. self._load_hwm: int = -1 @@ -72,9 +73,9 @@ def __init__( # (get_finished runs once per step, copy queue is FIFO). self._store_compute_done: torch.Event | None = None - # Pending event index sets, populated in bind_connector_metadata - self._pending_load_event_indices: set[int] = set() - self._pending_store_event_indices: set[int] = set() + # Ordered pending event indices, populated in bind_connector_metadata. + self._pending_load_event_indices: deque[int] = deque() + self._pending_store_event_indices: deque[int] = deque() # Completed store events to report via build_connector_worker_meta self._completed_store_events: dict[int, int] = {} @@ -223,9 +224,9 @@ def _init_cpu_mode( def bind_connector_metadata(self, metadata: SimpleCPUOffloadMetadata) -> None: self._connector_metadata = metadata if metadata.load_event >= 0: - self._pending_load_event_indices.add(metadata.load_event) + self._pending_load_event_indices.append(metadata.load_event) if metadata.store_event >= 0: - self._pending_store_event_indices.add(metadata.store_event) + self._pending_store_event_indices.append(metadata.store_event) def clear_connector_metadata(self) -> None: self._connector_metadata = None @@ -288,8 +289,11 @@ def get_finished( if self._pending_load_event_indices: load_wm = self._poll_stream_events(is_store=False) - for j in [j for j in self._pending_load_event_indices if j <= load_wm]: - self._pending_load_event_indices.discard(j) + while ( + self._pending_load_event_indices + and self._pending_load_event_indices[0] <= load_wm + ): + j = self._pending_load_event_indices.popleft() req_ids = ( metadata.load_event_to_reqs.get(j) if metadata is not None else None ) @@ -298,8 +302,11 @@ def get_finished( if self._pending_store_event_indices: store_wm = self._poll_stream_events(is_store=True) - for j in [j for j in self._pending_store_event_indices if j <= store_wm]: - self._pending_store_event_indices.discard(j) + while ( + self._pending_store_event_indices + and self._pending_store_event_indices[0] <= store_wm + ): + j = self._pending_store_event_indices.popleft() self._completed_store_events[j] = 1 return None, finished_recving or None @@ -324,15 +331,15 @@ def handle_preemptions( def _flush_and_sync_all(self) -> None: """Synchronize all in-flight transfer events.""" - for event_idx, event in self._load_events: + while self._load_events: + event_idx, event = self._load_events.popleft() event.synchronize() self._load_hwm = event_idx - self._load_events.clear() - for event_idx, event in self._store_events: + while self._store_events: + event_idx, event = self._store_events.popleft() event.synchronize() self._store_hwm = event_idx - self._store_events.clear() def _poll_stream_events(self, is_store: bool) -> int: """Non-blocking poll for completed events and return the high-water mark.""" @@ -343,7 +350,7 @@ def _poll_stream_events(self, is_store: bool) -> int: if not event.query(): break hwm = event_idx - events.pop(0) + events.popleft() if is_store: self._store_hwm = hwm else: