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 6, 12:00 UTC
[kv_offload+HMA][6/N]: Split offloading_connector.py (#37405)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Mar 18, 2026 at 15:42 UTC
525f2eeb0b6ea86c7b618466ef1830e2d7bd77f1
7 files changed
+733
-658
tests/v1/kv_connector/unit/test_offloading_connector.py
+6
-2
@@ -13,11 +13,15 @@ from vllm import SamplingParams
13
from vllm.config import KVTransferConfig, VllmConfig
14
from vllm.distributed.kv_events import BlockRemoved, BlockStored
15
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
16
-from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
17
- OffloadingConnector,
16
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
17
OffloadingConnectorMetadata,
18
+)
19
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
20
OffloadingConnectorStats,
21
)
22
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
23
+ OffloadingConnector,
24
+)
25
from vllm.forward_context import ForwardContext
26
from vllm.utils.hashing import sha256
27
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
vllm/distributed/kv_transfer/kv_connector/v1/offloading/__init__.py
vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py
new
+15
@@ -0,0 +1,15 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+
5
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
6
+from vllm.v1.kv_offload.worker.worker import TransferSpec
7
+
8
+ReqId = str
9
+
10
+
11
+@dataclass
12
+class OffloadingConnectorMetadata(KVConnectorMetadata):
13
+ reqs_to_load: dict[ReqId, TransferSpec]
14
+ reqs_to_store: dict[ReqId, TransferSpec]
15
+ reqs_to_flush: set[str] | None = None
vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py
new
+165
@@ -0,0 +1,165 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+from typing import Any
5
+
6
+from vllm.config import VllmConfig
7
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
8
+ KVConnectorPromMetrics,
9
+ KVConnectorStats,
10
+ PromMetric,
11
+ PromMetricT,
12
+)
13
+from vllm.logger import init_logger
14
+from vllm.v1.kv_offload.worker.worker import TransferType
15
+
16
+logger = init_logger(__name__)
17
+
18
+
19
+@dataclass
20
+class OffloadingOperationMetrics:
21
+ op_size: int
22
+ op_time: float
23
+
24
+
25
+@dataclass
26
+class OffloadingConnectorStats(KVConnectorStats):
27
+ def __post_init__(self):
28
+ if not self.data:
29
+ # Empty container init, no data is passed in.
30
+ self.reset()
31
+
32
+ def reset(self):
33
+ self.data: dict[str, list[OffloadingOperationMetrics]] = {}
34
+
35
+ def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
36
+ if not other.is_empty():
37
+ for k, v in other.data.items():
38
+ if k not in self.data:
39
+ self.data[k] = v
40
+ else:
41
+ accumulator = self.data[k]
42
+ assert isinstance(accumulator, list)
43
+ accumulator.extend(v)
44
+ return self
45
+
46
+ def reduce(self) -> dict[str, int | float]:
47
+ """
48
+ Reduce the observations collected during a time interval to one or
49
+ more representative values (eg avg/median/sum of the series).
50
+ This is meant to be called by the logger to produce a summary of the
51
+ stats for the last time interval.
52
+ """
53
+ return_dict: dict[str, int | float] = {}
54
+ for transfer_type, ops_list in self.data.items():
55
+ assert isinstance(ops_list, list)
56
+ total_bytes = 0
57
+ total_time = 0.0
58
+ for op in ops_list:
59
+ assert isinstance(op, dict)
60
+ total_bytes += op["op_size"]
61
+ total_time += op["op_time"]
62
+ return_dict[f"{transfer_type}_total_bytes"] = total_bytes
63
+ return_dict[f"{transfer_type}_total_time"] = total_time
64
+ return return_dict
65
+
66
+ def is_empty(self) -> bool:
67
+ return not self.data
68
+
69
+ def record_transfer(self, num_bytes: int, time: float, transfer_type: TransferType):
70
+ src, dst = transfer_type
71
+ transfer_type_key = src + "_to_" + dst
72
+ op = OffloadingOperationMetrics(num_bytes, time)
73
+ if transfer_type_key in self.data:
74
+ self.data[transfer_type_key].append(op)
75
+ else:
76
+ self.data[transfer_type_key] = [op]
77
+
78
+
79
+class OffloadPromMetrics(KVConnectorPromMetrics):
80
+ def __init__(
81
+ self,
82
+ vllm_config: VllmConfig,
83
+ metric_types: dict[type[PromMetric], type[PromMetricT]],
84
+ labelnames: list[str],
85
+ per_engine_labelvalues: dict[int, list[object]],
86
+ ):
87
+ super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
88
+ # (engine_idx, transfer_type) -> (metric with bounded labels)
89
+ self.histogram_transfer_size: dict[tuple[int, str], PromMetricT] = {}
90
+ self.counter_kv_bytes: dict[tuple[int, str], PromMetricT] = {}
91
+ self.counter_kv_transfer_time: dict[tuple[int, str], PromMetricT] = {}
92
+ buckets = [ # In bytes
93
+ 1e6,
94
+ 5e6,
95
+ 10e6,
96
+ 20e6,
97
+ 40e6,
98
+ 60e6,
99
+ 80e6,
100
+ 100e6,
101
+ 150e6,
102
+ 200e6,
103
+ ]
104
+
105
+ self._counter_kv_bytes = self._counter_cls(
106
+ name="vllm:kv_offload_total_bytes",
107
+ documentation="Number of bytes offloaded by KV connector",
108
+ labelnames=labelnames + ["transfer_type"],
109
+ )
110
+
111
+ self._counter_kv_transfer_time = self._counter_cls(
112
+ name="vllm:kv_offload_total_time",
113
+ documentation="Total time measured by all KV offloading operations",
114
+ labelnames=labelnames + ["transfer_type"],
115
+ )
116
+
117
+ self._histogram_transfer_size = self._histogram_cls(
118
+ name="vllm:kv_offload_size",
119
+ documentation="Histogram of KV offload transfer size, in bytes.",
120
+ buckets=buckets[:],
121
+ labelnames=labelnames + ["transfer_type"],
122
+ )
123
+
124
+ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
125
+ """
126
+ Observe transfer statistics from the new data structure.
127
+ transfer_stats_data is expected to be a dict where:
128
+ - keys are transfer type strings (e.g., "cpu_to_gpu", "gpu_to_cpu")
129
+ - values are lists of OffloadingOperationMetrics objects
130
+ """
131
+
132
+ for transfer_type, ops in transfer_stats_data.items():
133
+ # Cache:
134
+ if (engine_idx, transfer_type) not in self.histogram_transfer_size:
135
+ self.histogram_transfer_size[(engine_idx, transfer_type)] = (
136
+ self._histogram_transfer_size.labels(
137
+ *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
138
+ )
139
+ )
140
+ self.counter_kv_bytes[(engine_idx, transfer_type)] = (
141
+ self._counter_kv_bytes.labels(
142
+ *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
143
+ )
144
+ )
145
+ self.counter_kv_transfer_time[(engine_idx, transfer_type)] = (
146
+ self._counter_kv_transfer_time.labels(
147
+ *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
148
+ )
149
+ )
150
+
151
+ # Process ops:
152
+ assert isinstance(ops, list)
153
+ for op in ops: # ops is a list of serialized OffloadingOperationMetrics
154
+ assert isinstance(op, dict)
155
+ # Observe size histogram
156
+ self.histogram_transfer_size[(engine_idx, transfer_type)].observe(
157
+ op["op_size"]
158
+ )
159
+
160
+ # Increment byte and time counters
161
+ self.counter_kv_bytes[(engine_idx, transfer_type)].inc(op["op_size"])
162
+
163
+ self.counter_kv_transfer_time[(engine_idx, transfer_type)].inc(
164
+ op["op_time"]
165
+ )
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
new
+347
@@ -0,0 +1,347 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections import defaultdict
4
+from collections.abc import Iterable
5
+from itertools import islice
6
+from typing import Any
7
+
8
+from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
9
+from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
10
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
11
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
12
+ OffloadingConnectorMetadata,
13
+ ReqId,
14
+)
15
+from vllm.logger import init_logger
16
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
17
+from vllm.v1.core.kv_cache_utils import BlockHash
18
+from vllm.v1.core.sched.output import SchedulerOutput
19
+from vllm.v1.kv_offload.abstract import OffloadingManager
20
+from vllm.v1.kv_offload.mediums import GPULoadStoreSpec
21
+from vllm.v1.kv_offload.spec import OffloadingSpec
22
+from vllm.v1.kv_offload.worker.worker import TransferSpec
23
+from vllm.v1.outputs import KVConnectorOutput
24
+from vllm.v1.request import Request
25
+
26
+logger = init_logger(__name__)
27
+
28
+
29
+class OffloadingConnectorScheduler:
30
+ """Implementation of Scheduler side methods"""
31
+
32
+ def __init__(self, spec: OffloadingSpec):
33
+ assert len(spec.gpu_block_size) == 1
34
+ self.gpu_block_size = spec.gpu_block_size[0]
35
+ self.offloaded_block_size = self.gpu_block_size * spec.block_size_factor
36
+ self.block_size_factor = spec.block_size_factor
37
+ self.manager: OffloadingManager = spec.get_manager()
38
+
39
+ self._requests: dict[ReqId, Request] = {}
40
+ # list of GPU block IDs per request
41
+ self._request_block_ids: dict[ReqId, list[int]] = {}
42
+ # requests to load for the current scheduler step
43
+ self._reqs_to_load: dict[ReqId, TransferSpec] = {}
44
+ # request blocks are stored in order
45
+ # index of next block (of size offloaded_block_size) to offload
46
+ self._next_stored_block_idx: dict[ReqId, int] = {}
47
+ # if GPU prefix caching is enabled,
48
+ # track loaded blocks to avoid redundant loads
49
+ self._blocks_being_loaded: set[BlockHash] | None = (
50
+ set() if spec.vllm_config.cache_config.enable_prefix_caching else None
51
+ )
52
+
53
+ # request ID -> set(block hashes being stored/load)
54
+ self._reqs_being_stored = defaultdict[ReqId, set[BlockHash]](set)
55
+ self._reqs_being_loaded = defaultdict[ReqId, set[BlockHash]](set)
56
+
57
+ def _get_block_hashes(
58
+ self,
59
+ req: Request,
60
+ start_idx: int = 0,
61
+ end_idx: int | None = None,
62
+ ) -> Iterable[BlockHash]:
63
+ return islice(
64
+ req.block_hashes,
65
+ self.block_size_factor * start_idx + self.block_size_factor - 1,
66
+ self.block_size_factor * end_idx if end_idx else None,
67
+ self.block_size_factor,
68
+ )
69
+
70
+ def get_num_new_matched_tokens(
71
+ self, request: Request, num_computed_tokens: int
72
+ ) -> tuple[int | None, bool]:
73
+ """
74
+ Get number of new tokens that can be loaded beyond the
75
+ num_computed_tokens.
76
+
77
+ Args:
78
+ request (Request): the request object.
79
+ num_computed_tokens (int): the number of locally
80
+ computed tokens for this request
81
+
82
+ Returns:
83
+ A tuple with the following elements:
84
+ - The number of tokens that can be loaded beyond what is
85
+ already computed.
86
+ If None, it means that the connector needs more time to
87
+ determine the number of matched tokens, and the scheduler
88
+ should query for this request again later.
89
+ - `True` if tokens will be loaded asynchronously
90
+ (between scheduler steps).
91
+ """
92
+ num_blocks = request.num_tokens // self.offloaded_block_size
93
+
94
+ assert len(request.block_hashes) // self.block_size_factor == num_blocks
95
+ block_hashes = self._get_block_hashes(request)
96
+
97
+ self.manager.touch(block_hashes)
98
+
99
+ full_block_tokens = self.offloaded_block_size * num_blocks
100
+ if full_block_tokens - num_computed_tokens < self.offloaded_block_size:
101
+ # we can load less than a block, skip
102
+ return 0, False
103
+
104
+ start_block_idx = num_computed_tokens // self.offloaded_block_size
105
+ hits = self.manager.lookup(
106
+ self._get_block_hashes(request, start_idx=start_block_idx)
107
+ )
108
+ if hits is None:
109
+ # indicates a lookup that should be tried later
110
+ return None, False
111
+ if hits == 0:
112
+ return 0, False
113
+
114
+ num_hit_tokens = (
115
+ self.offloaded_block_size * (start_block_idx + hits) - num_computed_tokens
116
+ )
117
+ logger.debug(
118
+ "Request %s hit %s offloaded tokens after %s GPU hit tokens",
119
+ request.request_id,
120
+ num_hit_tokens,
121
+ num_computed_tokens,
122
+ )
123
+ if num_hit_tokens < self.offloaded_block_size:
124
+ return 0, False
125
+
126
+ if self._blocks_being_loaded:
127
+ block_hashes = self._get_block_hashes(
128
+ request, start_idx=start_block_idx, end_idx=start_block_idx + hits
129
+ )
130
+
131
+ if any(
132
+ block_hash in self._blocks_being_loaded for block_hash in block_hashes
133
+ ):
134
+ # hit blocks are being loaded, delay request
135
+ logger.debug(
136
+ "Delaying request %s since some of its blocks are already"
137
+ " being loaded",
138
+ request.request_id,
139
+ )
140
+ return None, False
141
+
142
+ return num_hit_tokens, True
143
+
144
+ def update_state_after_alloc(
145
+ self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int
146
+ ):
147
+ self._requests[request.request_id] = request
148
+ # the block ids are updated in _get_reqs_to_store
149
+ self._request_block_ids[request.request_id] = []
150
+
151
+ if num_external_tokens == 0:
152
+ return
153
+
154
+ block_groups = blocks.get_block_ids()
155
+ block_ids = block_groups[0]
156
+
157
+ num_computed_gpu_blocks = sum(
158
+ block.block_hash is not None for block in blocks.blocks[0]
159
+ )
160
+ num_computed_tokens = num_computed_gpu_blocks * self.gpu_block_size
161
+ full_block_tokens = num_computed_tokens + num_external_tokens
162
+ assert full_block_tokens % self.offloaded_block_size == 0
163
+
164
+ num_pending_gpu_blocks = len(block_ids) - num_computed_gpu_blocks
165
+ assert num_external_tokens == num_pending_gpu_blocks * self.gpu_block_size
166
+
167
+ start_block_idx = num_computed_tokens // self.offloaded_block_size
168
+ num_blocks = full_block_tokens // self.offloaded_block_size
169
+
170
+ assert len(request.block_hashes) // self.block_size_factor >= num_blocks
171
+ block_hashes = self._get_block_hashes(
172
+ request, start_idx=start_block_idx, end_idx=num_blocks
173
+ )
174
+
175
+ src_spec = self.manager.prepare_load(block_hashes)
176
+ dst_spec = GPULoadStoreSpec(block_ids[num_computed_gpu_blocks:])
177
+
178
+ block_hashes = self._get_block_hashes(
179
+ request, start_idx=start_block_idx, end_idx=num_blocks
180
+ )
181
+
182
+ self._reqs_to_load[request.request_id] = (src_spec, dst_spec)
183
+ req_blocks_being_loaded = self._reqs_being_loaded[request.request_id]
184
+ req_blocks_being_loaded.update(block_hashes)
185
+ self._next_stored_block_idx[request.request_id] = num_blocks
186
+
187
+ if self._blocks_being_loaded is not None:
188
+ self._blocks_being_loaded.update(req_blocks_being_loaded)
189
+
190
+ def _get_reqs_to_store(self, scheduler_output: SchedulerOutput):
191
+ reqs_to_store: dict[ReqId, TransferSpec] = {}
192
+ # iterate over both new and cached requests
193
+ for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
194
+ if preempted:
195
+ self._request_block_ids[req_id] = []
196
+
197
+ if new_block_id_groups:
198
+ new_block_ids = new_block_id_groups[0]
199
+ self._request_block_ids[req_id] += new_block_ids
200
+
201
+ block_ids = self._request_block_ids[req_id]
202
+
203
+ req = self._requests[req_id]
204
+ new_tokens = scheduler_output.num_scheduled_tokens[req_id]
205
+ expected_tokens = req.num_computed_tokens + new_tokens
206
+ # with async scheduling, some tokens may be missing
207
+ total_tokens = min(expected_tokens, req.num_tokens)
208
+ num_blocks = total_tokens // self.offloaded_block_size
209
+ start_block_idx = self._next_stored_block_idx.get(req_id, 0)
210
+ num_new_blocks = num_blocks - start_block_idx
211
+
212
+ if num_new_blocks <= 0:
213
+ continue
214
+
215
+ num_gpu_blocks = num_blocks * self.block_size_factor
216
+ assert len(req.block_hashes) >= num_gpu_blocks
217
+
218
+ new_block_hashes = self._get_block_hashes(
219
+ req, start_idx=start_block_idx, end_idx=num_blocks
220
+ )
221
+ store_output = self.manager.prepare_store(new_block_hashes)
222
+ if store_output is None:
223
+ logger.warning(
224
+ "Request %s: cannot store %s blocks", req_id, num_new_blocks
225
+ )
226
+ continue
227
+
228
+ self._next_stored_block_idx[req_id] = num_blocks
229
+
230
+ if not store_output.block_hashes_to_store:
231
+ continue
232
+ block_hashes_to_store = set(store_output.block_hashes_to_store)
233
+
234
+ block_hashes = self._get_block_hashes(req, end_idx=num_blocks)
235
+ self.manager.touch(block_hashes)
236
+
237
+ new_block_hashes = self._get_block_hashes(
238
+ req, start_idx=start_block_idx, end_idx=num_blocks
239
+ )
240
+ dst_spec = store_output.store_spec
241
+ src_block_ids: list[int] = []
242
+ for idx, blk_hash in enumerate(new_block_hashes):
243
+ if blk_hash not in block_hashes_to_store:
244
+ continue
245
+ offloaded_block_idx = start_block_idx + idx
246
+ gpu_block_idx = offloaded_block_idx * self.block_size_factor
247
+ for i in range(self.block_size_factor):
248
+ src_block_ids.append(block_ids[gpu_block_idx + i])
249
+ src_spec = GPULoadStoreSpec(src_block_ids)
250
+
251
+ reqs_to_store[req_id] = (src_spec, dst_spec)
252
+ self._reqs_being_stored[req_id] |= block_hashes_to_store
253
+
254
+ logger.debug(
255
+ "Request %s offloading %s blocks starting from block #%d",
256
+ req_id,
257
+ len(block_hashes_to_store),
258
+ start_block_idx,
259
+ )
260
+
261
+ return reqs_to_store
262
+
263
+ def build_connector_meta(
264
+ self, scheduler_output: SchedulerOutput
265
+ ) -> KVConnectorMetadata:
266
+ meta = OffloadingConnectorMetadata(
267
+ reqs_to_load=self._reqs_to_load,
268
+ reqs_to_store=self._get_reqs_to_store(scheduler_output),
269
+ reqs_to_flush=scheduler_output.preempted_req_ids,
270
+ )
271
+ self._reqs_to_load = {}
272
+
273
+ # NOTE (orozery): we should move this logic to update_connector_output
274
+ # once KVConnectorOutput allows us to report completed transfers
275
+ for req_id in scheduler_output.preempted_req_ids or ():
276
+ block_hashes = self._reqs_being_stored.get(req_id)
277
+ if block_hashes:
278
+ self.manager.complete_store(block_hashes)
279
+ block_hashes.clear()
280
+
281
+ return meta
282
+
283
+ def update_connector_output(self, connector_output: KVConnectorOutput):
284
+ """
285
+ Update KVConnector state from worker-side connectors output.
286
+
287
+ Args:
288
+ connector_output (KVConnectorOutput): the worker-side
289
+ connectors output.
290
+ """
291
+ for req_id in connector_output.finished_sending or []:
292
+ block_hashes = self._reqs_being_stored.pop(req_id, None)
293
+ if block_hashes:
294
+ self.manager.complete_store(block_hashes)
295
+
296
+ for req_id in connector_output.finished_recving or []:
297
+ block_hashes = self._reqs_being_loaded.pop(req_id, None)
298
+ if block_hashes:
299
+ if self._blocks_being_loaded:
300
+ self._blocks_being_loaded.difference_update(block_hashes)
301
+ self.manager.complete_load(block_hashes)
302
+
303
+ def request_finished(
304
+ self,
305
+ request: Request,
306
+ block_ids: list[int],
307
+ ) -> tuple[bool, dict[str, Any] | None]:
308
+ """
309
+ Called when a request has finished, before its blocks are freed.
310
+
311
+ Returns:
312
+ True if the request is being saved/sent asynchronously and blocks
313
+ should not be freed until the request_id is returned from
314
+ get_finished().
315
+ Optional KVTransferParams to be included in the request outputs
316
+ returned by the engine.
317
+ """
318
+ req_id = request.request_id
319
+ self._requests.pop(req_id, None)
320
+ self._request_block_ids.pop(req_id, None)
321
+
322
+ # TODO(orozery): possibly kickoff offload for last block
323
+ # which may have been deferred due to async scheduling
324
+ self._next_stored_block_idx.pop(req_id, None)
325
+
326
+ request_being_stored = req_id in self._reqs_being_stored
327
+ return request_being_stored, None
328
+
329
+ def take_events(self) -> Iterable[KVCacheEvent]:
330
+ """Take the KV cache events from the connector.
331
+
332
+ Returns:
333
+ A list of KV cache events.
334
+ """
335
+ for event in self.manager.take_events():
336
+ if event.removed:
337
+ yield BlockRemoved(block_hashes=event.block_hashes, medium=event.medium)
338
+ else:
339
+ yield BlockStored(
340
+ block_hashes=event.block_hashes,
341
+ parent_block_hash=None,
342
+ token_ids=[],
343
+ lora_id=None,
344
+ block_size=event.block_size,
345
+ medium=event.medium,
346
+ lora_name=None,
347
+ )
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
new
+185
@@ -0,0 +1,185 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections import defaultdict
4
+
5
+import torch
6
+
7
+from vllm.config import get_layers_from_vllm_config
8
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
9
+ KVConnectorStats,
10
+)
11
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
12
+ OffloadingConnectorMetadata,
13
+ ReqId,
14
+)
15
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
16
+ OffloadingConnectorStats,
17
+)
18
+from vllm.logger import init_logger
19
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
20
+from vllm.v1.attention.backend import AttentionBackend
21
+from vllm.v1.kv_offload.spec import OffloadingSpec
22
+from vllm.v1.kv_offload.worker.worker import (
23
+ OffloadingWorker,
24
+ TransferSpec,
25
+)
26
+
27
+logger = init_logger(__name__)
28
+
29
+
30
+class OffloadingConnectorWorker:
31
+ """Implementation of Worker side methods"""
32
+
33
+ def __init__(self, spec: OffloadingSpec):
34
+ self.spec = spec
35
+ self.worker = OffloadingWorker()
36
+
37
+ self._job_counter = 0
38
+
39
+ self.kv_connector_stats = OffloadingConnectorStats()
40
+ # req_id -> (job_id, store)
41
+ self._jobs: dict[int, tuple[ReqId, bool]] = {}
42
+ # req_id -> active job IDs
43
+ self._load_job: dict[ReqId, int] = {}
44
+ # req_id -> set(active job IDs)
45
+ self._store_jobs = defaultdict[ReqId, set[int]](set)
46
+ # list of store jobs pending submission (job_id, transfer_spec)
47
+ self._unsubmitted_store_jobs: list[tuple[int, TransferSpec]] = []
48
+
49
+ self._finished_reqs_waiting_for_store: set[ReqId] = set()
50
+
51
+ def _generate_job_id(self) -> int:
52
+ job_id = self._job_counter
53
+ self._job_counter = job_id + 1
54
+ return job_id
55
+
56
+ def _register_handlers(
57
+ self,
58
+ kv_caches: dict[str, torch.Tensor],
59
+ attn_backends: dict[str, type[AttentionBackend]],
60
+ ):
61
+ for src_cls, dst_cls, handler in self.spec.get_handlers(
62
+ kv_caches, attn_backends
63
+ ):
64
+ self.worker.register_handler(src_cls, dst_cls, handler)
65
+
66
+ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
67
+ layer_names = list(kv_caches.keys())
68
+ layers = get_layers_from_vllm_config(
69
+ self.spec.vllm_config,
70
+ AttentionLayerBase, # type: ignore[type-abstract]
71
+ layer_names,
72
+ )
73
+ attn_backends = {
74
+ layer_name: layers[layer_name].get_attn_backend()
75
+ for layer_name in layer_names
76
+ }
77
+ self._register_handlers(kv_caches, attn_backends)
78
+
79
+ def register_cross_layers_kv_cache(
80
+ self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
81
+ ):
82
+ cross_layer_name = "ALL_LAYERS"
83
+ kv_caches = {cross_layer_name: kv_cache}
84
+ attn_backends = {cross_layer_name: attn_backend}
85
+ self._register_handlers(kv_caches, attn_backends)
86
+
87
+ def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata):
88
+ for job_id, transfer_spec in self._unsubmitted_store_jobs:
89
+ success = self.worker.transfer_async(job_id, transfer_spec)
90
+ assert success
91
+ self._unsubmitted_store_jobs.clear()
92
+
93
+ for req_id in kv_connector_metadata.reqs_to_flush or ():
94
+ job_ids = self._store_jobs.get(req_id)
95
+ if job_ids:
96
+ self.worker.wait(job_ids)
97
+
98
+ def start_kv_transfers(self, metadata: OffloadingConnectorMetadata):
99
+ for job_id, transfer_spec in self._unsubmitted_store_jobs:
100
+ success = self.worker.transfer_async(job_id, transfer_spec)
101
+ assert success
102
+ self._unsubmitted_store_jobs.clear()
103
+
104
+ for req_id, transfer_spec in metadata.reqs_to_load.items():
105
+ job_id = self._generate_job_id()
106
+ self._jobs[job_id] = (req_id, False)
107
+ assert req_id not in self._load_job
108
+ self._load_job[req_id] = job_id
109
+ success = self.worker.transfer_async(job_id, transfer_spec)
110
+ assert success
111
+
112
+ def prepare_store_kv(self, metadata: OffloadingConnectorMetadata):
113
+ for req_id, transfer_spec in metadata.reqs_to_store.items():
114
+ job_id = self._generate_job_id()
115
+ self._jobs[job_id] = (req_id, True)
116
+ self._store_jobs[req_id].add(job_id)
117
+ # NOTE(orozery): defer the store to the beginning of the next engine step,
118
+ # so that offloading starts AFTER transfers related to token sampling,
119
+ # thereby avoiding delays to token generation due to offloading.
120
+ self._unsubmitted_store_jobs.append((job_id, transfer_spec))
121
+
122
+ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
123
+ """
124
+ Notifies worker-side connector ids of requests that have
125
+ finished generating tokens.
126
+ Returns a list of request IDs that finished loading or storing.
127
+
128
+ Returns:
129
+ ids of requests that have finished asynchronous transfer
130
+ tuple of (sending/saving ids, recving/loading ids).
131
+ """
132
+ finished_sending = set()
133
+ finished_recving = set()
134
+ for transfer_result in self.worker.get_finished():
135
+ # we currently do not support job failures
136
+ job_id = transfer_result.job_id
137
+ assert transfer_result.success
138
+ req_id, store = self._jobs.pop(job_id)
139
+ if (
140
+ transfer_result.transfer_time
141
+ and transfer_result.transfer_size is not None
142
+ and transfer_result.transfer_type is not None
143
+ ):
144
+ self.kv_connector_stats.record_transfer(
145
+ num_bytes=transfer_result.transfer_size,
146
+ time=transfer_result.transfer_time,
147
+ transfer_type=transfer_result.transfer_type,
148
+ )
149
+ if store:
150
+ req_jobs = self._store_jobs[req_id]
151
+ req_jobs.remove(job_id)
152
+ if req_jobs:
153
+ continue
154
+
155
+ if req_id in self._finished_reqs_waiting_for_store:
156
+ self._finished_reqs_waiting_for_store.remove(req_id)
157
+ finished_sending.add(req_id)
158
+ del self._store_jobs[req_id]
159
+ else:
160
+ req_job = self._load_job[req_id]
161
+ assert job_id == req_job
162
+ del self._load_job[req_id]
163
+ finished_recving.add(req_id)
164
+
165
+ for req_id in finished_req_ids:
166
+ pending_req_jobs = self._store_jobs.get(req_id)
167
+ if pending_req_jobs:
168
+ self._finished_reqs_waiting_for_store.add(req_id)
169
+ elif pending_req_jobs is not None:
170
+ finished_sending.add(req_id)
171
+ del self._store_jobs[req_id]
172
+
173
+ return finished_sending, finished_recving
174
+
175
+ def get_kv_connector_stats(self) -> KVConnectorStats | None:
176
+ """
177
+ Get the KV transfer stats for the connector.
178
+ """
179
+
180
+ if self.kv_connector_stats.is_empty():
181
+ return None
182
+ # Clear stats for next iteration
183
+ kv_connector_stats = self.kv_connector_stats
184
+ self.kv_connector_stats = OffloadingConnectorStats()
185
+ return kv_connector_stats
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+15
-656
@@ -1,16 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections import defaultdict
3
from collections.abc import Iterable
5
-from dataclasses import dataclass
6
-from itertools import islice
4
from typing import Any
5
6
import torch
7
11
-from vllm.config import VllmConfig, get_layers_from_vllm_config
12
-from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
13
-from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
8
+from vllm.config import VllmConfig
9
+from vllm.distributed.kv_events import KVCacheEvent
10
from vllm.distributed.kv_transfer.kv_connector.v1 import (
11
KVConnectorBase_V1,
12
KVConnectorRole,
@@ -22,97 +18,28 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
18
PromMetric,
19
PromMetricT,
20
)
21
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
22
+ OffloadingConnectorMetadata,
23
+)
24
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
25
+ OffloadingConnectorStats,
26
+ OffloadPromMetrics,
27
+)
28
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
29
+ OffloadingConnectorScheduler,
30
+)
31
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
32
+ OffloadingConnectorWorker,
33
+)
34
from vllm.forward_context import ForwardContext
26
-from vllm.logger import init_logger
27
-from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
35
from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
36
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
30
-from vllm.v1.core.kv_cache_utils import BlockHash
37
from vllm.v1.core.sched.output import SchedulerOutput
38
from vllm.v1.kv_cache_interface import KVCacheConfig
33
-from vllm.v1.kv_offload.abstract import OffloadingManager
39
from vllm.v1.kv_offload.factory import OffloadingSpecFactory
35
-from vllm.v1.kv_offload.mediums import GPULoadStoreSpec
36
-from vllm.v1.kv_offload.spec import OffloadingSpec
37
-from vllm.v1.kv_offload.worker.worker import (
38
- OffloadingWorker,
39
- TransferSpec,
40
- TransferType,
41
-)
40
from vllm.v1.outputs import KVConnectorOutput
41
from vllm.v1.request import Request
42
45
-ReqId = str
46
-
47
-logger = init_logger(__name__)
48
-
49
-
50
-@dataclass
51
-class OffloadingOperationMetrics:
52
- op_size: int
53
- op_time: float
54
-
55
-
56
-@dataclass
57
-class OffloadingConnectorStats(KVConnectorStats):
58
- def __post_init__(self):
59
- if not self.data:
60
- # Empty container init, no data is passed in.
61
- self.reset()
62
-
63
- def reset(self):
64
- self.data: dict[str, list[OffloadingOperationMetrics]] = {}
65
-
66
- def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
67
- if not other.is_empty():
68
- for k, v in other.data.items():
69
- if k not in self.data:
70
- self.data[k] = v
71
- else:
72
- accumulator = self.data[k]
73
- assert isinstance(accumulator, list)
74
- accumulator.extend(v)
75
- return self
76
-
77
- def reduce(self) -> dict[str, int | float]:
78
- """
79
- Reduce the observations collected during a time interval to one or
80
- more representative values (eg avg/median/sum of the series).
81
- This is meant to be called by the logger to produce a summary of the
82
- stats for the last time interval.
83
- """
84
- return_dict: dict[str, int | float] = {}
85
- for transfer_type, ops_list in self.data.items():
86
- assert isinstance(ops_list, list)
87
- total_bytes = 0
88
- total_time = 0.0
89
- for op in ops_list:
90
- assert isinstance(op, dict)
91
- total_bytes += op["op_size"]
92
- total_time += op["op_time"]
93
- return_dict[f"{transfer_type}_total_bytes"] = total_bytes
94
- return_dict[f"{transfer_type}_total_time"] = total_time
95
- return return_dict
96
-
97
- def is_empty(self) -> bool:
98
- return not self.data
99
-
100
- def record_transfer(self, num_bytes: int, time: float, transfer_type: TransferType):
101
- src, dst = transfer_type
102
- transfer_type_key = src + "_to_" + dst
103
- op = OffloadingOperationMetrics(num_bytes, time)
104
- if transfer_type_key in self.data:
105
- self.data[transfer_type_key].append(op)
106
- else:
107
- self.data[transfer_type_key] = [op]
108
-
109
-
110
-@dataclass
111
-class OffloadingConnectorMetadata(KVConnectorMetadata):
112
- reqs_to_load: dict[ReqId, TransferSpec]
113
- reqs_to_store: dict[ReqId, TransferSpec]
114
- reqs_to_flush: set[str] | None = None
115
-
43
44
class OffloadingConnector(KVConnectorBase_V1):
45
@property
@@ -242,571 +169,3 @@ class OffloadingConnector(KVConnectorBase_V1):
169
return OffloadPromMetrics(
170
vllm_config, metric_types, labelnames, per_engine_labelvalues
171
)
245
-
246
-
247
-class OffloadingConnectorScheduler:
248
- """Implementation of Scheduler side methods"""
249
-
250
- def __init__(self, spec: OffloadingSpec):
251
- assert len(spec.gpu_block_size) == 1
252
- self.gpu_block_size = spec.gpu_block_size[0]
253
- self.offloaded_block_size = self.gpu_block_size * spec.block_size_factor
254
- self.block_size_factor = spec.block_size_factor
255
- self.manager: OffloadingManager = spec.get_manager()
256
-
257
- self._requests: dict[ReqId, Request] = {}
258
- # list of GPU block IDs per request
259
- self._request_block_ids: dict[ReqId, list[int]] = {}
260
- # requests to load for the current scheduler step
261
- self._reqs_to_load: dict[ReqId, TransferSpec] = {}
262
- # request blocks are stored in order
263
- # index of next block (of size offloaded_block_size) to offload
264
- self._next_stored_block_idx: dict[ReqId, int] = {}
265
- # if GPU prefix caching is enabled,
266
- # track loaded blocks to avoid redundant loads
267
- self._blocks_being_loaded: set[BlockHash] | None = (
268
- set() if spec.vllm_config.cache_config.enable_prefix_caching else None
269
- )
270
-
271
- # request ID -> set(block hashes being stored/load)
272
- self._reqs_being_stored = defaultdict[ReqId, set[BlockHash]](set)
273
- self._reqs_being_loaded = defaultdict[ReqId, set[BlockHash]](set)
274
-
275
- def _get_block_hashes(
276
- self,
277
- req: Request,
278
- start_idx: int = 0,
279
- end_idx: int | None = None,
280
- ) -> Iterable[BlockHash]:
281
- return islice(
282
- req.block_hashes,
283
- self.block_size_factor * start_idx + self.block_size_factor - 1,
284
- self.block_size_factor * end_idx if end_idx else None,
285
- self.block_size_factor,
286
- )
287
-
288
- def get_num_new_matched_tokens(
289
- self, request: Request, num_computed_tokens: int
290
- ) -> tuple[int | None, bool]:
291
- """
292
- Get number of new tokens that can be loaded beyond the
293
- num_computed_tokens.
294
-
295
- Args:
296
- request (Request): the request object.
297
- num_computed_tokens (int): the number of locally
298
- computed tokens for this request
299
-
300
- Returns:
301
- A tuple with the following elements:
302
- - The number of tokens that can be loaded beyond what is
303
- already computed.
304
- If None, it means that the connector needs more time to
305
- determine the number of matched tokens, and the scheduler
306
- should query for this request again later.
307
- - `True` if tokens will be loaded asynchronously
308
- (between scheduler steps).
309
- """
310
- num_blocks = request.num_tokens // self.offloaded_block_size
311
-
312
- assert len(request.block_hashes) // self.block_size_factor == num_blocks
313
- block_hashes = self._get_block_hashes(request)
314
-
315
- self.manager.touch(block_hashes)
316
-
317
- full_block_tokens = self.offloaded_block_size * num_blocks
318
- if full_block_tokens - num_computed_tokens < self.offloaded_block_size:
319
- # we can load less than a block, skip
320
- return 0, False
321
-
322
- start_block_idx = num_computed_tokens // self.offloaded_block_size
323
- hits = self.manager.lookup(
324
- self._get_block_hashes(request, start_idx=start_block_idx)
325
- )
326
- if hits is None:
327
- # indicates a lookup that should be tried later
328
- return None, False
329
- if hits == 0:
330
- return 0, False
331
-
332
- num_hit_tokens = (
333
- self.offloaded_block_size * (start_block_idx + hits) - num_computed_tokens
334
- )
335
- logger.debug(
336
- "Request %s hit %s offloaded tokens after %s GPU hit tokens",
337
- request.request_id,
338
- num_hit_tokens,
339
- num_computed_tokens,
340
- )
341
- if num_hit_tokens < self.offloaded_block_size:
342
- return 0, False
343
-
344
- if self._blocks_being_loaded:
345
- block_hashes = self._get_block_hashes(
346
- request, start_idx=start_block_idx, end_idx=start_block_idx + hits
347
- )
348
-
349
- if any(
350
- block_hash in self._blocks_being_loaded for block_hash in block_hashes
351
- ):
352
- # hit blocks are being loaded, delay request
353
- logger.debug(
354
- "Delaying request %s since some of its blocks are already"
355
- " being loaded",
356
- request.request_id,
357
- )
358
- return None, False
359
-
360
- return num_hit_tokens, True
361
-
362
- def update_state_after_alloc(
363
- self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int
364
- ):
365
- self._requests[request.request_id] = request
366
- # the block ids are updated in _get_reqs_to_store
367
- self._request_block_ids[request.request_id] = []
368
-
369
- if num_external_tokens == 0:
370
- return
371
-
372
- block_groups = blocks.get_block_ids()
373
- block_ids = block_groups[0]
374
-
375
- num_computed_gpu_blocks = sum(
376
- block.block_hash is not None for block in blocks.blocks[0]
377
- )
378
- num_computed_tokens = num_computed_gpu_blocks * self.gpu_block_size
379
- full_block_tokens = num_computed_tokens + num_external_tokens
380
- assert full_block_tokens % self.offloaded_block_size == 0
381
-
382
- num_pending_gpu_blocks = len(block_ids) - num_computed_gpu_blocks
383
- assert num_external_tokens == num_pending_gpu_blocks * self.gpu_block_size
384
-
385
- start_block_idx = num_computed_tokens // self.offloaded_block_size
386
- num_blocks = full_block_tokens // self.offloaded_block_size
387
-
388
- assert len(request.block_hashes) // self.block_size_factor >= num_blocks
389
- block_hashes = self._get_block_hashes(
390
- request, start_idx=start_block_idx, end_idx=num_blocks
391
- )
392
-
393
- src_spec = self.manager.prepare_load(block_hashes)
394
- dst_spec = GPULoadStoreSpec(block_ids[num_computed_gpu_blocks:])
395
-
396
- block_hashes = self._get_block_hashes(
397
- request, start_idx=start_block_idx, end_idx=num_blocks
398
- )
399
-
400
- self._reqs_to_load[request.request_id] = (src_spec, dst_spec)
401
- req_blocks_being_loaded = self._reqs_being_loaded[request.request_id]
402
- req_blocks_being_loaded.update(block_hashes)
403
- self._next_stored_block_idx[request.request_id] = num_blocks
404
-
405
- if self._blocks_being_loaded is not None:
406
- self._blocks_being_loaded.update(req_blocks_being_loaded)
407
-
408
- def _get_reqs_to_store(self, scheduler_output: SchedulerOutput):
409
- reqs_to_store: dict[ReqId, TransferSpec] = {}
410
- # iterate over both new and cached requests
411
- for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
412
- if preempted:
413
- self._request_block_ids[req_id] = []
414
-
415
- if new_block_id_groups:
416
- new_block_ids = new_block_id_groups[0]
417
- self._request_block_ids[req_id] += new_block_ids
418
-
419
- block_ids = self._request_block_ids[req_id]
420
-
421
- req = self._requests[req_id]
422
- new_tokens = scheduler_output.num_scheduled_tokens[req_id]
423
- expected_tokens = req.num_computed_tokens + new_tokens
424
- # with async scheduling, some tokens may be missing
425
- total_tokens = min(expected_tokens, req.num_tokens)
426
- num_blocks = total_tokens // self.offloaded_block_size
427
- start_block_idx = self._next_stored_block_idx.get(req_id, 0)
428
- num_new_blocks = num_blocks - start_block_idx
429
-
430
- if num_new_blocks <= 0:
431
- continue
432
-
433
- num_gpu_blocks = num_blocks * self.block_size_factor
434
- assert len(req.block_hashes) >= num_gpu_blocks
435
-
436
- new_block_hashes = self._get_block_hashes(
437
- req, start_idx=start_block_idx, end_idx=num_blocks
438
- )
439
- store_output = self.manager.prepare_store(new_block_hashes)
440
- if store_output is None:
441
- logger.warning(
442
- "Request %s: cannot store %s blocks", req_id, num_new_blocks
443
- )
444
- continue
445
-
446
- self._next_stored_block_idx[req_id] = num_blocks
447
-
448
- if not store_output.block_hashes_to_store:
449
- continue
450
- block_hashes_to_store = set(store_output.block_hashes_to_store)
451
-
452
- block_hashes = self._get_block_hashes(req, end_idx=num_blocks)
453
- self.manager.touch(block_hashes)
454
-
455
- new_block_hashes = self._get_block_hashes(
456
- req, start_idx=start_block_idx, end_idx=num_blocks
457
- )
458
- dst_spec = store_output.store_spec
459
- src_block_ids: list[int] = []
460
- for idx, blk_hash in enumerate(new_block_hashes):
461
- if blk_hash not in block_hashes_to_store:
462
- continue
463
- offloaded_block_idx = start_block_idx + idx
464
- gpu_block_idx = offloaded_block_idx * self.block_size_factor
465
- for i in range(self.block_size_factor):
466
- src_block_ids.append(block_ids[gpu_block_idx + i])
467
- src_spec = GPULoadStoreSpec(src_block_ids)
468
-
469
- reqs_to_store[req_id] = (src_spec, dst_spec)
470
- self._reqs_being_stored[req_id] |= block_hashes_to_store
471
-
472
- logger.debug(
473
- "Request %s offloading %s blocks starting from block #%d",
474
- req_id,
475
- len(block_hashes_to_store),
476
- start_block_idx,
477
- )
478
-
479
- return reqs_to_store
480
-
481
- def build_connector_meta(
482
- self, scheduler_output: SchedulerOutput
483
- ) -> KVConnectorMetadata:
484
- meta = OffloadingConnectorMetadata(
485
- reqs_to_load=self._reqs_to_load,
486
- reqs_to_store=self._get_reqs_to_store(scheduler_output),
487
- reqs_to_flush=scheduler_output.preempted_req_ids,
488
- )
489
- self._reqs_to_load = {}
490
-
491
- # NOTE (orozery): we should move this logic to update_connector_output
492
- # once KVConnectorOutput allows us to report completed transfers
493
- for req_id in scheduler_output.preempted_req_ids or ():
494
- block_hashes = self._reqs_being_stored.get(req_id)
495
- if block_hashes:
496
- self.manager.complete_store(block_hashes)
497
- block_hashes.clear()
498
-
499
- return meta
500
-
501
- def update_connector_output(self, connector_output: KVConnectorOutput):
502
- """
503
- Update KVConnector state from worker-side connectors output.
504
-
505
- Args:
506
- connector_output (KVConnectorOutput): the worker-side
507
- connectors output.
508
- """
509
- for req_id in connector_output.finished_sending or []:
510
- block_hashes = self._reqs_being_stored.pop(req_id, None)
511
- if block_hashes:
512
- self.manager.complete_store(block_hashes)
513
-
514
- for req_id in connector_output.finished_recving or []:
515
- block_hashes = self._reqs_being_loaded.pop(req_id, None)
516
- if block_hashes:
517
- if self._blocks_being_loaded:
518
- self._blocks_being_loaded.difference_update(block_hashes)
519
- self.manager.complete_load(block_hashes)
520
-
521
- def request_finished(
522
- self,
523
- request: Request,
524
- block_ids: list[int],
525
- ) -> tuple[bool, dict[str, Any] | None]:
526
- """
527
- Called when a request has finished, before its blocks are freed.
528
-
529
- Returns:
530
- True if the request is being saved/sent asynchronously and blocks
531
- should not be freed until the request_id is returned from
532
- get_finished().
533
- Optional KVTransferParams to be included in the request outputs
534
- returned by the engine.
535
- """
536
- req_id = request.request_id
537
- self._requests.pop(req_id, None)
538
- self._request_block_ids.pop(req_id, None)
539
-
540
- # TODO(orozery): possibly kickoff offload for last block
541
- # which may have been deferred due to async scheduling
542
- self._next_stored_block_idx.pop(req_id, None)
543
-
544
- request_being_stored = req_id in self._reqs_being_stored
545
- return request_being_stored, None
546
-
547
- def take_events(self) -> Iterable[KVCacheEvent]:
548
- """Take the KV cache events from the connector.
549
-
550
- Returns:
551
- A list of KV cache events.
552
- """
553
- for event in self.manager.take_events():
554
- if event.removed:
555
- yield BlockRemoved(block_hashes=event.block_hashes, medium=event.medium)
556
- else:
557
- yield BlockStored(
558
- block_hashes=event.block_hashes,
559
- parent_block_hash=None,
560
- token_ids=[],
561
- lora_id=None,
562
- block_size=event.block_size,
563
- medium=event.medium,
564
- lora_name=None,
565
- )
566
-
567
-
568
-class OffloadingConnectorWorker:
569
- """Implementation of Worker side methods"""
570
-
571
- def __init__(self, spec: OffloadingSpec):
572
- self.spec = spec
573
- self.worker = OffloadingWorker()
574
-
575
- self._job_counter = 0
576
-
577
- self.kv_connector_stats = OffloadingConnectorStats()
578
- # req_id -> (job_id, store)
579
- self._jobs: dict[int, tuple[ReqId, bool]] = {}
580
- # req_id -> active job IDs
581
- self._load_job: dict[ReqId, int] = {}
582
- # req_id -> set(active job IDs)
583
- self._store_jobs = defaultdict[ReqId, set[int]](set)
584
- # list of store jobs pending submission (job_id, transfer_spec)
585
- self._unsubmitted_store_jobs: list[tuple[int, TransferSpec]] = []
586
-
587
- self._finished_reqs_waiting_for_store: set[ReqId] = set()
588
-
589
- def _generate_job_id(self) -> int:
590
- job_id = self._job_counter
591
- self._job_counter = job_id + 1
592
- return job_id
593
-
594
- def _register_handlers(
595
- self,
596
- kv_caches: dict[str, torch.Tensor],
597
- attn_backends: dict[str, type[AttentionBackend]],
598
- ):
599
- for src_cls, dst_cls, handler in self.spec.get_handlers(
600
- kv_caches, attn_backends
601
- ):
602
- self.worker.register_handler(src_cls, dst_cls, handler)
603
-
604
- def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
605
- layer_names = list(kv_caches.keys())
606
- layers = get_layers_from_vllm_config(
607
- self.spec.vllm_config,
608
- AttentionLayerBase, # type: ignore[type-abstract]
609
- layer_names,
610
- )
611
- attn_backends = {
612
- layer_name: layers[layer_name].get_attn_backend()
613
- for layer_name in layer_names
614
- }
615
- self._register_handlers(kv_caches, attn_backends)
616
-
617
- def register_cross_layers_kv_cache(
618
- self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
619
- ):
620
- cross_layer_name = "ALL_LAYERS"
621
- kv_caches = {cross_layer_name: kv_cache}
622
- attn_backends = {cross_layer_name: attn_backend}
623
- self._register_handlers(kv_caches, attn_backends)
624
-
625
- def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata):
626
- for job_id, transfer_spec in self._unsubmitted_store_jobs:
627
- success = self.worker.transfer_async(job_id, transfer_spec)
628
- assert success
629
- self._unsubmitted_store_jobs.clear()
630
-
631
- for req_id in kv_connector_metadata.reqs_to_flush or ():
632
- job_ids = self._store_jobs.get(req_id)
633
- if job_ids:
634
- self.worker.wait(job_ids)
635
-
636
- def start_kv_transfers(self, metadata: OffloadingConnectorMetadata):
637
- for job_id, transfer_spec in self._unsubmitted_store_jobs:
638
- success = self.worker.transfer_async(job_id, transfer_spec)
639
- assert success
640
- self._unsubmitted_store_jobs.clear()
641
-
642
- for req_id, transfer_spec in metadata.reqs_to_load.items():
643
- job_id = self._generate_job_id()
644
- self._jobs[job_id] = (req_id, False)
645
- assert req_id not in self._load_job
646
- self._load_job[req_id] = job_id
647
- success = self.worker.transfer_async(job_id, transfer_spec)
648
- assert success
649
-
650
- def prepare_store_kv(self, metadata: OffloadingConnectorMetadata):
651
- for req_id, transfer_spec in metadata.reqs_to_store.items():
652
- job_id = self._generate_job_id()
653
- self._jobs[job_id] = (req_id, True)
654
- self._store_jobs[req_id].add(job_id)
655
- # NOTE(orozery): defer the store to the beginning of the next engine step,
656
- # so that offloading starts AFTER transfers related to token sampling,
657
- # thereby avoiding delays to token generation due to offloading.
658
- self._unsubmitted_store_jobs.append((job_id, transfer_spec))
659
-
660
- def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
661
- """
662
- Notifies worker-side connector ids of requests that have
663
- finished generating tokens.
664
- Returns a list of request IDs that finished loading or storing.
665
-
666
- Returns:
667
- ids of requests that have finished asynchronous transfer
668
- tuple of (sending/saving ids, recving/loading ids).
669
- """
670
- finished_sending = set()
671
- finished_recving = set()
672
- for transfer_result in self.worker.get_finished():
673
- # we currently do not support job failures
674
- job_id = transfer_result.job_id
675
- assert transfer_result.success
676
- req_id, store = self._jobs.pop(job_id)
677
- if (
678
- transfer_result.transfer_time
679
- and transfer_result.transfer_size is not None
680
- and transfer_result.transfer_type is not None
681
- ):
682
- self.kv_connector_stats.record_transfer(
683
- num_bytes=transfer_result.transfer_size,
684
- time=transfer_result.transfer_time,
685
- transfer_type=transfer_result.transfer_type,
686
- )
687
- if store:
688
- req_jobs = self._store_jobs[req_id]
689
- req_jobs.remove(job_id)
690
- if req_jobs:
691
- continue
692
-
693
- if req_id in self._finished_reqs_waiting_for_store:
694
- self._finished_reqs_waiting_for_store.remove(req_id)
695
- finished_sending.add(req_id)
696
- del self._store_jobs[req_id]
697
- else:
698
- req_job = self._load_job[req_id]
699
- assert job_id == req_job
700
- del self._load_job[req_id]
701
- finished_recving.add(req_id)
702
-
703
- for req_id in finished_req_ids:
704
- pending_req_jobs = self._store_jobs.get(req_id)
705
- if pending_req_jobs:
706
- self._finished_reqs_waiting_for_store.add(req_id)
707
- elif pending_req_jobs is not None:
708
- finished_sending.add(req_id)
709
- del self._store_jobs[req_id]
710
-
711
- return finished_sending, finished_recving
712
-
713
- def get_kv_connector_stats(self) -> KVConnectorStats | None:
714
- """
715
- Get the KV transfer stats for the connector.
716
- """
717
-
718
- if self.kv_connector_stats.is_empty():
719
- return None
720
- # Clear stats for next iteration
721
- kv_connector_stats = self.kv_connector_stats
722
- self.kv_connector_stats = OffloadingConnectorStats()
723
- return kv_connector_stats
724
-
725
-
726
-class OffloadPromMetrics(KVConnectorPromMetrics):
727
- def __init__(
728
- self,
729
- vllm_config: VllmConfig,
730
- metric_types: dict[type[PromMetric], type[PromMetricT]],
731
- labelnames: list[str],
732
- per_engine_labelvalues: dict[int, list[object]],
733
- ):
734
- super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
735
- # (engine_idx, transfer_type) -> (metric with bounded labels)
736
- self.histogram_transfer_size: dict[tuple[int, str], PromMetricT] = {}
737
- self.counter_kv_bytes: dict[tuple[int, str], PromMetricT] = {}
738
- self.counter_kv_transfer_time: dict[tuple[int, str], PromMetricT] = {}
739
- buckets = [ # In bytes
740
- 1e6,
741
- 5e6,
742
- 10e6,
743
- 20e6,
744
- 40e6,
745
- 60e6,
746
- 80e6,
747
- 100e6,
748
- 150e6,
749
- 200e6,
750
- ]
751
-
752
- self._counter_kv_bytes = self._counter_cls(
753
- name="vllm:kv_offload_total_bytes",
754
- documentation="Number of bytes offloaded by KV connector",
755
- labelnames=labelnames + ["transfer_type"],
756
- )
757
-
758
- self._counter_kv_transfer_time = self._counter_cls(
759
- name="vllm:kv_offload_total_time",
760
- documentation="Total time measured by all KV offloading operations",
761
- labelnames=labelnames + ["transfer_type"],
762
- )
763
-
764
- self._histogram_transfer_size = self._histogram_cls(
765
- name="vllm:kv_offload_size",
766
- documentation="Histogram of KV offload transfer size, in bytes.",
767
- buckets=buckets[:],
768
- labelnames=labelnames + ["transfer_type"],
769
- )
770
-
771
- def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
772
- """
773
- Observe transfer statistics from the new data structure.
774
- transfer_stats_data is expected to be a dict where:
775
- - keys are transfer type strings (e.g., "cpu_to_gpu", "gpu_to_cpu")
776
- - values are lists of OffloadingOperationMetrics objects
777
- """
778
-
779
- for transfer_type, ops in transfer_stats_data.items():
780
- # Cache:
781
- if (engine_idx, transfer_type) not in self.histogram_transfer_size:
782
- self.histogram_transfer_size[(engine_idx, transfer_type)] = (
783
- self._histogram_transfer_size.labels(
784
- *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
785
- )
786
- )
787
- self.counter_kv_bytes[(engine_idx, transfer_type)] = (
788
- self._counter_kv_bytes.labels(
789
- *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
790
- )
791
- )
792
- self.counter_kv_transfer_time[(engine_idx, transfer_type)] = (
793
- self._counter_kv_transfer_time.labels(
794
- *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
795
- )
796
- )
797
-
798
- # Process ops:
799
- assert isinstance(ops, list)
800
- for op in ops: # ops is a list of serialized OffloadingOperationMetrics
801
- assert isinstance(op, dict)
802
- # Observe size histogram
803
- self.histogram_transfer_size[(engine_idx, transfer_type)].observe(
804
- op["op_size"]
805
- )
806
-
807
- # Increment byte and time counters
808
- self.counter_kv_bytes[(engine_idx, transfer_type)].inc(op["op_size"])
809
-
810
- self.counter_kv_transfer_time[(engine_idx, transfer_type)].inc(
811
- op["op_time"]
812
- )