@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0be6d05b5e
Fork 0

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 15, 12:15 UTC

[V1][Metrics] add support for kv event publishing (#16750)

Signed-off-by: alec-flowers <aflowers@nvidia.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com>

Alec committed Apr 30, 2025 at 16:44 UTC 0be6d05b5ea92dae61cd83f26fea05a48d227802
15 files changed +1183 -51
examples/online_serving/kv_events.sh new
+86
@@ -0,0 +1,86 @@
1 +#!/bin/bash
2 +# This file demonstrates the KV cache event publishing
3 +# We will launch a vllm instances configured to publish KV cache
4 +# events and launch a simple subscriber to log those events.
5 +
6 +set -xe
7 +
8 +echo "🚧🚧 Warning: The usage of KV cache events is experimental and subject to change 🚧🚧"
9 +sleep 1
10 +
11 +MODEL_NAME=${HF_MODEL_NAME:-meta-llama/Meta-Llama-3.1-8B-Instruct}
12 +
13 +# Trap the SIGINT signal (triggered by Ctrl+C)
14 +trap 'cleanup' INT
15 +
16 +# Cleanup function
17 +cleanup() {
18 + echo "Caught Ctrl+C, cleaning up..."
19 + # Cleanup commands
20 + pgrep python | xargs kill -9
21 + pkill -f python
22 + echo "Cleanup complete. Exiting."
23 + exit 0
24 +}
25 +
26 +export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
27 +
28 +# a function that waits vLLM server to start
29 +wait_for_server() {
30 + local port=$1
31 + timeout 1200 bash -c "
32 + until curl -s localhost:${port}/v1/completions > /dev/null; do
33 + sleep 1
34 + done" && return 0 || return 1
35 +}
36 +
37 +vllm serve $MODEL_NAME \
38 + --port 8100 \
39 + --max-model-len 100 \
40 + --enforce-eager \
41 + --gpu-memory-utilization 0.8 \
42 + --trust-remote-code \
43 + --kv-events-config \
44 + '{"enable_kv_cache_events": true, "publisher": "zmq", "topic": "kv-events"}' &
45 +
46 +wait_for_server 8100
47 +
48 +SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
49 +
50 +python3 "$SCRIPT_DIR/kv_events_subscriber.py" &
51 +sleep 1
52 +
53 +# serve two example requests
54 +output1=$(curl -X POST -s http://localhost:8100/v1/completions \
55 +-H "Content-Type: application/json" \
56 +-d '{
57 +"model": "'"$MODEL_NAME"'",
58 +"prompt": "Explain quantum computing in simple terms a 5-year-old could understand.",
59 +"max_tokens": 80,
60 +"temperature": 0
61 +}')
62 +
63 +output2=$(curl -X POST -s http://localhost:8100/v1/completions \
64 +-H "Content-Type: application/json" \
65 +-d '{
66 +"model": "'"$MODEL_NAME"'",
67 +"prompt": "Explain quantum computing in simple terms a 50-year-old could understand.",
68 +"max_tokens": 80,
69 +"temperature": 0
70 +}')
71 +
72 +# Cleanup commands
73 +pkill -9 -u "$USER" -f python
74 +pkill -9 -u "$USER" -f vllm
75 +
76 +sleep 1
77 +
78 +echo "Cleaned up"
79 +
80 +# Print the outputs of the curl requests
81 +echo ""
82 +echo "Output of first request: $output1"
83 +echo "Output of second request: $output2"
84 +
85 +echo "🎉🎉 Successfully finished 2 test requests! 🎉🎉"
86 +echo ""
examples/online_serving/kv_events_subscriber.py new
+114
@@ -0,0 +1,114 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Any, Optional, Union
3 +
4 +import msgspec
5 +import zmq
6 +from msgspec.msgpack import Decoder
7 +
8 +
9 +#
10 +# Types copied from vllm.distributed.kv_events
11 +#
12 +class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True,
13 + gc=False):
14 + ts: float
15 + events: list[Any]
16 +
17 +
18 +class KVCacheEvent(msgspec.Struct,
19 + array_like=True,
20 + omit_defaults=True,
21 + gc=False,
22 + tag=True):
23 + """Base class for all KV cache-related events"""
24 +
25 +
26 +class BlockStored(KVCacheEvent):
27 + block_hashes: list[int]
28 + parent_block_hash: Optional[int]
29 + token_ids: list[int]
30 + block_size: int
31 + lora_id: Optional[int]
32 +
33 +
34 +class BlockRemoved(KVCacheEvent):
35 + block_hashes: list[int]
36 +
37 +
38 +class AllBlocksCleared(KVCacheEvent):
39 + pass
40 +
41 +
42 +class KVEventBatch(EventBatch):
43 + events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]]
44 +
45 +
46 +def process_event(event_batch):
47 + print(f"Received event batch at {event_batch.ts}:")
48 + for event in event_batch.events:
49 + print(f" - {event}")
50 +
51 +
52 +def main():
53 + decoder = Decoder(type=KVEventBatch)
54 + last_seq = -1
55 +
56 + context = zmq.Context()
57 +
58 + # Set up the main subscription socket
59 + sub = context.socket(zmq.SUB)
60 + sub.connect("tcp://localhost:5557")
61 + topic = "kv-events"
62 + sub.setsockopt_string(zmq.SUBSCRIBE, topic)
63 +
64 + # Initialize replay socket
65 + replay = context.socket(zmq.REQ)
66 + replay.connect("tcp://localhost:5558")
67 + poller = zmq.Poller()
68 + poller.register(replay, zmq.POLLIN)
69 +
70 + print("Listening for KV cache events on topic:", topic)
71 +
72 + while True:
73 + try:
74 + if sub.poll(50):
75 + _, seq_bytes, payload = sub.recv_multipart()
76 + seq = int.from_bytes(seq_bytes, "big")
77 +
78 + if last_seq >= 0 and seq > last_seq + 1:
79 + missed = seq - last_seq - 1
80 + print(f"Missed {missed} messages"
81 + f" (last: {last_seq}, current: {seq})")
82 +
83 + replay.send((last_seq + 1).to_bytes(8, "big"))
84 +
85 + while poller.poll(timeout=200):
86 + seq_bytes, replay_payload = replay.recv_multipart()
87 + if not replay_payload:
88 + # End of replay marker is sent as an empty frame
89 + # for the payload
90 + break
91 +
92 + replay_seq = int.from_bytes(seq_bytes, "big")
93 +
94 + if replay_seq > last_seq:
95 + event_batch = decoder.decode(replay_payload)
96 + process_event(event_batch)
97 + last_seq = replay_seq
98 + if replay_seq >= seq - 1:
99 + break
100 +
101 + event_batch = decoder.decode(payload)
102 + process_event(event_batch)
103 +
104 + # ... do other periodic work or check for shutdown ...
105 +
106 + except KeyboardInterrupt:
107 + print("Interrupted")
108 + break
109 + except Exception as e:
110 + print("Error decoding message:", e)
111 +
112 +
113 +if __name__ == "__main__":
114 + main()
tests/distributed/conftest.py new
+145
@@ -0,0 +1,145 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import random
3 +from typing import Optional, Union
4 +
5 +import msgspec
6 +import msgspec.msgpack
7 +import pytest
8 +import zmq
9 +
10 +from vllm.config import KVEventsConfig
11 +from vllm.distributed.kv_events import EventPublisherFactory
12 +
13 +from .test_events import SampleBatch
14 +
15 +
16 +@pytest.fixture
17 +def random_port():
18 + """Generate a random port number for testing"""
19 + return random.randint(10000, 60000)
20 +
21 +
22 +@pytest.fixture
23 +def publisher_config(random_port, request):
24 + """Create a publisher config with inproc transport"""
25 + how = request.param if hasattr(request, "param") else "inproc"
26 +
27 + if how == "inproc":
28 + endpoint = f"inproc://test-{random_port}"
29 + replay_endpoint = endpoint + "-replay"
30 + else:
31 + endpoint = f"tcp://*:{random_port}"
32 + replay_endpoint = f"tcp://*:{random_port + 1}"
33 +
34 + return KVEventsConfig(enable_kv_cache_events=True,
35 + publisher="zmq",
36 + endpoint=endpoint,
37 + replay_endpoint=replay_endpoint,
38 + buffer_steps=100,
39 + hwm=1000,
40 + topic="test")
41 +
42 +
43 +@pytest.fixture
44 +def publisher(publisher_config):
45 + """Create and return a publisher instance"""
46 + pub = EventPublisherFactory.create(publisher_config)
47 + yield pub
48 + pub.shutdown()
49 +
50 +
51 +@pytest.fixture
52 +def subscriber(publisher_config):
53 + """Create and return a subscriber for testing"""
54 + endpoint = publisher_config.endpoint
55 + replay_endpoint = publisher_config.replay_endpoint
56 +
57 + if endpoint.startswith("tcp://*"):
58 + endpoint = endpoint.replace("*", "127.0.0.1")
59 + if replay_endpoint and replay_endpoint.startswith("tcp://*"):
60 + replay_endpoint = replay_endpoint.replace("*", "127.0.0.1")
61 +
62 + sub = MockSubscriber(endpoint, replay_endpoint, publisher_config.topic)
63 + yield sub
64 + sub.close()
65 +
66 +
67 +class MockSubscriber:
68 + """Helper class to receive and verify published events"""
69 +
70 + def __init__(self,
71 + pub_endpoint: str,
72 + replay_endpoint: Optional[str] = None,
73 + topic: str = "",
74 + decode_type=SampleBatch):
75 + self.ctx = zmq.Context.instance()
76 +
77 + # Set up subscriber socket
78 + self.sub = self.ctx.socket(zmq.SUB)
79 + self.sub.setsockopt(zmq.SUBSCRIBE, topic.encode('utf-8'))
80 + self.sub.connect(pub_endpoint)
81 +
82 + # Set up replay socket if provided
83 + self.replay = None
84 + if replay_endpoint:
85 + self.replay = self.ctx.socket(zmq.REQ)
86 + self.replay.connect(replay_endpoint)
87 +
88 + self.topic = topic
89 + self.topic_bytes = topic.encode('utf-8')
90 + self.received_msgs: list[tuple[int, SampleBatch]] = []
91 + self.last_seq = -1
92 + self.decoder = msgspec.msgpack.Decoder(type=decode_type)
93 +
94 + def receive_one(self,
95 + timeout=1000) -> Union[tuple[int, SampleBatch], None]:
96 + """Receive a single message with timeout"""
97 + if not self.sub.poll(timeout):
98 + return None
99 +
100 + topic_bytes, seq_bytes, payload = self.sub.recv_multipart()
101 + assert topic_bytes == self.topic_bytes
102 +
103 + seq = int.from_bytes(seq_bytes, "big")
104 + data = self.decoder.decode(payload)
105 + self.last_seq = seq
106 + self.received_msgs.append((seq, data))
107 + return seq, data
108 +
109 + def request_replay(self, start_seq: int) -> None:
110 + """Request replay of messages starting from start_seq"""
111 + if not self.replay:
112 + raise ValueError("Replay socket not initialized")
113 +
114 + self.replay.send(start_seq.to_bytes(8, "big"))
115 +
116 + def receive_replay(self) -> list[tuple[int, SampleBatch]]:
117 + """Receive replayed messages"""
118 + if not self.replay:
119 + raise ValueError("Replay socket not initialized")
120 +
121 + replayed: list[tuple[int, SampleBatch]] = []
122 + while True:
123 + try:
124 + if not self.replay.poll(1000):
125 + break
126 +
127 + frames = self.replay.recv_multipart()
128 + if not frames or not frames[-1]:
129 + # End of replay marker
130 + break
131 +
132 + seq_bytes, payload = frames
133 + seq = int.from_bytes(seq_bytes, "big")
134 + data = self.decoder.decode(payload)
135 + replayed.append((seq, data))
136 + except zmq.ZMQError as _:
137 + break
138 +
139 + return replayed
140 +
141 + def close(self):
142 + """Clean up resources"""
143 + self.sub.close()
144 + if self.replay:
145 + self.replay.close()
tests/distributed/test_events.py new
+193
@@ -0,0 +1,193 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import threading
3 +import time
4 +
5 +import msgspec
6 +import pytest
7 +
8 +from vllm.distributed.kv_events import (EventBatch, EventPublisherFactory,
9 + NullEventPublisher)
10 +
11 +
12 +class EventSample(
13 + msgspec.Struct,
14 + tag=True, # type: ignore
15 + array_like=True # type: ignore
16 +):
17 + """Test event for publisher testing"""
18 + id: int
19 + value: str
20 +
21 +
22 +class SampleBatch(EventBatch):
23 + """Test event batch for publisher testing"""
24 + events: list[EventSample]
25 +
26 +
27 +def create_test_events(count: int) -> SampleBatch:
28 + """Create a batch of test events"""
29 + events = [EventSample(id=i, value=f"test-{i}") for i in range(count)]
30 + return SampleBatch(ts=time.time(), events=events)
31 +
32 +
33 +def test_basic_publishing(publisher, subscriber):
34 + """Test basic event publishing works"""
35 +
36 + test_batch = create_test_events(5)
37 + publisher.publish(test_batch)
38 +
39 + result = subscriber.receive_one(timeout=1000)
40 + assert result is not None, "No message received"
41 +
42 + seq, received = result
43 + assert seq == 0, "Sequence number mismatch"
44 + assert received.ts == pytest.approx(test_batch.ts,
45 + abs=0.1), ("Timestamp mismatch")
46 + assert len(received.events) == len(
47 + test_batch.events), ("Number of events mismatch")
48 +
49 + for i, event in enumerate(received.events):
50 + assert event.id == i, "Event id mismatch"
51 + assert event.value == f"test-{i}", "Event value mismatch"
52 +
53 +
54 +def test_multiple_events(publisher, subscriber):
55 + """Test publishing and receiving multiple event batches"""
56 + for _ in range(10):
57 + batch = create_test_events(2)
58 + publisher.publish(batch)
59 +
60 + received = []
61 + for _ in range(10):
62 + data = subscriber.receive_one(timeout=100)
63 + if data:
64 + received.append(data)
65 +
66 + assert len(received) == 10, "Number of messages mismatch"
67 + seqs = [seq for seq, _ in received]
68 + assert seqs == list(range(10)), "Sequence numbers mismatch"
69 +
70 +
71 +def test_replay_mechanism(publisher, subscriber):
72 + """Test the replay mechanism works correctly"""
73 + for _ in range(19):
74 + batch = create_test_events(1)
75 + publisher.publish(batch)
76 +
77 + time.sleep(0.5) # Need publisher to process above requests
78 + subscriber.request_replay(10)
79 +
80 + batch = create_test_events(1)
81 + publisher.publish(batch) # 20th message
82 +
83 + replayed = subscriber.receive_replay()
84 +
85 + assert len(replayed) > 0, "No replayed messages received"
86 + seqs = [seq for seq, _ in replayed]
87 + assert all(seq >= 10 for seq in seqs), "Replayed messages not in order"
88 + assert seqs == list(range(min(seqs),
89 + max(seqs) +
90 + 1)), ("Replayed messages not consecutive")
91 +
92 +
93 +def test_buffer_limit(publisher, subscriber, publisher_config):
94 + """Test buffer limit behavior"""
95 + buffer_size = publisher_config.buffer_steps
96 +
97 + # Publish more events than the buffer can hold
98 + for i in range(buffer_size + 10):
99 + batch = create_test_events(1)
100 + publisher.publish(batch)
101 +
102 + time.sleep(0.5) # Need publisher to process above requests
103 + subscriber.request_replay(0)
104 +
105 + batch = create_test_events(1)
106 + publisher.publish(batch)
107 +
108 + replayed = subscriber.receive_replay()
109 +
110 + assert len(replayed) <= buffer_size, "Can't replay more than buffer size"
111 +
112 + oldest_seq = min(seq for seq, _ in replayed)
113 + assert oldest_seq >= 10, "The oldest sequence should be at least 10"
114 +
115 +
116 +def test_topic_filtering(publisher_config):
117 + """
118 + Test that a subscriber only receives messages matching its topic filter
119 + """
120 + publisher_config.replay_endpoint = None
121 +
122 + cfg = publisher_config.model_copy()
123 + cfg.topic = "foo"
124 + pub = EventPublisherFactory.create(cfg)
125 +
126 + from .conftest import MockSubscriber
127 + sub_foo = MockSubscriber(cfg.endpoint, None, "foo")
128 + sub_bar = MockSubscriber(cfg.endpoint, None, "bar")
129 +
130 + try:
131 + time.sleep(0.1)
132 +
133 + for _ in range(3):
134 + pub.publish(create_test_events(1))
135 +
136 + foo_received = [sub_foo.receive_one(timeout=200) for _ in range(3)]
137 + assert all(msg is not None for msg in foo_received), (
138 + "Subscriber with matching topic should receive messages")
139 +
140 + bar_received = [sub_bar.receive_one(timeout=200) for _ in range(3)]
141 + assert all(msg is None for msg in bar_received), (
142 + "Subscriber with non-matching topic should receive no messages")
143 + finally:
144 + pub.shutdown()
145 + sub_foo.close()
146 + sub_bar.close()
147 +
148 +
149 +def test_high_volume(publisher, subscriber):
150 + """Test publishing and receiving a high volume of events"""
151 + num_batches = 10_000
152 + events_per_batch = 100
153 +
154 + # Publish events in a separate thread to not block
155 + def publish_events():
156 + for i in range(num_batches):
157 + batch = create_test_events(events_per_batch)
158 + publisher.publish(batch)
159 + # Small delay to avoid overwhelming
160 + if i % 100 == 0:
161 + time.sleep(0.01)
162 +
163 + received: list[tuple[int, SampleBatch]] = []
164 +
165 + publisher_thread = threading.Thread(target=publish_events)
166 + publisher_thread.start()
167 +
168 + start_time = time.time()
169 + while len(received) < num_batches:
170 + if time.time() - start_time > 10: # Timeout after 10 seconds
171 + break
172 +
173 + result = subscriber.receive_one(timeout=100)
174 + if result:
175 + received.append(result)
176 +
177 + publisher_thread.join()
178 +
179 + assert len(received) >= num_batches * 0.9, (
180 + "We should have received most messages")
181 +
182 + seqs = [seq for seq, _ in received]
183 + assert sorted(seqs) == seqs, "Sequence numbers should be in order"
184 +
185 +
186 +def test_null_publisher():
187 + """Test that NullEventPublisher can be used without errors"""
188 + publisher = NullEventPublisher()
189 +
190 + # This should not raise any errors
191 + batch = create_test_events(5)
192 + publisher.publish(batch)
193 + publisher.shutdown()
tests/v1/core/test_prefix_caching.py
+59 -3
@@ -6,6 +6,7 @@ from typing import Optional
6 import pytest
7 import torch
8
9 +from vllm.distributed.kv_events import AllBlocksCleared, BlockRemoved
10 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
11 from vllm.sampling_params import SamplingParams
12 from vllm.utils import sha256
@@ -48,9 +49,10 @@ def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
49 num_blocks=num_blocks,
50 tensors={},
51 kv_cache_groups=[
51 - KVCacheGroupSpec(['layer'],
52 - FullAttentionSpec(block_size, 1, 1, torch.float32,
53 - False))
52 + KVCacheGroupSpec(
53 + ["layer"],
54 + FullAttentionSpec(block_size, 1, 1, torch.float32, False),
55 + )
56 ],
57 )
58
@@ -783,6 +785,60 @@ def test_prefix_cache_stats_disabled():
785 assert manager.prefix_cache_stats is None
786
787
788 +@pytest.mark.parametrize("blocks_to_cache", [2, 3, 10])
789 +def test_kv_cache_events(blocks_to_cache: int):
790 + block_size = 16
791 + num_blocks = blocks_to_cache + 1
792 +
793 + # Allocate Blocks
794 + # Should see a single block stored event with a blocks_to_cache number of
795 + # block hashes
796 + # take_events should reset the kv_event_queue
797 + manager = KVCacheManager(
798 + make_kv_cache_config(block_size, num_blocks),
799 + max_model_len=8192,
800 + enable_caching=True,
801 + enable_kv_cache_events=True,
802 + )
803 +
804 + num_tokens = block_size * blocks_to_cache
805 + req0 = make_request("0", list(range(num_tokens)))
806 + _ = manager.allocate_slots(req0, num_tokens)
807 + events = manager.take_events()
808 +
809 + block = events[-1]
810 + assert (len(block.block_hashes) == blocks_to_cache == len(
811 + manager.block_pool.cached_block_hash_to_block))
812 + assert len(block.token_ids) == block.block_size * len(block.block_hashes)
813 + assert len(manager.block_pool.kv_event_queue) == 0
814 +
815 + stored_block_hash = block.block_hashes
816 +
817 + # Remove blocks and send another request
818 + # Should see block_to_cache number of removed block events and a new block
819 + # stored event
820 + manager.free(req0)
821 + req1 = make_request("1", list(range(num_tokens)))
822 + _ = manager.allocate_slots(req1, num_tokens)
823 + events = manager.take_events()
824 +
825 + for blocks in events[:-1]:
826 + assert blocks.block_hashes[0] in stored_block_hash
827 + assert len(events) == blocks_to_cache + 1
828 + assert (isinstance(events[-2], BlockRemoved))
829 + assert (len(events[-1].block_hashes) == blocks_to_cache == len(
830 + manager.block_pool.cached_block_hash_to_block))
831 +
832 + # All Blocks Cleared
833 + # Should see a single all blocks cleared event
834 + manager.free(req1)
835 + manager.reset_prefix_cache()
836 + events = manager.take_events()
837 +
838 + assert isinstance(events[-1], AllBlocksCleared)
839 + assert len(manager.block_pool.cached_block_hash_to_block) == 0
840 +
841 +
842 def test_eagle_enabled_removes_last_block():
843 """Verify Eagle does NOT remove blocks when request
844 length is divisible by block size."""
tests/v1/engine/conftest.py
+2
@@ -13,6 +13,8 @@ from tests.v1.engine.utils import (NUM_PROMPT_LOGPROBS_UNDER_TEST,
13 from vllm.engine.arg_utils import EngineArgs
14 from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
15
16 +from ...distributed.conftest import publisher_config, random_port # noqa: F401
17 +
18 from tests.v1.engine.utils import FULL_STRINGS # isort: skip
19
20 EngineCoreSampleLogprobsType = list[tuple[torch.Tensor, torch.Tensor]]
tests/v1/engine/test_engine_core_client.py
+127 -37
@@ -11,6 +11,7 @@ import pytest
11 from transformers import AutoTokenizer
12
13 from vllm import SamplingParams
14 +from vllm.distributed.kv_events import BlockStored, KVEventBatch
15 from vllm.engine.arg_utils import EngineArgs
16 from vllm.platforms import current_platform
17 from vllm.usage.usage_lib import UsageContext
@@ -20,6 +21,7 @@ from vllm.v1.engine.core_client import (AsyncMPClient, EngineCoreClient,
21 SyncMPClient)
22 from vllm.v1.executor.abstract import Executor
23
24 +from ...distributed.conftest import MockSubscriber
25 from ...utils import create_new_process_for_each_test
26
27 if not current_platform.is_cuda():
@@ -199,54 +201,142 @@ async def test_engine_core_client_asyncio(monkeypatch: pytest.MonkeyPatch):
201 log_stats=True,
202 )
203
202 - MAX_TOKENS = 20
203 - params = SamplingParams(max_tokens=MAX_TOKENS)
204 - """Normal Request Cycle."""
204 + try:
205 + MAX_TOKENS = 20
206 + params = SamplingParams(max_tokens=MAX_TOKENS)
207 + """Normal Request Cycle."""
208
206 - requests = [make_request(params) for _ in range(10)]
207 - request_ids = [req.request_id for req in requests]
209 + requests = [make_request(params) for _ in range(10)]
210 + request_ids = [req.request_id for req in requests]
211
209 - # Add requests to the engine.
210 - for request in requests:
211 - await client.add_request_async(request)
212 - await asyncio.sleep(0.01)
212 + # Add requests to the engine.
213 + for request in requests:
214 + await client.add_request_async(request)
215 + await asyncio.sleep(0.01)
216
214 - outputs: dict[str, list] = {req_id: [] for req_id in request_ids}
215 - await loop_until_done_async(client, outputs)
217 + outputs: dict[str, list] = {req_id: [] for req_id in request_ids}
218 + await loop_until_done_async(client, outputs)
219
217 - for req_id in request_ids:
218 - assert len(outputs[req_id]) == MAX_TOKENS, (
219 - f"{outputs[req_id]=}, {MAX_TOKENS=}")
220 - """Abort Request Cycle."""
220 + for req_id in request_ids:
221 + assert len(outputs[req_id]) == MAX_TOKENS, (
222 + f"{outputs[req_id]=}, {MAX_TOKENS=}")
223 + """Abort Request Cycle."""
224 +
225 + # Add requests to the engine.
226 + for idx, request in enumerate(requests):
227 + await client.add_request_async(request)
228 + await asyncio.sleep(0.01)
229 + if idx % 2 == 0:
230 + await client.abort_requests_async([request.request_id])
231 +
232 + outputs = {req_id: [] for req_id in request_ids}
233 + await loop_until_done_async(client, outputs)
234 +
235 + for idx, req_id in enumerate(request_ids):
236 + if idx % 2 == 0:
237 + assert len(outputs[req_id]) < MAX_TOKENS, (
238 + f"{len(outputs[req_id])=}, {MAX_TOKENS=}")
239 + else:
240 + assert len(outputs[req_id]) == MAX_TOKENS, (
241 + f"{len(outputs[req_id])=}, {MAX_TOKENS=}")
242 + """Utility method invocation"""
243
222 - # Add requests to the engine.
223 - for idx, request in enumerate(requests):
224 - await client.add_request_async(request)
225 - await asyncio.sleep(0.01)
226 - if idx % 2 == 0:
227 - await client.abort_requests_async([request.request_id])
244 + core_client: AsyncMPClient = client
245
229 - outputs = {req_id: [] for req_id in request_ids}
230 - await loop_until_done_async(client, outputs)
246 + result = await core_client.call_utility_async("echo", "testarg")
247 + assert result == "testarg"
248
232 - for idx, req_id in enumerate(request_ids):
233 - if idx % 2 == 0:
234 - assert len(outputs[req_id]) < MAX_TOKENS, (
235 - f"{len(outputs[req_id])=}, {MAX_TOKENS=}")
236 - else:
237 - assert len(outputs[req_id]) == MAX_TOKENS, (
238 - f"{len(outputs[req_id])=}, {MAX_TOKENS=}")
239 - """Utility method invocation"""
249 + with pytest.raises(Exception) as e_info:
250 + await core_client.call_utility_async("echo", None, "help!")
251 +
252 + assert str(e_info.value) == "Call to echo method failed: help!"
253 + finally:
254 + client.shutdown()
255 +
256 +
257 +@pytest.mark.parametrize(
258 + "multiprocessing_mode,publisher_config",
259 + [(True, "tcp"), (False, "inproc")],
260 + indirect=["publisher_config"],
261 +)
262 +def test_kv_cache_events(
263 + monkeypatch: pytest.MonkeyPatch,
264 + multiprocessing_mode: bool,
265 + publisher_config,
266 +):
267 +
268 + with monkeypatch.context() as m:
269 + m.setenv("VLLM_USE_V1", "1")
270 + block_size = 16
271 + num_blocks = 2
272
241 - core_client: AsyncMPClient = client
273 + engine_args = EngineArgs(model=MODEL_NAME,
274 + enforce_eager=True,
275 + enable_prefix_caching=True,
276 + block_size=block_size)
277 + engine_args.kv_events_config = publisher_config
278
243 - result = await core_client.call_utility_async("echo", "testarg")
244 - assert result == "testarg"
279 + vllm_config = engine_args.create_engine_config(
280 + UsageContext.UNKNOWN_CONTEXT)
281
246 - with pytest.raises(Exception) as e_info:
247 - await core_client.call_utility_async("echo", None, "help!")
282 + executor_class = Executor.get_class(vllm_config)
283 + client = EngineCoreClient.make_client(
284 + multiprocess_mode=multiprocessing_mode,
285 + asyncio_mode=False,
286 + vllm_config=vllm_config,
287 + executor_class=executor_class,
288 + log_stats=False,
289 + )
290 + endpoint = publisher_config.endpoint.replace("*", "127.0.0.1")
291 + time.sleep(0.1)
292 + subscriber = MockSubscriber(endpoint,
293 + topic=publisher_config.topic,
294 + decode_type=KVEventBatch)
295 +
296 + try:
297 + custom_tokens = list(range(num_blocks * block_size))
298 + request = EngineCoreRequest(
299 + request_id=str(uuid.uuid4()),
300 + prompt_token_ids=custom_tokens,
301 + mm_inputs=None,
302 + mm_hashes=None,
303 + mm_placeholders=None,
304 + sampling_params=SamplingParams(
305 + max_tokens=1), # Short completion for speed
306 + eos_token_id=None,
307 + arrival_time=time.time(),
308 + lora_request=None,
309 + )
310 + client.add_request(request)
311
249 - assert str(e_info.value) == "Call to echo method failed: help!"
312 + outputs: dict[str, list] = {request.request_id: []}
313 + loop_until_done(client, outputs)
314 +
315 + result = subscriber.receive_one(timeout=1000)
316 + assert result is not None, "No message received"
317 +
318 + seq, received = result
319 +
320 + assert seq == 0, "Sequence number mismatch"
321 + assert len(received.events) == 1, (
322 + "We should have exactly one BlockStored event")
323 + event = received.events[0]
324 + assert isinstance(
325 + event, BlockStored), ("We should have a BlockStored event")
326 + assert len(event.block_hashes) == num_blocks, (
327 + "We should have a BlockStored event with 2 block_hashes")
328 + assert event.block_size == block_size, (
329 + "Block size should be the same as the block size")
330 + assert event.parent_block_hash is None, (
331 + "Parent block hash should be None")
332 + assert event.lora_id is None, "Lora id should be None"
333 + assert len(event.token_ids) == num_blocks * block_size, (
334 + "Token ids should be the same as the custom tokens")
335 + assert event.token_ids == custom_tokens, (
336 + "Token ids should be the same as the custom tokens")
337 + finally:
338 + client.shutdown()
339 + return
340
341
342 @pytest.mark.timeout(10)
vllm/config.py
+60
@@ -1958,6 +1958,8 @@ class SchedulerConfig:
1958 some image tokens can be scheduled (like TTTTIIIII, leaving IIIII),
1959 it will be scheduled as TTTT in one step and IIIIIIIIII in the next."""
1960
1961 + # scheduler class or path. "vllm.core.scheduler.Scheduler" (default)
1962 + # or "mod.custom_class".
1963 scheduler_cls: Union[str, type[object]] = "vllm.core.scheduler.Scheduler"
1964 """The scheduler class to use. "vllm.core.scheduler.Scheduler" is the
1965 default scheduler. Can be a class directly or the path to a class of form
@@ -3417,6 +3419,51 @@ class KVTransferConfig(BaseModel):
3419 return self.kv_connector_extra_config.get(key, default)
3420
3421
3422 +class KVEventsConfig(BaseModel):
3423 + """Configuration for KV event publishing."""
3424 +
3425 + enable_kv_cache_events: bool = False
3426 + """If True, enable KV cache events for tracking block storage and removal.
3427 + Events can be published externally by zmq using the event publisher config.
3428 + """
3429 +
3430 + publisher: str = "null"
3431 + """The publisher to use for publishing kv events. Can be "null", "zmq".
3432 + """
3433 +
3434 + endpoint: str = "tcp://*:5557"
3435 + """The zmq endpoint to use for publishing kv events.
3436 + """
3437 +
3438 + replay_endpoint: Optional[str] = None
3439 + """The zmq endpoint to use for replaying kv events.
3440 + """
3441 +
3442 + buffer_steps: int = 10_000
3443 + """The number of steps to cache for replay endpoint. Will only save
3444 + events from the last N steps for the replay endpoint.
3445 + """
3446 +
3447 + hwm: int = 100_000
3448 + """The zmq high water mark for the event publisher. After queueing N events,
3449 + events will start dropping if the consumer is not keeping up.
3450 + """
3451 +
3452 + max_queue_size: int = 100_000
3453 + """The maximum number of events to queue while waiting for publishing.
3454 + """
3455 +
3456 + topic: str = ""
3457 + """The topic to use for the event publisher. Consumers can subscribe to
3458 + this topic to receive events.
3459 + """
3460 +
3461 + @classmethod
3462 + def from_cli(cls, cli_value: str) -> "KVEventsConfig":
3463 + """Parse the CLI value for the event publisher config."""
3464 + return KVEventsConfig.model_validate_json(cli_value)
3465 +
3466 +
3467 class CompilationLevel:
3468 # constants for the levels of the compilation process
3469 NO_COMPILATION = 0
@@ -3779,6 +3826,7 @@ class VllmConfig:
3826 init=True) # type: ignore
3827 kv_transfer_config: KVTransferConfig = field(default=None,
3828 init=True) # type: ignore
3829 + kv_events_config: Optional[KVEventsConfig] = None
3830 # some opaque config, only used to provide additional information
3831 # for the hash computation, mainly used for testing, debugging or out of
3832 # tree config registration.
@@ -4038,6 +4086,18 @@ class VllmConfig:
4086 if self.cache_config is not None:
4087 self.cache_config.enable_prefix_caching = False
4088
4089 + if (self.kv_events_config
4090 + and self.kv_events_config.enable_kv_cache_events
4091 + and not self.cache_config.enable_prefix_caching):
4092 + logger.warning(
4093 + "KV cache events are on, but prefix caching is not enabled."
4094 + "Use --enable-prefix-caching to enable.")
4095 + if (self.kv_events_config and self.kv_events_config.publisher != "null"
4096 + and not self.kv_events_config.enable_kv_cache_events):
4097 + logger.warning("KV cache events are disabled,"
4098 + "but the scheduler is configured to publish them."
4099 + "Modify KVEventsConfig.enable_kv_cache_events"
4100 + "to True to enable.")
4101 current_platform.check_and_update_config(self)
4102
4103 if not self.instance_id:
vllm/distributed/kv_events.py new
+295
@@ -0,0 +1,295 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import queue
4 +import threading
5 +import time
6 +from abc import ABC, abstractmethod
7 +from collections import deque
8 +from itertools import count
9 +from queue import Queue
10 +from typing import Any, Callable, Optional, Union
11 +
12 +import msgspec
13 +import zmq
14 +
15 +from vllm.config import KVEventsConfig
16 +from vllm.logger import init_logger
17 +
18 +logger = init_logger(__name__)
19 +
20 +
21 +class EventBatch(
22 + msgspec.Struct,
23 + array_like=True, # type: ignore[call-arg]
24 + omit_defaults=True, # type: ignore[call-arg]
25 + gc=False, # type: ignore[call-arg]
26 +):
27 + ts: float
28 + events: list[Any]
29 +
30 +
31 +class KVCacheEvent(
32 + msgspec.Struct,
33 + array_like=True, # type: ignore[call-arg]
34 + omit_defaults=True, # type: ignore[call-arg]
35 + gc=False, # type: ignore[call-arg]
36 + tag=True):
37 + """Base class for all KV cache-related events"""
38 +
39 +
40 +class BlockStored(KVCacheEvent):
41 + block_hashes: list[int]
42 + parent_block_hash: Optional[int]
43 + token_ids: list[int]
44 + block_size: int
45 + lora_id: Optional[int]
46 +
47 +
48 +class BlockRemoved(KVCacheEvent):
49 + block_hashes: list[int]
50 +
51 +
52 +class AllBlocksCleared(KVCacheEvent):
53 + pass
54 +
55 +
56 +class KVEventBatch(EventBatch):
57 + events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]]
58 +
59 +
60 +class EventPublisher(ABC):
61 + """Lightweight publisher for EventBatch batches."""
62 +
63 + @abstractmethod
64 + def publish(self, events: EventBatch) -> None:
65 + """Emit events in order.
66 +
67 + Implementations should guarantee at-least-once delivery and
68 + monotonic ordering (e.g., via sequence numbers).
69 + """
70 +
71 + @abstractmethod
72 + def shutdown(self) -> None:
73 + """Shutdown the publisher."""
74 +
75 +
76 +class NullEventPublisher(EventPublisher):
77 + """No-op implementation (default when disabled)."""
78 +
79 + def publish(self, events) -> None:
80 + return
81 +
82 + def shutdown(self) -> None:
83 + return
84 +
85 +
86 +class ZmqEventPublisher(EventPublisher):
87 + """Reliable PUB/ROUTER publisher with an in-memory replay buffer.
88 +
89 + Spawns a separate thread to handle publishing from a queue.
90 +
91 + Parameters
92 + ----------
93 + endpoint:
94 + PUB address. Use ``tcp://*:5557`` to bind or ``tcp://host:5557`` to
95 + connect.
96 + replay_endpoint:
97 + Optional ROUTER address for replay requests. When given, subscribers can
98 + request missed batches by sending the starting sequence number as an
99 + 8-byte big-endian integer.
100 + buffer_steps:
101 + Number of past batches to keep for replay.
102 + hwm:
103 + ZeroMQ high-water-mark for PUB socket.
104 + max_queue_size:
105 + Maximum number of events to buffer in memory.
106 + topic:
107 + Topic to publish events to.
108 + """
109 + SHUTDOWN_TIMEOUT: float = 1.0
110 + END_SEQ = (-1).to_bytes(8, "big", signed=True)
111 +
112 + def __init__(
113 + self,
114 + endpoint: str = "tcp://*:5557",
115 + replay_endpoint: Optional[str] = None,
116 + buffer_steps: int = 10_000,
117 + hwm: int = 100_000,
118 + max_queue_size: int = 100_000,
119 + topic: str = "",
120 + ) -> None:
121 + # Storage
122 + self._event_queue = Queue[Optional[EventBatch]](maxsize=max_queue_size)
123 + self._buffer = deque[tuple[int, bytes]](maxlen=buffer_steps)
124 +
125 + # ZMQ sockets
126 + self._ctx = zmq.Context.instance()
127 + self._pub: Optional[zmq.Socket] = None
128 + self._replay: Optional[zmq.Socket] = None
129 + self._endpoint = endpoint
130 + self._replay_endpoint = replay_endpoint
131 + self._hwm = hwm
132 +
133 + # Payload
134 + self._seq_gen = count()
135 + self._topic_bytes = topic.encode('utf-8')
136 +
137 + # Thread
138 + self._running = True
139 + logger.info("Starting ZMQ publisher thread")
140 +
141 + self._thread = threading.Thread(target=self._publisher_thread,
142 + daemon=True,
143 + name="zmq-publisher")
144 + self._thread.start()
145 +
146 + def publish(self, events: EventBatch) -> None:
147 + if not self._running:
148 + raise RuntimeError("Publisher is closed")
149 + self._event_queue.put(events)
150 +
151 + def shutdown(self) -> None:
152 + """Stop the publisher thread and clean up resources."""
153 + self._running = False
154 + self._event_queue.put_nowait(None)
155 +
156 + start = time.time()
157 + pending_items = True
158 + while pending_items and (time.time() - start < self.SHUTDOWN_TIMEOUT):
159 + pending_items = not self._event_queue.empty()
160 + if pending_items:
161 + time.sleep(0.1)
162 +
163 + if pending_items:
164 + logger.warning(
165 + "Warning: Queue still has %s items after %s seconds timeout",
166 + self._event_queue.qsize(),
167 + self.SHUTDOWN_TIMEOUT,
168 + )
169 +
170 + if self._thread.is_alive():
171 + self._thread.join(timeout=self.SHUTDOWN_TIMEOUT)
172 +
173 + # Clean up ZMQ resources
174 + try:
175 + if self._pub is not None:
176 + self._pub.close(linger=0)
177 + if self._replay is not None:
178 + self._replay.close(linger=0)
179 + finally:
180 + pass # Do not terminate context; other sockets may use it
181 +
182 + def _socket_setup(self) -> None:
183 + """Initialize sockets
184 + https://pyzmq.readthedocs.io/en/v19.0.0/morethanbindings.html#thread-safety
185 + """
186 + if self._pub is None:
187 + self._pub = self._ctx.socket(zmq.PUB)
188 + self._pub.set_hwm(self._hwm)
189 + # Heuristic: bind if wildcard / * present, else connect.
190 + # bind stable, connect volatile convention
191 + if ("*" in self._endpoint or "::" in self._endpoint
192 + or self._endpoint.startswith("ipc://")
193 + or self._endpoint.startswith("inproc://")):
194 + self._pub.bind(self._endpoint)
195 + else:
196 + self._pub.connect(self._endpoint)
197 +
198 + # Set up replay socket: use ROUTER
199 + # 1) handles multiple REQ clients (identities)
200 + # 2) lets us send back one request → many replies (streamed events)
201 + # 3) works in our non‑blocking poll loop alongside PUB
202 + if self._replay_endpoint is not None:
203 + self._replay = self._ctx.socket(zmq.ROUTER)
204 + self._replay.bind(self._replay_endpoint)
205 +
206 + def _publisher_thread(self) -> None:
207 + """Background thread that processes the event queue."""
208 + self._pack = msgspec.msgpack.Encoder()
209 + self._socket_setup()
210 +
211 + assert self._pub is not None # narrows type for mypy
212 +
213 + while self._running or self._event_queue.qsize() > 0:
214 + # --- replay (non-critical) ---------------------------------
215 + if self._replay is not None and self._replay.poll(0):
216 + try:
217 + self._service_replay()
218 + except Exception as e:
219 + logger.exception("Error in replay: %s", e)
220 +
221 + # --- main queue (critical) ---------------------------------
222 + try:
223 + event = self._event_queue.get(timeout=0.1)
224 + if event is None:
225 + break # Sentinel received, exit thread
226 + except queue.Empty:
227 + continue
228 +
229 + try:
230 + seq = next(self._seq_gen)
231 +
232 + payload = self._pack.encode(event)
233 + seq_bytes = seq.to_bytes(8, "big")
234 + self._pub.send_multipart(
235 + (self._topic_bytes, seq_bytes, payload))
236 +
237 + self._buffer.append((seq, payload))
238 + self._event_queue.task_done()
239 +
240 + except Exception as e:
241 + # Publishing failed; back-off a bit to avoid a tight error loop
242 + logger.exception("Error in publisher thread: %s", e)
243 + time.sleep(0.1)
244 +
245 + def _service_replay(self) -> None:
246 + """If a replay request is waiting, send buffered batches."""
247 + assert self._replay is not None # narrows type for mypy
248 +
249 + frame = self._replay.recv_multipart()
250 + if len(frame) != 3:
251 + logger.warning("Invalid replay request: %s", frame)
252 + return
253 + client_id, _, start_seq_bytes = frame
254 + start_seq = int.from_bytes(start_seq_bytes, "big")
255 +
256 + for seq, buf in self._buffer:
257 + if seq >= start_seq:
258 + # [identity, empty_delim, seq_bytes, payload]
259 + # (identity, empty_delim) are stripped off by the router
260 + # receiving payload is (seq_bytes, payload)
261 + self._replay.send_multipart(
262 + (client_id, b"", seq.to_bytes(8, "big"), buf))
263 + # Send end of sequence marker
264 + # receiving payload is (-1, b""")
265 + self._replay.send_multipart((client_id, b"", self.END_SEQ, b""))
266 +
267 +
268 +class EventPublisherFactory:
269 + _registry: dict[str, Callable[..., EventPublisher]] = {
270 + "null": NullEventPublisher,
271 + "zmq": ZmqEventPublisher,
272 + }
273 +
274 + @classmethod
275 + def register_publisher(cls, name: str,
276 + ctor: Callable[..., EventPublisher]) -> None:
277 + if name in cls._registry:
278 + raise KeyError(f"publisher '{name}' already registered")
279 + cls._registry[name] = ctor
280 +
281 + @classmethod
282 + def create(cls, config: Optional[KVEventsConfig]) -> EventPublisher:
283 + """Create publisher from a config mapping."""
284 + if not config:
285 + return NullEventPublisher()
286 +
287 + config_dict = config.model_dump()
288 +
289 + kind = config_dict.pop("publisher", "null")
290 + config_dict.pop("enable_kv_cache_events")
291 + try:
292 + constructor = cls._registry[kind]
293 + except KeyError as exc:
294 + raise ValueError(f"Unknown event publisher '{kind}'") from exc
295 + return constructor(**config_dict)
vllm/engine/arg_utils.py
+14 -8
@@ -19,14 +19,14 @@ from vllm.config import (BlockSize, CacheConfig, CacheDType, CompilationConfig,
19 ConfigFormat, ConfigType, DecodingConfig, Device,
20 DeviceConfig, DistributedExecutorBackend,
21 GuidedDecodingBackend, GuidedDecodingBackendV1,
22 - HfOverrides, KVTransferConfig, LoadConfig, LoadFormat,
23 - LoRAConfig, ModelConfig, ModelDType, ModelImpl,
24 - MultiModalConfig, ObservabilityConfig, ParallelConfig,
25 - PoolerConfig, PrefixCachingHashAlgo,
26 - PromptAdapterConfig, SchedulerConfig, SchedulerPolicy,
27 - SpeculativeConfig, TaskOption, TokenizerMode,
28 - TokenizerPoolConfig, VllmConfig, get_attr_docs,
29 - get_field)
22 + HfOverrides, KVEventsConfig, KVTransferConfig,
23 + LoadConfig, LoadFormat, LoRAConfig, ModelConfig,
24 + ModelDType, ModelImpl, MultiModalConfig,
25 + ObservabilityConfig, ParallelConfig, PoolerConfig,
26 + PrefixCachingHashAlgo, PromptAdapterConfig,
27 + SchedulerConfig, SchedulerPolicy, SpeculativeConfig,
28 + TaskOption, TokenizerMode, TokenizerPoolConfig,
29 + VllmConfig, get_attr_docs, get_field)
30 from vllm.executor.executor_base import ExecutorBase
31 from vllm.logger import init_logger
32 from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -353,6 +353,7 @@ class EngineArgs:
353 worker_extension_cls: str = ParallelConfig.worker_extension_cls
354
355 kv_transfer_config: Optional[KVTransferConfig] = None
356 + kv_events_config: Optional[KVEventsConfig] = None
357
358 generation_config: str = ModelConfig.generation_config
359 enable_sleep_mode: bool = ModelConfig.enable_sleep_mode
@@ -769,6 +770,10 @@ class EngineArgs:
770 default=None,
771 help='The configurations for distributed KV cache '
772 'transfer. Should be a JSON string.')
773 + parser.add_argument('--kv-events-config',
774 + type=KVEventsConfig.from_cli,
775 + default=None,
776 + help='The configurations for event publishing.')
777
778 parser.add_argument(
779 '--worker-cls',
@@ -1125,6 +1130,7 @@ class EngineArgs:
1130 prompt_adapter_config=prompt_adapter_config,
1131 compilation_config=self.compilation_config,
1132 kv_transfer_config=self.kv_transfer_config,
1133 + kv_events_config=self.kv_events_config,
1134 additional_config=self.additional_config,
1135 )
1136
vllm/v1/core/block_pool.py
+48 -1
@@ -3,6 +3,8 @@ from collections import defaultdict
3 from collections.abc import Iterable
4 from typing import Callable, Optional
5
6 +from vllm.distributed.kv_events import (AllBlocksCleared, BlockRemoved,
7 + BlockStored, KVCacheEvent)
8 from vllm.logger import init_logger
9 from vllm.v1.core.kv_cache_utils import (BlockHashType, FreeKVCacheBlockQueue,
10 KVCacheBlock,
@@ -26,7 +28,12 @@ class BlockPool:
28 enable_caching: Whether to enable prefix caching.
29 """
30
29 - def __init__(self, num_gpu_blocks: int, enable_caching: bool):
31 + def __init__(
32 + self,
33 + num_gpu_blocks: int,
34 + enable_caching: bool,
35 + enable_kv_cache_events: bool = False,
36 + ):
37 assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0
38 self.num_gpu_blocks = num_gpu_blocks
39 self.enable_caching = enable_caching
@@ -56,6 +63,9 @@ class BlockPool:
63 # avoid freeing it.
64 self.null_block = self.free_block_queue.popleft()
65
66 + self.enable_kv_cache_events = enable_kv_cache_events
67 + self.kv_event_queue: list[KVCacheEvent] = []
68 +
69 def get_cached_block(self,
70 block_hash: BlockHashType) -> Optional[KVCacheBlock]:
71 """Get a cached block by the block hash, or None if cache miss.
@@ -116,6 +126,9 @@ class BlockPool:
126 assert prev_block.block_hash is not None
127 prev_block_hash_value = prev_block.block_hash.hash_value
128
129 + parent_block_hash = prev_block_hash_value
130 + new_hashes: Optional[list[int]] = ([] if self.enable_kv_cache_events
131 + else None)
132 for i, blk in enumerate(new_full_blocks):
133 assert blk.block_hash is None
134
@@ -153,8 +166,23 @@ class BlockPool:
166 # Update and added the full block to the cache.
167 blk.block_hash = block_hash
168 self.cached_block_hash_to_block[block_hash][blk.block_id] = blk
169 + if new_hashes is not None:
170 + new_hashes.append(block_hash.hash_value)
171 prev_block_hash_value = block_hash.hash_value
172
173 + if self.enable_kv_cache_events:
174 + self.kv_event_queue.append(
175 + BlockStored(
176 + block_hashes=new_hashes,
177 + parent_block_hash=parent_block_hash,
178 + token_ids=request.
179 + all_token_ids[num_cached_blocks *
180 + block_size:num_full_blocks * block_size],
181 + block_size=block_size,
182 + lora_id=request.lora_request.id
183 + if request.lora_request else None,
184 + ))
185 +
186 def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
187 """Get new blocks from the free block pool.
188
@@ -206,6 +234,9 @@ class BlockPool:
234 if len(self.cached_block_hash_to_block[block_hash]) == 0:
235 del self.cached_block_hash_to_block[block_hash]
236
237 + if self.enable_kv_cache_events:
238 + self.kv_event_queue.append(
239 + BlockRemoved(block_hashes=[block_hash.hash_value]))
240 return True
241 return False
242
@@ -262,6 +293,10 @@ class BlockPool:
293 block.reset_hash()
294
295 logger.info("Successfully reset prefix cache")
296 +
297 + if self.enable_kv_cache_events:
298 + self.kv_event_queue.append(AllBlocksCleared())
299 +
300 return True
301
302 def get_num_free_blocks(self) -> int:
@@ -279,3 +314,15 @@ class BlockPool:
314 The KV cache usage (between 0.0 and 1.0).
315 """
316 return 1.0 - (self.get_num_free_blocks() / self.num_gpu_blocks)
317 +
318 + def take_events(self) -> list[KVCacheEvent]:
319 + """Atomically takes all events and clears the queue.
320 +
321 + Returns:
322 + A list of KV cache events.
323 + """
324 + if not self.enable_kv_cache_events:
325 + return []
326 + events = self.kv_event_queue
327 + self.kv_event_queue = []
328 + return events
vllm/v1/core/kv_cache_manager.py
+13 -1
@@ -4,6 +4,7 @@ from collections import defaultdict
4 from collections.abc import Iterable
5 from typing import Optional
6
7 +from vllm.distributed.kv_events import KVCacheEvent
8 from vllm.logger import init_logger
9 from vllm.utils import cdiv, sha256
10 from vllm.v1.core.block_pool import BlockPool
@@ -27,6 +28,7 @@ class KVCacheManager:
28 caching_hash_algo: str = "builtin",
29 use_eagle: bool = False,
30 log_stats: bool = False,
31 + enable_kv_cache_events: bool = False,
32 ) -> None:
33 assert len(kv_cache_config.kv_cache_groups) == 1, (
34 "KVCacheManager does not support hybrid models with more than 1 "
@@ -44,7 +46,9 @@ class KVCacheManager:
46 # FIXME: make prefix cache stats conditional on log_stats
47 self.prefix_cache_stats = PrefixCacheStats() if log_stats else None
48
47 - self.block_pool = BlockPool(self.num_gpu_blocks, enable_caching)
49 + self.block_pool = BlockPool(self.num_gpu_blocks, enable_caching,
50 + enable_kv_cache_events)
51 +
52 self.specialized_manager = get_specialized_manager(
53 kv_cache_spec=kv_cache_spec,
54 block_pool=self.block_pool,
@@ -383,3 +387,11 @@ class KVCacheManager:
387 is finished, not when it is preempted.
388 """
389 self.req_to_block_hashes.pop(request.request_id, None)
390 +
391 + def take_events(self) -> list[KVCacheEvent]:
392 + """Take the KV cache events from the block pool.
393 +
394 + Returns:
395 + A list of KV cache events.
396 + """
397 + return self.block_pool.take_events()
vllm/v1/core/sched/interface.py
+5
@@ -132,3 +132,8 @@ class SchedulerInterface(ABC):
132 The SchedulerStats object is created for every scheduling step.
133 """
134 raise NotImplementedError
135 +
136 + @abstractmethod
137 + def shutdown(self) -> None:
138 + """Shutdown the scheduler."""
139 + raise NotImplementedError
vllm/v1/core/sched/scheduler.py
+20 -1
@@ -8,6 +8,7 @@ from collections.abc import Iterable
8 from typing import Optional, Union
9
10 from vllm.config import VllmConfig
11 +from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch
12 from vllm.distributed.kv_transfer.kv_connector.factory import (
13 KVConnectorFactory)
14 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
@@ -48,6 +49,7 @@ class Scheduler(SchedulerInterface):
49 self.cache_config = vllm_config.cache_config
50 self.lora_config = vllm_config.lora_config
51 self.kv_cache_config = kv_cache_config
52 + self.kv_events_config = vllm_config.kv_events_config
53 self.log_stats = log_stats
54 self.structured_output_manager = structured_output_manager
55
@@ -62,6 +64,9 @@ class Scheduler(SchedulerInterface):
64 self.max_num_scheduled_tokens = \
65 self.scheduler_config.max_num_batched_tokens
66 self.max_model_len = self.scheduler_config.max_model_len
67 + self.enable_kv_cache_events = (
68 + self.kv_events_config is not None
69 + and self.kv_events_config.enable_kv_cache_events)
70
71 # Create KVConnector for the Scheduler. Note that each Worker
72 # will have a corresponding KVConnector with Role=WORKER.
@@ -71,6 +76,9 @@ class Scheduler(SchedulerInterface):
76 self.connector = KVConnectorFactory.create_connector_v1(
77 config=self.vllm_config, role=KVConnectorRole.SCHEDULER)
78
79 + self.kv_event_publisher = EventPublisherFactory.create(
80 + self.kv_events_config)
81 +
82 num_gpu_blocks = self.cache_config.num_gpu_blocks
83 assert num_gpu_blocks is not None and num_gpu_blocks > 0
84
@@ -132,7 +140,9 @@ class Scheduler(SchedulerInterface):
140 enable_caching=self.cache_config.enable_prefix_caching,
141 caching_hash_algo=self.cache_config.prefix_caching_hash_algo,
142 use_eagle=self.use_eagle,
135 - log_stats=self.log_stats)
143 + log_stats=self.log_stats,
144 + enable_kv_cache_events=self.enable_kv_cache_events,
145 + )
146
147 def schedule(self) -> SchedulerOutput:
148 # NOTE(woosuk) on the scheduling algorithm:
@@ -493,6 +503,11 @@ class Scheduler(SchedulerInterface):
503 meta = self.connector.build_connector_meta(scheduler_output)
504 scheduler_output.kv_connector_metadata = meta
505
506 + events = self.kv_cache_manager.take_events()
507 + if events:
508 + batch = KVEventBatch(ts=time.time(), events=events)
509 + self.kv_event_publisher.publish(batch)
510 +
511 # Advance the number of computed tokens for the request AFTER
512 # the request is scheduled.
513 # 1. The scheduler_output of the current step has to include the
@@ -843,3 +858,7 @@ class Scheduler(SchedulerInterface):
858 num_draft_tokens=num_draft_tokens,
859 num_accepted_tokens=num_accepted_tokens)
860 return spec_decoding_stats
861 +
862 + def shutdown(self) -> None:
863 + if self.kv_event_publisher:
864 + self.kv_event_publisher.shutdown()
vllm/v1/engine/core.py
+2
@@ -259,6 +259,8 @@ class EngineCore:
259 self.structured_output_manager.clear_backend()
260 if self.model_executor:
261 self.model_executor.shutdown()
262 + if self.scheduler:
263 + self.scheduler.shutdown()
264
265 def profile(self, is_start: bool = True):
266 self.model_executor.profile(is_start)