@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7227d06156
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 9, 13:33 UTC

[Metrics] [KVConnector] Add Offloading Connector metrics (#27942)

Added queries and hits metrics for the Offloading Connector. Also added timing metrics for store and load operations, which take the average time it takes to load/store, per-token. The metrics are available from Prometheus and from the StatLogger. Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Co-authored-by: Omer Paz <Omer.Paz@ibm.com>

omerpaz95 committed Jan 27, 2026 at 15:34 UTC 7227d06156e12b6f6ac2d04f34d8316d2f60e0f5
6 files changed +450 -29
tests/v1/kv_connector/unit/test_offloading_connector.py
+150 -1
@@ -16,6 +16,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
16 from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
17 OffloadingConnector,
18 OffloadingConnectorMetadata,
19 + OffloadingConnectorStats,
20 )
21 from vllm.forward_context import ForwardContext
22 from vllm.utils.hashing import sha256
@@ -86,7 +87,14 @@ class MockOffloadingHandler(OffloadingHandler):
87 if job_id in self.waiting_jobs:
88 self.waiting_jobs.remove(job_id)
89 self.completed_jobs.append(job_id)
89 - self.completed_transfers.append((job_id, True))
90 + result = TransferResult(
91 + job_id=job_id,
92 + success=True,
93 + transfer_size=None,
94 + transfer_time=None,
95 + transfer_type=None,
96 + )
97 + self.completed_transfers.append(result)
98
99 def wait(self, job_ids: set[int]) -> None:
100 self.flushed_jobs |= job_ids
@@ -720,3 +728,144 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner):
728
729 # second request will use the GPU prefix cache
730 assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
731 +
732 +
733 +class TestOffloadingConnectorStats:
734 + """Tests for OffloadingConnector stats reconstruction and operations."""
735 +
736 + def test_build_kv_connector_stats_with_none(self):
737 + """Test that build_kv_connector_stats returns empty stats when given None."""
738 + stats = OffloadingConnector.build_kv_connector_stats(data=None)
739 +
740 + assert stats is not None
741 + assert isinstance(stats, OffloadingConnectorStats)
742 + assert len(stats.data) == 0
743 + assert stats.is_empty()
744 +
745 + def test_build_kv_connector_stats_with_empty_dict(self):
746 + """Test that build_kv_connector_stats returns empty stats with empty dict."""
747 + stats = OffloadingConnector.build_kv_connector_stats(data={})
748 +
749 + assert stats is not None
750 + assert isinstance(stats, OffloadingConnectorStats)
751 + assert len(stats.data) == 0
752 + assert stats.is_empty()
753 +
754 + def test_build_kv_connector_stats_reconstructs_offload_stats(self):
755 + """Test that OffloadingConnector stats are properly reconstructed with
756 + correct data."""
757 + serialized_data = {
758 + "CPU_to_GPU": [
759 + {"op_size": 16, "op_time": 1.0},
760 + {"op_size": 8, "op_time": 0.5},
761 + ],
762 + "GPU_to_CPU": [
763 + {"op_size": 1, "op_time": 0.1},
764 + {"op_size": 2, "op_time": 0.2},
765 + ],
766 + }
767 +
768 + stats = OffloadingConnector.build_kv_connector_stats(data=serialized_data)
769 +
770 + offload_connector_stats = stats
771 + assert isinstance(offload_connector_stats, OffloadingConnectorStats)
772 + assert offload_connector_stats.data["CPU_to_GPU"] == [
773 + {"op_size": 16, "op_time": 1.0},
774 + {"op_size": 8, "op_time": 0.5},
775 + ]
776 + assert offload_connector_stats.data["GPU_to_CPU"] == [
777 + {"op_size": 1, "op_time": 0.1},
778 + {"op_size": 2, "op_time": 0.2},
779 + ]
780 +
781 + def test_aggregate_same_connector(self):
782 + """Test aggregating stats from the same connector type."""
783 + stats1 = OffloadingConnectorStats(
784 + data={
785 + "CPU_to_GPU": [
786 + {"op_size": 16, "op_time": 1.0},
787 + {"op_size": 8, "op_time": 0.5},
788 + ],
789 + "GPU_to_CPU": [
790 + {"op_size": 1, "op_time": 0.1},
791 + {"op_size": 2, "op_time": 0.2},
792 + ],
793 + }
794 + )
795 +
796 + stats2 = OffloadingConnectorStats(
797 + data={
798 + "CPU_to_GPU": [
799 + {"op_size": 3, "op_time": 0.2},
800 + {"op_size": 7, "op_time": 0.9},
801 + ],
802 + "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
803 + }
804 + )
805 +
806 + result = stats1.aggregate(stats2)
807 +
808 + assert result is stats1 # Should return self
809 + offload_connector_stats = result
810 + assert offload_connector_stats.data["CPU_to_GPU"] == [
811 + {"op_size": 16, "op_time": 1.0},
812 + {"op_size": 8, "op_time": 0.5},
813 + {"op_size": 3, "op_time": 0.2},
814 + {"op_size": 7, "op_time": 0.9},
815 + ]
816 + assert offload_connector_stats.data["GPU_to_CPU"] == [
817 + {"op_size": 1, "op_time": 0.1},
818 + {"op_size": 2, "op_time": 0.2},
819 + {"op_size": 16, "op_time": 2},
820 + ]
821 +
822 + def test_reduce(self):
823 + """Test that reduce() correctly reduces all nested connector stats."""
824 + stats = OffloadingConnectorStats(
825 + data={
826 + "CPU_to_GPU": [
827 + {"op_size": 16, "op_time": 1.0},
828 + {"op_size": 8, "op_time": 0.5},
829 + {"op_size": 3, "op_time": 0.2},
830 + {"op_size": 7, "op_time": 0.9},
831 + ],
832 + "GPU_to_CPU": [
833 + {"op_size": 1, "op_time": 0.1},
834 + {"op_size": 2, "op_time": 0.2},
835 + {"op_size": 16, "op_time": 2},
836 + ],
837 + }
838 + )
839 +
840 + reduced = stats.reduce()
841 +
842 + assert isinstance(reduced, dict)
843 + # Check that the stats were reduced (should have aggregated values)
844 + assert "CPU_to_GPU_total_bytes" in reduced
845 + assert "CPU_to_GPU_total_time" in reduced
846 + assert "GPU_to_CPU_total_bytes" in reduced
847 + assert "GPU_to_CPU_total_time" in reduced
848 + assert reduced["CPU_to_GPU_total_bytes"] == 34
849 + assert reduced["CPU_to_GPU_total_time"] == 2.6
850 + assert reduced["GPU_to_CPU_total_time"] == 2.3
851 + assert reduced["GPU_to_CPU_total_bytes"] == 19
852 +
853 + def test_reset(self):
854 + """Test that reset() resets all nested connector stats."""
855 + offload_connector_stats = OffloadingConnectorStats(
856 + data={
857 + "CPU_to_GPU": [
858 + {"op_size": 3, "op_time": 0.2},
859 + {"op_size": 7, "op_time": 0.9},
860 + ],
861 + "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
862 + }
863 + )
864 +
865 + assert not offload_connector_stats.is_empty()
866 +
867 + offload_connector_stats.reset()
868 +
869 + # After reset, stats should be empty
870 + assert offload_connector_stats.is_empty()
871 + assert len(offload_connector_stats.data) == 0
tests/v1/kv_offload/test_cpu_gpu.py
+17 -2
@@ -168,15 +168,30 @@ def test_transfer(
168 orig_dst_caches = [x.clone() for x in handler.dst_tensors]
169
170 # call transfer function
171 + start_time = time.time()
172 assert handler.transfer_async(1, (src_spec, dst_spec))
172 - assert set({x[0] for x in handler._transfers}) == {1}
173 + assert set({x.job_id for x in handler._transfers}) == {1}
174
175 # wait for transfer to complete
176 end_time = time.time() + 10
177 while time.time() < end_time:
178 finished = handler.get_finished()
179 if finished:
179 - assert finished == [(1, True)]
180 + assert finished[0].job_id == 1
181 + assert finished[0].success
182 + assert (
183 + finished[0].transfer_type == ("GPU", "CPU")
184 + if gpu_to_cpu
185 + else ("CPU", "GPU")
186 + )
187 + assert (
188 + finished[0].transfer_size
189 + == handler.total_block_size_in_bytes
190 + * handler.dst_block_size_factor
191 + * len(dst_blocks)
192 + )
193 + assert finished[0].transfer_time > 0
194 + assert finished[0].transfer_time < (time.time() - start_time)
195 break
196 time.sleep(0.1)
197
vllm/distributed/kv_transfer/kv_connector/v1/metrics.py
+2 -2
@@ -124,7 +124,7 @@ class KVConnectorPromMetrics:
124 self._counter_cls = metric_types[Counter]
125 self._histogram_cls = metric_types[Histogram]
126 self._labelnames = labelnames
127 - self._per_engine_labelvalues = per_engine_labelvalues
127 + self.per_engine_labelvalues = per_engine_labelvalues
128
129 def make_per_engine(self, metric: PromMetric) -> dict[int, PromMetric]:
130 """
@@ -134,7 +134,7 @@ class KVConnectorPromMetrics:
134 """
135 return {
136 idx: metric.labels(*labelvalues)
137 - for idx, labelvalues in self._per_engine_labelvalues.items()
137 + for idx, labelvalues in self.per_engine_labelvalues.items()
138 }
139
140 def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+214 -4
@@ -17,6 +17,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
17 KVConnectorRole,
18 )
19 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
20 +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
21 + KVConnectorPromMetrics,
22 + KVConnectorStats,
23 + PromMetric,
24 + PromMetricT,
25 +)
26 from vllm.forward_context import ForwardContext
27 from vllm.logger import init_logger
28 from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
@@ -28,7 +34,11 @@ from vllm.v1.kv_offload.abstract import OffloadingManager
34 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
31 -from vllm.v1.kv_offload.worker.worker import OffloadingWorker, TransferSpec
37 +from vllm.v1.kv_offload.worker.worker import (
38 + OffloadingWorker,
39 + TransferSpec,
40 + TransferType,
41 +)
42 from vllm.v1.outputs import KVConnectorOutput
43 from vllm.v1.request import Request
44
@@ -37,6 +47,66 @@ ReqId = str
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
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]
@@ -143,6 +213,33 @@ class OffloadingConnector(KVConnectorBase_V1):
213 assert self.connector_scheduler is not None
214 return self.connector_scheduler.take_events()
215
216 + def get_kv_connector_stats(self) -> KVConnectorStats | None:
217 + if self.connector_worker is None:
218 + return None # We only emit stats from the worker-side
219 + return self.connector_worker.get_kv_connector_stats()
220 +
221 + @classmethod
222 + def build_kv_connector_stats(
223 + cls, data: dict[str, Any] | None = None
224 + ) -> KVConnectorStats | None:
225 + return (
226 + OffloadingConnectorStats(data=data)
227 + if data is not None
228 + else OffloadingConnectorStats()
229 + )
230 +
231 + @classmethod
232 + def build_prom_metrics(
233 + cls,
234 + vllm_config: VllmConfig,
235 + metric_types: dict[type[PromMetric], type[PromMetricT]],
236 + labelnames: list[str],
237 + per_engine_labelvalues: dict[int, list[object]],
238 + ) -> KVConnectorPromMetrics:
239 + return OffloadPromMetrics(
240 + vllm_config, metric_types, labelnames, per_engine_labelvalues
241 + )
242 +
243
244 class OffloadingConnectorScheduler:
245 """Implementation of Scheduler side methods"""
@@ -467,7 +564,8 @@ class OffloadingConnectorWorker:
564
565 self._job_counter = 0
566
470 - # job_id -> (req_id, store)
567 + self.kv_connector_stats = OffloadingConnectorStats()
568 + # req_id -> (job_id, store)
569 self._jobs: dict[int, tuple[ReqId, bool]] = {}
570 # req_id -> active job IDs
571 self._load_job: dict[ReqId, int] = {}
@@ -559,10 +657,21 @@ class OffloadingConnectorWorker:
657 """
658 finished_sending = set()
659 finished_recving = set()
562 - for job_id, success in self.worker.get_finished():
660 + for transfer_result in self.worker.get_finished():
661 # we currently do not support job failures
564 - assert success
662 + job_id = transfer_result.job_id
663 + assert transfer_result.success
664 req_id, store = self._jobs.pop(job_id)
665 + if (
666 + transfer_result.transfer_time
667 + and transfer_result.transfer_size is not None
668 + and transfer_result.transfer_type is not None
669 + ):
670 + self.kv_connector_stats.record_transfer(
671 + num_bytes=transfer_result.transfer_size,
672 + time=transfer_result.transfer_time,
673 + transfer_type=transfer_result.transfer_type,
674 + )
675 if store:
676 req_jobs = self._store_jobs[req_id]
677 req_jobs.remove(job_id)
@@ -588,3 +697,104 @@ class OffloadingConnectorWorker:
697 del self._store_jobs[req_id]
698
699 return finished_sending, finished_recving
700 +
701 + def get_kv_connector_stats(self) -> KVConnectorStats | None:
702 + """
703 + Get the KV transfer stats for the connector.
704 + """
705 +
706 + if self.kv_connector_stats.is_empty():
707 + return None
708 + # Clear stats for next iteration
709 + kv_connector_stats = self.kv_connector_stats
710 + self.kv_connector_stats = OffloadingConnectorStats()
711 + return kv_connector_stats
712 +
713 +
714 +class OffloadPromMetrics(KVConnectorPromMetrics):
715 + def __init__(
716 + self,
717 + vllm_config: VllmConfig,
718 + metric_types: dict[type[PromMetric], type[PromMetricT]],
719 + labelnames: list[str],
720 + per_engine_labelvalues: dict[int, list[object]],
721 + ):
722 + super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
723 + # (engine_idx, transfer_tupe) -> (metric with bounded labels)
724 + self.histogram_transfer_size: dict[tuple[int, str], PromMetricT] = {}
725 + self.counter_kv_bytes: dict[tuple[int, str], PromMetricT] = {}
726 + self.counter_kv_transfer_time: dict[tuple[int, str], PromMetricT] = {}
727 + buckets = [ # In bytes
728 + 1e6,
729 + 5e6,
730 + 10e6,
731 + 20e6,
732 + 40e6,
733 + 60e6,
734 + 80e6,
735 + 100e6,
736 + 150e6,
737 + 200e6,
738 + ]
739 +
740 + self._counter_kv_bytes = self._counter_cls(
741 + name="vllm:kv_offload_total_bytes",
742 + documentation="Number of bytes offloaded by KV connector",
743 + labelnames=labelnames + ["transfer_type"],
744 + )
745 +
746 + self._counter_kv_transfer_time = self._counter_cls(
747 + name="vllm:kv_offload_total_time",
748 + documentation="Total time measured by all KV offloading operations",
749 + labelnames=labelnames + ["transfer_type"],
750 + )
751 +
752 + self._histogram_transfer_size = self._histogram_cls(
753 + name="vllm:kv_offload_size",
754 + documentation="Histogram of KV offload transfer size, in bytes.",
755 + buckets=buckets[:],
756 + labelnames=labelnames + ["transfer_type"],
757 + )
758 +
759 + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
760 + """
761 + Observe transfer statistics from the new data structure.
762 + transfer_stats_data is expected to be a dict where:
763 + - keys are transfer type strings (e.g., "cpu_to_gpu", "gpu_to_cpu")
764 + - values are lists of OffloadingOperationMetrics objects
765 + """
766 +
767 + for transfer_type, ops in transfer_stats_data.items():
768 + # Cache:
769 + if (engine_idx, transfer_type) not in self.histogram_transfer_size:
770 + self.histogram_transfer_size[(engine_idx, transfer_type)] = (
771 + self._histogram_transfer_size.labels(
772 + *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
773 + )
774 + )
775 + self.counter_kv_bytes[(engine_idx, transfer_type)] = (
776 + self._counter_kv_bytes.labels(
777 + *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
778 + )
779 + )
780 + self.counter_kv_transfer_time[(engine_idx, transfer_type)] = (
781 + self._counter_kv_transfer_time.labels(
782 + *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
783 + )
784 + )
785 +
786 + # Process ops:
787 + assert isinstance(ops, list)
788 + for op in ops: # ops is a list of serialized OffloadingOperationMetrics
789 + assert isinstance(op, dict)
790 + # Observe size histogram
791 + self.histogram_transfer_size[(engine_idx, transfer_type)].observe(
792 + op["op_size"]
793 + )
794 +
795 + # Increment byte and time counters
796 + self.counter_kv_bytes[(engine_idx, transfer_type)].inc(op["op_size"])
797 +
798 + self.counter_kv_transfer_time[(engine_idx, transfer_type)].inc(
799 + op["op_time"]
800 + )
vllm/v1/kv_offload/worker/cpu_gpu.py
+56 -14
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections import deque
4 +from dataclasses import dataclass
5
6 import numpy as np
7 import torch
@@ -19,6 +20,15 @@ from vllm.v1.kv_offload.worker.worker import (
20 logger = init_logger(__name__)
21
22
23 +@dataclass
24 +class Transfer:
25 + job_id: int
26 + stream: torch.cuda.Stream
27 + start_event: torch.Event
28 + end_event: torch.Event
29 + num_bytes: int
30 +
31 +
32 def expand_block_ids(
33 block_ids: np.ndarray,
34 block_size_factor: int,
@@ -92,14 +102,15 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
102 tensor.element_size() * tensor.stride(0) * min_block_size_factor
103 for tensor in src_tensors
104 ]
105 + self.total_block_size_in_bytes = sum(self.block_size_in_bytes)
106
107 assert len(src_tensors) > 0
108 self.gpu_to_cpu: bool = self.src_tensors[0].is_cuda
98 -
109 + self.transfer_type = ("GPU", "CPU") if self.gpu_to_cpu else ("CPU", "GPU")
110 # job_id -> event
111 self._transfer_events: dict[int, torch.Event] = {}
112 # queue of transfers (job_id, stream, event)
102 - self._transfers: deque[tuple[int, torch.cuda.Stream, torch.Event]] = deque()
113 + self._transfers: deque[Transfer] = deque()
114 # list of CUDA streams available for re-use
115 self._stream_pool: list[torch.cuda.Stream] = []
116 # list of CUDA events available for re-use
@@ -132,16 +143,27 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
143 src_to_dst_tensor = torch.from_numpy(src_to_dst)
144
145 stream = self._stream_pool.pop() if self._stream_pool else torch.cuda.Stream()
135 - event = self._event_pool.pop() if self._event_pool else torch.Event()
146 + start_event = (
147 + self._event_pool.pop()
148 + if self._event_pool
149 + else torch.Event(enable_timing=True)
150 + )
151 + end_event = (
152 + self._event_pool.pop()
153 + if self._event_pool
154 + else torch.Event(enable_timing=True)
155 + )
156
157 if self.gpu_to_cpu:
158 # wait for model computation to finish before offloading
159 stream.wait_stream(torch.cuda.current_stream())
160 if self._transfers:
141 - _, _, last_event = self._transfers[-1]
161 + last_transfer: Transfer = self._transfers[-1]
162 + last_event = last_transfer.end_event
163 # assure job will start only after the previous one completes
164 stream.wait_event(last_event)
165 with torch.cuda.stream(stream):
166 + start_event.record(stream)
167 for src_tensor, dst_tensor, block_size_in_bytes in zip(
168 self.src_tensors,
169 self.dst_tensors,
@@ -153,22 +175,42 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
175 block_size_in_bytes,
176 src_to_dst_tensor,
177 )
156 - event.record(stream)
157 -
158 - self._transfer_events[job_id] = event
159 - self._transfers.append((job_id, stream, event))
178 + end_event.record(stream)
179 +
180 + self._transfer_events[job_id] = end_event
181 + self._transfers.append(
182 + Transfer(
183 + job_id=job_id,
184 + stream=stream,
185 + start_event=start_event,
186 + end_event=end_event,
187 + num_bytes=dst_sub_block_count * self.total_block_size_in_bytes,
188 + )
189 + )
190
191 # success
192 return True
193
194 def get_finished(self) -> list[TransferResult]:
195 results: list[TransferResult] = []
166 - while self._transfers and self._transfers[0][2].query():
167 - job_id, stream, event = self._transfers.popleft()
168 - results.append((job_id, True))
169 - self._stream_pool.append(stream)
170 - self._event_pool.append(event)
171 - del self._transfer_events[job_id]
196 + while self._transfers and self._transfers[0].end_event.query():
197 + transfer = self._transfers.popleft()
198 + transfer_time = (
199 + transfer.start_event.elapsed_time(transfer.end_event) * 1e-3
200 + ) # elapsed_time is in miliseconds
201 + result = TransferResult(
202 + job_id=transfer.job_id,
203 + success=True,
204 + transfer_size=transfer.num_bytes,
205 + transfer_time=transfer_time,
206 + transfer_type=self.transfer_type,
207 + )
208 +
209 + results.append(result)
210 + self._stream_pool.append(transfer.stream)
211 + self._event_pool.append(transfer.end_event)
212 + self._event_pool.append(transfer.start_event)
213 + del self._transfer_events[transfer.job_id]
214 return results
215
216 def wait(self, job_ids: set[int]):
vllm/v1/kv_offload/worker/worker.py
+11 -6
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from abc import ABC, abstractmethod
4 +from dataclasses import dataclass
5
6 from vllm.logger import init_logger
7 from vllm.v1.kv_offload.abstract import LoadStoreSpec
@@ -9,12 +10,19 @@ from vllm.v1.kv_offload.abstract import LoadStoreSpec
10 TransferSpec = tuple[LoadStoreSpec, LoadStoreSpec]
11 # transfers are forwarded to workers by (src_medium, dst_medium)
12 TransferType = tuple[str, str]
12 -# transfer result (job_id, success)
13 -TransferResult = tuple[int, bool]
13
14 logger = init_logger(__name__)
15
16
17 +@dataclass
18 +class TransferResult:
19 + job_id: int
20 + success: bool
21 + transfer_size: int | None = None # Size in bytes
22 + transfer_time: float | None = None
23 + transfer_type: TransferType | None = None
24 +
25 +
26 class OffloadingHandler(ABC):
27 """
28 OffloadingHandler class for managing asynchronous KV data transfers
@@ -57,7 +65,6 @@ class OffloadingHandler(ABC):
65 def wait(self, job_ids: set[int]) -> None:
66 """
67 Wait for jobs to finish (blocking).
60 -
68 Args:
69 job_ids: The set of job IDs to wait for.
70 """
@@ -120,7 +127,6 @@ class OffloadingWorker:
127 transfer_type = (src.medium(), dst.medium())
128 handler = self.transfer_type_to_handler.get(transfer_type)
129 assert handler is not None
123 -
130 try:
131 success = handler.transfer_async(job_id, spec)
132 except Exception as e:
@@ -137,7 +143,6 @@ class OffloadingWorker:
143 logger.warning("Failed to submit %r transfer %d", transfer_type, job_id)
144 else:
145 logger.debug("Submitted %r transfer %d: %r", transfer_type, job_id, spec)
140 -
146 return success
147
148 def get_finished(self) -> list[TransferResult]:
@@ -145,7 +150,7 @@ class OffloadingWorker:
150 Get transfers finished since last call.
151
152 Returns:
148 - A list of (job_id, success) of transfers.
153 + A list of TransferResults
154 """
155 finished = []
156 for handler in self.handlers: