[Feat][Core] Add disk offloading support to SimpleCPUOffloadConnector (#49644)
Signed-off-by: Guanyi Chen <939416532@qq.com> Co-authored-by: Claude <noreply@anthropic.com>
Guanyi Chen committed
Aug 7, 2026 at 15:37 UTC
58fcaa0baaa32ba0c34e1119f6ce4554ef8a6256
4 files changed
+466
-19
vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
+51
-2
@@ -41,6 +41,15 @@ logger = init_logger(__name__)
41
# Default CPU capacity: 8 GB
42
DEFAULT_CPU_CAPACITY_BYTES = 8 * (1024**3)
43
44
+VALID_KV_OFFLOAD_BACKENDS = ("cpu", "disk")
45
+# Keys that only apply to the disk backend, warned about under "cpu".
46
+_DISK_ONLY_KEYS = (
47
+ "disk_path",
48
+ "disk_capacity_bytes",
49
+ "disk_buffer_slots",
50
+ "use_page_cache",
51
+)
52
+
53
54
class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
55
"""CPU KV cache offloading with custom kernel transfers and BlockPool LRU."""
@@ -82,6 +91,36 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
91
92
lazy_offload = bool(extra_config.get("lazy_offload", False))
93
94
+ kv_offload_backend = str(extra_config.get("kv_offload_backend", "cpu"))
95
+ if kv_offload_backend not in VALID_KV_OFFLOAD_BACKENDS:
96
+ raise ValueError(
97
+ f"Unknown kv_offload_backend {kv_offload_backend!r}; "
98
+ f"expected one of {VALID_KV_OFFLOAD_BACKENDS}"
99
+ )
100
+ disk_mode = kv_offload_backend == "disk"
101
+
102
+ disk_path = extra_config.get("disk_path", None) or None
103
+ disk_capacity_bytes = int(
104
+ extra_config.get("disk_capacity_bytes", 100 * (1024**3))
105
+ )
106
+ disk_buffer_slots = max(1, int(extra_config.get("disk_buffer_slots", 2)))
107
+ use_page_cache = bool(extra_config.get("use_page_cache", False))
108
+
109
+ if disk_mode:
110
+ if disk_path is None:
111
+ raise ValueError(
112
+ 'kv_offload_backend="disk" requires disk_path to be set.'
113
+ )
114
+ else:
115
+ ignored = [k for k in _DISK_ONLY_KEYS if k in extra_config]
116
+ if ignored:
117
+ logger.warning(
118
+ 'kv_offload_backend is "cpu", ignoring disk-only config: %s. '
119
+ 'Set kv_offload_backend="disk" to enable the disk backend.',
120
+ ", ".join(ignored),
121
+ )
122
+ disk_path = None
123
+
124
self.scheduler_manager: SimpleCPUOffloadScheduler | None = None
125
self.worker_handler: SimpleCPUOffloadWorker | None = None
126
@@ -94,11 +133,13 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
133
134
logger.info(
135
"SimpleCPUOffloadConnector: role=%s, "
97
- "per_rank=%.2f GB, world_size=%d, mode=%s",
136
+ "per_rank=%.2f GB, world_size=%d, mode=%s, backend=%s, disk=%s",
137
role.name,
138
cpu_capacity_per_rank / (1024**3),
139
world_size,
140
"lazy" if lazy_offload else "eager",
141
+ kv_offload_backend,
142
+ disk_path or "none",
143
)
144
145
if role == KVConnectorRole.SCHEDULER:
@@ -115,10 +156,18 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
156
scheduler_block_size=scheduler_block_size,
157
hash_block_size=hash_block_size,
158
lazy_offload=lazy_offload,
159
+ disk_capacity_bytes=disk_capacity_bytes if disk_mode else 0,
160
)
161
elif role == KVConnectorRole.WORKER:
162
self.worker_handler = SimpleCPUOffloadWorker(
121
- vllm_config, kv_cache_config, cpu_capacity_per_rank
163
+ vllm_config,
164
+ kv_cache_config,
165
+ cpu_capacity_per_rank,
166
+ kv_offload_backend=kv_offload_backend,
167
+ disk_path=disk_path,
168
+ disk_capacity_bytes=disk_capacity_bytes,
169
+ disk_buffer_slots=disk_buffer_slots,
170
+ use_page_cache=use_page_cache,
171
)
172
173
# --- Worker-side methods ---
vllm/v1/simple_kv_offload/disk_backend.py
new
+336
@@ -0,0 +1,336 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Disk I/O backend for GPU<->NVMe block transfers via pinned staging buffers.
4
+
5
+Uses separate IO threads for store and load so that loads (latency-critical)
6
+never block behind stores (background work). Each thread owns its own pinned
7
+staging buffers to avoid contention.
8
+"""
9
+
10
+from __future__ import annotations
11
+
12
+import contextlib
13
+import os
14
+import queue
15
+import threading
16
+
17
+import torch
18
+
19
+from vllm.logger import init_logger
20
+from vllm.platforms import current_platform
21
+from vllm.v1.simple_kv_offload.cuda_mem_ops import (
22
+ CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
23
+ CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
24
+ BatchMemcpyParams,
25
+ build_params,
26
+ copy_blocks,
27
+ pin_tensor,
28
+)
29
+
30
+logger = init_logger(__name__)
31
+
32
+O_DIRECT = getattr(os, "O_DIRECT", 0)
33
+_ALIGNMENT = 4096
34
+
35
+
36
+def _alloc_aligned(num_slots: int, bpb: int) -> torch.Tensor:
37
+ """Allocate a staging buffer whose base address is O_DIRECT aligned.
38
+
39
+ The CPU allocator only guarantees 64-byte alignment, so over-allocate by
40
+ one alignment unit and return an aligned view. The view keeps the backing
41
+ storage alive.
42
+ """
43
+ nbytes = num_slots * bpb
44
+ raw = torch.zeros(nbytes + _ALIGNMENT, dtype=torch.int8, device="cpu")
45
+ offset = -raw.data_ptr() % _ALIGNMENT
46
+ return raw[offset : offset + nbytes].view(num_slots, bpb)
47
+
48
+
49
+class DiskBackend:
50
+ """Async disk offload backend with pipelined GPU DMA and interleaved IO.
51
+
52
+ Architecture:
53
+ - Separate coordinator threads for store and load (never block each other)
54
+ - Interleaved pipeline: DMA slot N while preadv/pwritev slot N-1
55
+ - O_DIRECT by default; page cache is opt-in via use_page_cache
56
+
57
+ Same launch_copy interface as DmaCopyBackend so the worker can swap
58
+ backends without changing calling code.
59
+ """
60
+
61
+ def __init__(self) -> None:
62
+ self._store_params: BatchMemcpyParams | None = None
63
+ self._load_params: BatchMemcpyParams | None = None
64
+ self._load_stream: torch.cuda.Stream | None = None
65
+ self._store_stream: torch.cuda.Stream | None = None
66
+ self._store_queue: queue.SimpleQueue = queue.SimpleQueue()
67
+ self._load_queue: queue.SimpleQueue = queue.SimpleQueue()
68
+ self._store_thread: threading.Thread | None = None
69
+ self._load_thread: threading.Thread | None = None
70
+ self._shutdown: bool = False
71
+ self._fd: int = -1
72
+ self._disk_path: str = ""
73
+ self._total_block_bytes: int = 0
74
+ self._store_buffer_caches: dict[str, torch.Tensor] = {}
75
+ self._load_buffer_caches: dict[str, torch.Tensor] = {}
76
+ self._store_slot_views: list[list[memoryview]] = []
77
+ self._load_slot_views: list[list[memoryview]] = []
78
+ self._per_tensor_bpb: list[int] = []
79
+ self._tensor_names: list[str] = []
80
+
81
+ def init(
82
+ self,
83
+ gpu_caches: dict[str, torch.Tensor],
84
+ device: torch.device,
85
+ load_stream: torch.cuda.Stream,
86
+ store_stream: torch.cuda.Stream,
87
+ disk_path: str,
88
+ num_disk_slots: int,
89
+ total_block_bytes: int,
90
+ num_buffer_slots: int = 2,
91
+ use_page_cache: bool = False,
92
+ ) -> None:
93
+ self._load_stream = load_stream
94
+ self._store_stream = store_stream
95
+ self._total_block_bytes = total_block_bytes
96
+ self._num_buffer_slots = num_buffer_slots
97
+ self._tensor_names = list(gpu_caches.keys())
98
+ self._per_tensor_bpb = [
99
+ t.stride(0) * t.element_size() for t in gpu_caches.values()
100
+ ]
101
+
102
+ assert total_block_bytes % _ALIGNMENT == 0, (
103
+ f"total_block_bytes={total_block_bytes} not aligned to {_ALIGNMENT}"
104
+ )
105
+
106
+ # Separate buffer pools for store and load threads
107
+ self._store_buffer_caches = {}
108
+ self._load_buffer_caches = {}
109
+ for name, gpu_t in gpu_caches.items():
110
+ bpb = gpu_t.stride(0) * gpu_t.element_size()
111
+ store_buf = _alloc_aligned(num_buffer_slots, bpb)
112
+ pin_tensor(store_buf)
113
+ self._store_buffer_caches[name] = store_buf
114
+ load_buf = _alloc_aligned(num_buffer_slots, bpb)
115
+ pin_tensor(load_buf)
116
+ self._load_buffer_caches[name] = load_buf
117
+
118
+ # Pre-built iovec views per slot (avoid per-transfer .numpy() calls)
119
+ self._store_slot_views = [
120
+ [
121
+ memoryview(self._store_buffer_caches[name][slot].numpy())
122
+ for name in self._tensor_names
123
+ ]
124
+ for slot in range(num_buffer_slots)
125
+ ]
126
+ self._load_slot_views = [
127
+ [
128
+ memoryview(self._load_buffer_caches[name][slot].numpy())
129
+ for name in self._tensor_names
130
+ ]
131
+ for slot in range(num_buffer_slots)
132
+ ]
133
+
134
+ self._store_params = build_params(
135
+ gpu_caches,
136
+ self._store_buffer_caches,
137
+ store_stream,
138
+ src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
139
+ )
140
+ self._load_params = build_params(
141
+ self._load_buffer_caches,
142
+ gpu_caches,
143
+ load_stream,
144
+ src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
145
+ )
146
+
147
+ os.makedirs(os.path.dirname(disk_path) or ".", exist_ok=True)
148
+ # Slot contents never outlive the process, so unlink then O_EXCL rather
149
+ # than reopening: a pre-existing file would otherwise keep its own
150
+ # (possibly world-readable) mode, and blocks may encode user prompts.
151
+ with contextlib.suppress(FileNotFoundError):
152
+ os.unlink(disk_path)
153
+ # O_DIRECT by default: page cache would consume the very host DRAM this
154
+ # backend exists to conserve, and doubles the copy on the store path.
155
+ flags = os.O_RDWR | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW
156
+ if not use_page_cache:
157
+ flags |= O_DIRECT
158
+ self._fd = os.open(disk_path, flags, 0o600)
159
+ self._disk_path = disk_path
160
+ os.ftruncate(self._fd, num_disk_slots * total_block_bytes)
161
+
162
+ logger.info(
163
+ "DiskBackend: path=%s, slots=%d, total=%.2f GB, buf=%dx%d bytes"
164
+ " (page_cache=%s)",
165
+ disk_path,
166
+ num_disk_slots,
167
+ (num_disk_slots * total_block_bytes) / (1024**3),
168
+ num_buffer_slots,
169
+ total_block_bytes,
170
+ use_page_cache,
171
+ )
172
+
173
+ self._store_thread = threading.Thread(
174
+ target=self._store_loop,
175
+ args=(device, store_stream),
176
+ daemon=True,
177
+ )
178
+ self._load_thread = threading.Thread(
179
+ target=self._load_loop,
180
+ args=(device, load_stream),
181
+ daemon=True,
182
+ )
183
+ self._store_thread.start()
184
+ self._load_thread.start()
185
+
186
+ def launch_copy(
187
+ self,
188
+ src_blocks: list[int],
189
+ dst_blocks: list[int],
190
+ is_store: bool,
191
+ event_idx: int,
192
+ events_list: list[tuple[int, torch.Event]],
193
+ wait_event: torch.Event | None = None,
194
+ ) -> None:
195
+ q = self._store_queue if is_store else self._load_queue
196
+ q.put((src_blocks, dst_blocks, event_idx, events_list, wait_event))
197
+
198
+ def shutdown(self) -> None:
199
+ if self._shutdown:
200
+ return
201
+ self._shutdown = True
202
+ self._store_queue.put(None)
203
+ self._load_queue.put(None)
204
+ if self._store_thread is not None:
205
+ self._store_thread.join(timeout=10.0)
206
+ if self._load_thread is not None:
207
+ self._load_thread.join(timeout=10.0)
208
+ if self._fd < 0:
209
+ return
210
+ # Slot contents can encode user prompts, so drop the name now rather
211
+ # than leaving them readable until the next run overwrites the file.
212
+ # Unlinking only removes the directory entry: any thread still holding
213
+ # the fd keeps writing to the (now anonymous) inode, which the kernel
214
+ # frees once the last fd goes away.
215
+ with contextlib.suppress(OSError):
216
+ os.unlink(self._disk_path)
217
+ # Closing under a still-running IO thread would let the fd number be
218
+ # reused by an unrelated open(), turning its next pwritev into a write
219
+ # into that file. Leaking one fd for the remaining process lifetime is
220
+ # the cheaper failure.
221
+ if any(
222
+ t is not None and t.is_alive()
223
+ for t in (self._store_thread, self._load_thread)
224
+ ):
225
+ logger.warning(
226
+ "IO thread still running after shutdown timeout; leaking fd %d",
227
+ self._fd,
228
+ )
229
+ return
230
+ os.close(self._fd)
231
+ self._fd = -1
232
+
233
+ def _store_loop(
234
+ self,
235
+ device: torch.device,
236
+ stream: torch.cuda.Stream,
237
+ ) -> None:
238
+ current_platform.set_device(device)
239
+ while True:
240
+ item = self._store_queue.get()
241
+ if item is None:
242
+ return
243
+ (src_blocks, dst_blocks, event_idx, events_list, wait_event) = item
244
+ if wait_event is not None:
245
+ stream.wait_event(wait_event)
246
+ self._do_store(src_blocks, dst_blocks, stream)
247
+ event = torch.Event()
248
+ event.record(stream)
249
+ events_list.append((event_idx, event))
250
+
251
+ def _writev_slot(self, buf_slot: int, file_offset: int) -> None:
252
+ written = os.pwritev(self._fd, self._store_slot_views[buf_slot], file_offset)
253
+ if written < self._total_block_bytes:
254
+ raise OSError(
255
+ f"Short write: expected {self._total_block_bytes} bytes, "
256
+ f"wrote {written}"
257
+ )
258
+
259
+ def _readv_slot(self, buf_slot: int, file_offset: int) -> None:
260
+ bytes_read = os.preadv(self._fd, self._load_slot_views[buf_slot], file_offset)
261
+ if bytes_read < self._total_block_bytes:
262
+ raise OSError(
263
+ f"Short read: expected {self._total_block_bytes} bytes, "
264
+ f"read {bytes_read}"
265
+ )
266
+
267
+ def _load_loop(
268
+ self,
269
+ device: torch.device,
270
+ stream: torch.cuda.Stream,
271
+ ) -> None:
272
+ current_platform.set_device(device)
273
+ while True:
274
+ item = self._load_queue.get()
275
+ if item is None:
276
+ return
277
+ (src_blocks, dst_blocks, event_idx, events_list, wait_event) = item
278
+ if wait_event is not None:
279
+ stream.wait_event(wait_event)
280
+ self._do_load(src_blocks, dst_blocks, stream)
281
+ event = torch.Event()
282
+ event.record(stream)
283
+ events_list.append((event_idx, event))
284
+
285
+ def _do_store(
286
+ self,
287
+ gpu_blocks: list[int],
288
+ disk_slots: list[int],
289
+ stream: torch.cuda.Stream,
290
+ ) -> None:
291
+ """GPU -> buffer (DMA) -> disk (pwritev), interleaved double-buffer."""
292
+ assert self._store_params is not None
293
+ n = self._num_buffer_slots
294
+ # (DMA event, file offset) of the block already staged in each slot.
295
+ pending: list[tuple[torch.Event, int] | None] = [None] * n
296
+
297
+ for i, (gpu_blk, disk_slot) in enumerate(zip(gpu_blocks, disk_slots)):
298
+ buf_slot = i % n
299
+ prev = pending[buf_slot]
300
+ if prev is not None:
301
+ prev[0].synchronize()
302
+ self._writev_slot(buf_slot, prev[1])
303
+
304
+ copy_blocks([gpu_blk], [buf_slot], self._store_params)
305
+ ev = torch.Event()
306
+ ev.record(stream)
307
+ pending[buf_slot] = (ev, disk_slot * self._total_block_bytes)
308
+
309
+ for slot, last in enumerate(pending):
310
+ if last is not None:
311
+ last[0].synchronize()
312
+ self._writev_slot(slot, last[1])
313
+
314
+ def _do_load(
315
+ self,
316
+ disk_slots: list[int],
317
+ gpu_blocks: list[int],
318
+ stream: torch.cuda.Stream,
319
+ ) -> None:
320
+ """Disk (preadv) -> buffer -> GPU (DMA), interleaved double-buffer."""
321
+ assert self._load_params is not None
322
+ n = self._num_buffer_slots
323
+ prev_dma_events: list[torch.Event | None] = [None] * n
324
+
325
+ for i, (disk_slot, gpu_blk) in enumerate(zip(disk_slots, gpu_blocks)):
326
+ buf_slot = i % n
327
+ prev = prev_dma_events[buf_slot]
328
+ if prev is not None:
329
+ prev.synchronize()
330
+
331
+ self._readv_slot(buf_slot, disk_slot * self._total_block_bytes)
332
+
333
+ copy_blocks([buf_slot], [gpu_blk], self._load_params)
334
+ ev = torch.Event()
335
+ ev.record(stream)
336
+ prev_dma_events[buf_slot] = ev
vllm/v1/simple_kv_offload/manager.py
+10
-3
@@ -75,9 +75,14 @@ class SimpleCPUOffloadScheduler:
75
scheduler_block_size: int,
76
hash_block_size: int,
77
lazy_offload: bool = False,
78
+ disk_capacity_bytes: int = 0,
79
):
80
self.vllm_config = vllm_config
81
self.kv_cache_config = kv_cache_config
82
+ # When disk mode is active, the offload pool size is disk-based.
83
+ offload_capacity = (
84
+ disk_capacity_bytes if disk_capacity_bytes > 0 else cpu_capacity_bytes
85
+ )
86
self.enable_kv_cache_events = (
87
vllm_config.kv_events_config is not None
88
and vllm_config.kv_events_config.enable_kv_cache_events
@@ -90,7 +95,7 @@ class SimpleCPUOffloadScheduler:
95
# Derive a CPU KVCacheConfig from the GPU config and build a coordinator
96
assert kv_cache_config is not None
97
self.cpu_kv_cache_config = self._derive_cpu_config(
93
- kv_cache_config, cpu_capacity_bytes
98
+ kv_cache_config, offload_capacity
99
)
100
self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks
101
# Find the full attention kv group for prefix cache matching.
@@ -111,10 +116,12 @@ class SimpleCPUOffloadScheduler:
116
assert self.block_size % self.fa_block_size == 0
117
118
logger.info(
114
- "SimpleCPUOffloadScheduler: Allocating %d CPU blocks (%.2f GB, mode=%s)",
119
+ "SimpleCPUOffloadScheduler: Allocating %d offload blocks "
120
+ "(%.2f GB, mode=%s, backend=%s)",
121
self.num_cpu_blocks,
116
- cpu_capacity_bytes / (1024**3),
122
+ offload_capacity / (1024**3),
123
"lazy" if lazy_offload else "eager",
124
+ "disk" if disk_capacity_bytes > 0 else "cpu",
125
)
126
127
spec_config = vllm_config.speculative_config
vllm/v1/simple_kv_offload/worker.py
+69
-14
@@ -11,6 +11,7 @@ from vllm.logger import init_logger
11
from vllm.utils.torch_utils import PIN_MEMORY
12
from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend
13
from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor
14
+from vllm.v1.simple_kv_offload.disk_backend import DiskBackend
15
from vllm.v1.simple_kv_offload.metadata import (
16
SimpleCPUOffloadMetadata,
17
SimpleCPUOffloadWorkerMetadata,
@@ -30,10 +31,20 @@ class SimpleCPUOffloadWorker:
31
vllm_config: VllmConfig,
32
kv_cache_config: "KVCacheConfig | None",
33
cpu_capacity_bytes: int,
34
+ kv_offload_backend: str = "cpu",
35
+ disk_path: str | None = None,
36
+ disk_capacity_bytes: int = 0,
37
+ disk_buffer_slots: int = 2,
38
+ use_page_cache: bool = False,
39
):
40
self.vllm_config = vllm_config
41
self.kv_cache_config = kv_cache_config
42
self.cpu_capacity_bytes = cpu_capacity_bytes
43
+ self.disk_path = disk_path
44
+ self.disk_capacity_bytes = disk_capacity_bytes
45
+ self.disk_buffer_slots = disk_buffer_slots
46
+ self.use_page_cache = use_page_cache
47
+ self.disk_mode = kv_offload_backend == "disk"
48
49
self.gpu_kv_caches: dict[str, torch.Tensor] | None = None
50
self.cpu_kv_caches: dict[str, torch.Tensor] | None = None
@@ -44,7 +55,7 @@ class SimpleCPUOffloadWorker:
55
self.load_stream: torch.cuda.Stream | None = None
56
self.store_stream: torch.cuda.Stream | None = None
57
47
- self._backend = DmaCopyBackend()
58
+ self._backend: DmaCopyBackend | DiskBackend | None = None
59
60
# Ordered (event_idx, Event). Events pre-allocated on main thread.
61
self._load_events: list[tuple[int, torch.Event]] = []
@@ -145,9 +156,57 @@ class SimpleCPUOffloadWorker:
156
157
self.num_cpu_blocks = max(1, self.cpu_capacity_bytes // total_bytes_per_block)
158
159
+ # Use lowest priority so KV cache I/O yields to compute streams.
160
+ low_pri, _ = torch.cuda.Stream.priority_range()
161
+ self.load_stream = torch.cuda.Stream(priority=low_pri)
162
+ self.store_stream = torch.cuda.Stream(priority=low_pri)
163
+
164
+ self.gpu_kv_caches = unique_gpu_caches
165
+
166
+ if self.disk_mode:
167
+ self._init_disk_mode(unique_gpu_caches, total_bytes_per_block, self.device)
168
+ else:
169
+ self._init_cpu_mode(unique_gpu_caches, total_bytes_per_block, self.device)
170
+
171
+ def _init_disk_mode(
172
+ self,
173
+ unique_gpu_caches: dict[str, torch.Tensor],
174
+ total_bytes_per_block: int,
175
+ device: torch.device,
176
+ ) -> None:
177
+ num_disk_slots = max(1, self.disk_capacity_bytes // total_bytes_per_block)
178
+ self.num_cpu_blocks = num_disk_slots
179
+
180
logger.info(
149
- "SimpleCPUOffloadWorker: %d unique GPU KV tensors, "
150
- "allocating %d CPU blocks (%.2f GB)",
181
+ "SimpleCPUOffloadWorker [DISK]: %d tensors, %d disk slots (%.2f GB)",
182
+ len(unique_gpu_caches),
183
+ num_disk_slots,
184
+ (num_disk_slots * total_bytes_per_block) / (1024**3),
185
+ )
186
+
187
+ assert self.disk_path is not None
188
+ rank_path = f"{self.disk_path}.rank_{device.index or 0}"
189
+ self._backend = DiskBackend()
190
+ self._backend.init(
191
+ unique_gpu_caches,
192
+ device,
193
+ self.load_stream,
194
+ self.store_stream,
195
+ rank_path,
196
+ num_disk_slots,
197
+ total_bytes_per_block,
198
+ self.disk_buffer_slots,
199
+ self.use_page_cache,
200
+ )
201
+
202
+ def _init_cpu_mode(
203
+ self,
204
+ unique_gpu_caches: dict[str, torch.Tensor],
205
+ total_bytes_per_block: int,
206
+ device: torch.device,
207
+ ) -> None:
208
+ logger.info(
209
+ "SimpleCPUOffloadWorker [CPU]: %d tensors, %d CPU blocks (%.2f GB)",
210
len(unique_gpu_caches),
211
self.num_cpu_blocks,
212
(self.num_cpu_blocks * total_bytes_per_block) / (1024**3),
@@ -159,7 +218,6 @@ class SimpleCPUOffloadWorker:
218
"Pinned memory not available. CPU offload performance may be degraded."
219
)
220
162
- self.gpu_kv_caches = unique_gpu_caches
221
self.cpu_kv_caches = {}
222
for name, gpu_tensor in unique_gpu_caches.items():
223
cpu_shape = (self.num_cpu_blocks,) + gpu_tensor.shape[1:]
@@ -171,16 +229,11 @@ class SimpleCPUOffloadWorker:
229
pin_tensor(tensor)
230
self.cpu_kv_caches[name] = tensor
231
174
- # Use lowest priority so KV cache I/O yields to compute streams.
175
- low_pri, _ = torch.cuda.Stream.priority_range()
176
- self.load_stream = torch.cuda.Stream(priority=low_pri)
177
- self.store_stream = torch.cuda.Stream(priority=low_pri)
178
-
179
- # Initialize copy backend with caches and streams.
232
+ self._backend = DmaCopyBackend()
233
self._backend.init(
181
- self.gpu_kv_caches,
234
+ unique_gpu_caches,
235
self.cpu_kv_caches,
183
- self.device,
236
+ device,
237
self.load_stream,
238
self.store_stream,
239
)
@@ -224,8 +277,10 @@ class SimpleCPUOffloadWorker:
277
# (1) Submit transfers
278
metadata = self._connector_metadata
279
if metadata is not None:
280
+ backend = self._backend
281
+ assert backend is not None
282
if metadata.load_cpu_blocks:
228
- self._backend.launch_copy(
283
+ backend.launch_copy(
284
metadata.load_cpu_blocks,
285
metadata.load_gpu_blocks,
286
is_store=False,
@@ -236,7 +291,7 @@ class SimpleCPUOffloadWorker:
291
if self._store_compute_done is None:
292
self._store_compute_done = torch.Event()
293
self._store_compute_done.record(torch.cuda.current_stream())
239
- self._backend.launch_copy(
294
+ backend.launch_copy(
295
metadata.store_gpu_blocks,
296
metadata.store_cpu_blocks,
297
is_store=True,