A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
v1: Support KV events from connectors (#19737)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Sep 1, 2025 at 04:13 UTC
14b4326b9470c098d537cce3834033a7f3b2c386
6 files changed
+44
-3
examples/online_serving/kv_events_subscriber.py
+2
@@ -27,10 +27,12 @@ class BlockStored(KVCacheEvent):
27
token_ids: list[int]
28
block_size: int
29
lora_id: Optional[int]
30
+ medium: Optional[str]
31
32
33
class BlockRemoved(KVCacheEvent):
34
block_hashes: list[int]
35
+ medium: Optional[str]
36
37
38
class AllBlocksCleared(KVCacheEvent):
vllm/distributed/kv_events.py
+5
@@ -40,16 +40,21 @@ class KVCacheEvent(
40
"""Base class for all KV cache-related events"""
41
42
43
+MEDIUM_GPU = "GPU"
44
+
45
+
46
class BlockStored(KVCacheEvent):
47
block_hashes: list[int]
48
parent_block_hash: Optional[int]
49
token_ids: list[int]
50
block_size: int
51
lora_id: Optional[int]
52
+ medium: Optional[str]
53
54
55
class BlockRemoved(KVCacheEvent):
56
block_hashes: list[int]
57
+ medium: Optional[str]
58
59
60
class AllBlocksCleared(KVCacheEvent):
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+13
@@ -19,6 +19,8 @@ The class provides the following primitives:
19
Returns whether KV cache should be freed now or will be
20
freed asynchronously and optionally returns KV transfer
21
params.
22
+ take_events() - returns new KV events that were collected
23
+ by the connector since the last call.
24
25
Worker-side: runs in each worker, loads/saves KV cache to/from
26
the Connector based on the metadata.
@@ -34,6 +36,7 @@ The class provides the following primitives:
36
37
import enum
38
from abc import ABC, abstractmethod
39
+from collections.abc import Iterable
40
from typing import TYPE_CHECKING, Any, Callable, Literal, Optional
41
42
import torch
@@ -45,6 +48,7 @@ from vllm.v1.outputs import KVConnectorOutput
48
if TYPE_CHECKING:
49
from vllm.attention.backends.abstract import AttentionMetadata
50
from vllm.config import VllmConfig
51
+ from vllm.distributed.kv_events import KVCacheEvent
52
from vllm.forward_context import ForwardContext
53
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
54
from vllm.v1.request import Request
@@ -313,6 +317,15 @@ class KVConnectorBase_V1(ABC):
317
"""
318
return False, None
319
320
+ def take_events(self) -> Iterable["KVCacheEvent"]:
321
+ """
322
+ Take the KV cache events from the connector.
323
+
324
+ Yields:
325
+ New KV cache events since the last call.
326
+ """
327
+ return ()
328
+
329
@classmethod
330
def get_required_kvcache_layout(
331
cls, vllm_config: "VllmConfig") -> Optional[str]:
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+6
@@ -1,12 +1,14 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import copy
4
+from collections.abc import Iterable
5
from dataclasses import dataclass
6
from typing import TYPE_CHECKING, Any, Optional
7
8
import torch
9
10
from vllm.config import KVTransferConfig, VllmConfig
11
+from vllm.distributed.kv_events import KVCacheEvent
12
from vllm.distributed.kv_transfer.kv_connector.factory import (
13
KVConnectorFactory)
14
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
@@ -208,6 +210,10 @@ class MultiConnector(KVConnectorBase_V1):
210
211
return async_saves > 0, kv_txfer_params
212
213
+ def take_events(self) -> Iterable[KVCacheEvent]:
214
+ for c in self._connectors:
215
+ yield from c.take_events()
216
+
217
@classmethod
218
def get_required_kvcache_layout(
219
cls, vllm_config: "VllmConfig") -> Optional[str]:
vllm/v1/core/block_pool.py
+6
-3
@@ -4,8 +4,9 @@ from collections import defaultdict
4
from collections.abc import Iterable
5
from typing import Optional
6
7
-from vllm.distributed.kv_events import (AllBlocksCleared, BlockRemoved,
8
- BlockStored, KVCacheEvent)
7
+from vllm.distributed.kv_events import (MEDIUM_GPU, AllBlocksCleared,
8
+ BlockRemoved, BlockStored,
9
+ KVCacheEvent)
10
from vllm.logger import init_logger
11
from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,
12
FreeKVCacheBlockQueue, KVCacheBlock)
@@ -156,6 +157,7 @@ class BlockPool:
157
block_size=block_size,
158
lora_id=request.lora_request.id
159
if request.lora_request else None,
160
+ medium=MEDIUM_GPU,
161
))
162
163
def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
@@ -218,7 +220,8 @@ class BlockPool:
220
# we disable hybrid kv cache manager when kv cache event is
221
# enabled, so there is only one group.
222
self.kv_event_queue.append(
221
- BlockRemoved(block_hashes=[block_hash.get_hash_value()]))
223
+ BlockRemoved(block_hashes=[block_hash.get_hash_value()],
224
+ medium=MEDIUM_GPU))
225
return True
226
227
def touch(self, blocks: tuple[list[KVCacheBlock], ...]) -> None:
vllm/v1/core/sched/scheduler.py
+12
@@ -589,7 +589,19 @@ class Scheduler(SchedulerInterface):
589
meta = self.connector.build_connector_meta(scheduler_output)
590
scheduler_output.kv_connector_metadata = meta
591
592
+ # collect KV cache events from KV cache manager
593
events = self.kv_cache_manager.take_events()
594
+
595
+ # collect KV cache events from connector
596
+ if self.connector is not None:
597
+ connector_events = self.connector.take_events()
598
+ if connector_events:
599
+ if events is None:
600
+ events = list(connector_events)
601
+ else:
602
+ events.extend(connector_events)
603
+
604
+ # publish collected KV cache events
605
if events:
606
batch = KVEventBatch(ts=time.time(), events=events)
607
self.kv_event_publisher.publish(batch)