[KV Offload] Avoid quadratic event polling

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]

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

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

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: