[KV Offload] Use background thread for mmap / cpu_tensors pinning (#45850)
Signed-off-by: Sorryhorizon <arikara6666@gmail.com>
Yiwei Hu committed
Jun 24, 2026 at 23:13 UTC
f889325c511b9e647f432a5838e82ca1494f7ac9
1 file changed
+89
-38
vllm/v1/kv_offload/cpu/gpu_worker.py
+89
-38
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import functools
4
+import threading
5
import time
6
from collections import deque
7
from dataclasses import dataclass
@@ -120,36 +121,6 @@ def compute_sub_block_ptrs(
121
output[:] = flat[skip_count : skip_count + num_sub_blocks]
122
123
123
-def pin_mmap_region(region: SharedOffloadRegion) -> None:
124
- """Register the entire mmap as CUDA pinned memory via cudaHostRegister."""
125
- if not current_platform.is_cuda_alike():
126
- logger.info(
127
- "Skipping mmap host registration on %s; cudaHostRegister is only "
128
- "available on CUDA/ROCm.",
129
- current_platform.device_name,
130
- )
131
- return
132
-
133
- rank = region.rank
134
-
135
- base_ptr = region._base.data_ptr()
136
- result = torch.cuda.cudart().cudaHostRegister(base_ptr, region.total_size_bytes, 0)
137
- if result.value != 0:
138
- logger.warning(
139
- "cudaHostRegister failed for rank=%d (code=%d) — "
140
- "transfers will still work but may be slower (unpinned DMA)",
141
- rank,
142
- result,
143
- )
144
- else:
145
- logger.debug(
146
- "cudaHostRegister rank=%d %.2f GB",
147
- rank,
148
- region.total_size_bytes / 1e9,
149
- )
150
- region.is_pinned = True
151
-
152
-
124
def _new_descriptor_buffers(
125
num_copy_ops: int,
126
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -179,6 +150,8 @@ class SingleDirectionOffloadingHandler:
150
kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]],
151
gpu_to_cpu: bool,
152
mmap_region: SharedOffloadRegion | None = None,
153
+ pin_thread: threading.Thread | None = None,
154
+ manually_pinned_tensors: list[torch.Tensor] | None = None,
155
):
156
"""
157
Initialize a SingleDirectionOffloadingHandler.
@@ -226,6 +199,8 @@ class SingleDirectionOffloadingHandler:
199
200
# mmap_region to clean up on shutdown (gpu_to_cpu handler owns it)
201
self._mmap_region = mmap_region
202
+ self._pin_thread = pin_thread
203
+ self._manually_pinned_tensors = manually_pinned_tensors
204
# job_id -> event
205
self._transfer_events: dict[int, torch.Event] = {}
206
# queue of transfers (job_id, stream, event)
@@ -458,8 +433,23 @@ class SingleDirectionOffloadingHandler:
433
self._stream_pool.clear()
434
self._event_pool.clear()
435
self._buffer_pool.clear()
436
+
437
+ if self._pin_thread is not None:
438
+ self._pin_thread.join()
439
+ self._pin_thread = None
440
+
441
+ if self._manually_pinned_tensors is not None:
442
+ for tensor in self._manually_pinned_tensors:
443
+ result = torch.cuda.cudart().cudaHostUnregister(tensor.data_ptr())
444
+ if result.value != 0:
445
+ logger.warning(
446
+ "cudaHostUnregister failed for CPU tensor (code=%d)",
447
+ result.value,
448
+ )
449
+
450
self.src_tensors.clear()
451
self.dst_tensors.clear()
452
+
453
if self._mmap_region is not None:
454
self._mmap_region.cleanup()
455
self._mmap_region = None
@@ -481,12 +471,14 @@ class CPUOffloadingWorker(OffloadingWorker):
471
mmap_region: SharedOffloadRegion | None = None,
472
):
473
pin_memory = PIN_MEMORY
474
+ self.pin_thread: threading.Thread | None = None
475
+ self._manually_pinned_tensors: list[torch.Tensor] = []
476
+
477
logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors))
485
- if mmap_region is not None and pin_memory:
486
- pin_mmap_region(mmap_region)
478
+ self._mmap_region = mmap_region
479
480
gpu_tensors: list[torch.Tensor] = []
489
- cpu_tensors: list[torch.Tensor] = []
481
+ self.cpu_tensors: list[torch.Tensor] = []
482
for kv_cache_tensor in kv_caches.tensors:
483
gpu_page_size_bytes = kv_cache_tensor.page_size_bytes
484
gpu_tensor = kv_cache_tensor.tensor.view(torch.int8).view(
@@ -502,10 +494,13 @@ class CPUOffloadingWorker(OffloadingWorker):
494
(num_cpu_blocks, cpu_page_size_bytes),
495
dtype=torch.int8,
496
device="cpu",
505
- pin_memory=pin_memory,
497
+ # CUDA/ROCm memory is registered asynchronously below.
498
+ # Pinning here would block worker initialization; other
499
+ # hardware need PyTorch allocation-time pinning.
500
+ pin_memory=PIN_MEMORY and not current_platform.is_cuda_alike(),
501
)
502
logger.debug(
508
- "torch.zeros pinned tensor %d×%d (%.2f GB): %.3f s",
503
+ "torch.zeros tensor %d×%d (%.2f GB): %.3f s",
504
num_cpu_blocks,
505
cpu_page_size_bytes,
506
num_cpu_blocks * cpu_page_size_bytes / 1e9,
@@ -513,25 +508,81 @@ class CPUOffloadingWorker(OffloadingWorker):
508
)
509
510
gpu_tensors.append(gpu_tensor)
516
- cpu_tensors.append(cpu_tensor)
511
+ self.cpu_tensors.append(cpu_tensor)
512
+
513
+ if pin_memory:
514
+ if not current_platform.is_cuda_alike():
515
+ logger.info(
516
+ "Skipping host registration on %s; cudaHostRegister is only "
517
+ "available on CUDA/ROCm.",
518
+ current_platform.device_name,
519
+ )
520
+ else:
521
+ self.pin_thread = threading.Thread(
522
+ target=self._pin_cpu_tensors,
523
+ name="CPUTensorPinThread",
524
+ )
525
+ self.pin_thread.start()
526
+ logger.info("Starting to pin memory in background...")
527
528
self._store_handler = SingleDirectionOffloadingHandler(
529
gpu_tensors=gpu_tensors,
520
- cpu_tensors=cpu_tensors,
530
+ cpu_tensors=self.cpu_tensors,
531
block_size_factor=block_size_factor,
532
kv_cache_groups_data_refs=kv_caches.group_data_refs,
533
gpu_to_cpu=True,
534
mmap_region=mmap_region,
535
+ pin_thread=self.pin_thread,
536
+ manually_pinned_tensors=self._manually_pinned_tensors,
537
)
538
539
self._load_handler = SingleDirectionOffloadingHandler(
540
gpu_tensors=gpu_tensors,
529
- cpu_tensors=cpu_tensors,
541
+ cpu_tensors=self.cpu_tensors,
542
block_size_factor=block_size_factor,
543
kv_cache_groups_data_refs=kv_caches.group_data_refs,
544
gpu_to_cpu=False,
545
)
546
547
+ def _pin_cpu_tensors(self) -> None:
548
+ """Register the CPU offload memory as CUDA pinned memory."""
549
+
550
+ t0 = time.monotonic()
551
+ tensors_to_pin = (
552
+ [self._mmap_region._base]
553
+ if self._mmap_region is not None
554
+ else self.cpu_tensors
555
+ )
556
+ num_pinned = 0
557
+ for tensor in tensors_to_pin:
558
+ total_size_bytes = tensor.numel() * tensor.element_size()
559
+ result = torch.cuda.cudart().cudaHostRegister(
560
+ tensor.data_ptr(), total_size_bytes, 0
561
+ )
562
+ if result.value != 0:
563
+ logger.warning(
564
+ "cudaHostRegister failed for host tensor (code=%d) "
565
+ "- transfers will still work but may be slower (unpinned DMA)",
566
+ result.value,
567
+ )
568
+ continue
569
+ if self._mmap_region is not None:
570
+ self._mmap_region.is_pinned = True
571
+ else:
572
+ self._manually_pinned_tensors.append(tensor)
573
+ num_pinned += 1
574
+
575
+ logger.debug(
576
+ "cudaHostRegister pin %.2f GB",
577
+ total_size_bytes / 1e9,
578
+ )
579
+
580
+ logger.info(
581
+ "Completed CPU memory pinning: %d tensors pinned in %.3f s",
582
+ num_pinned,
583
+ time.monotonic() - t0,
584
+ )
585
+
586
def submit_store(
587
self, job_id: int, src_spec: GPULoadStoreSpec, dst_spec: LoadStoreSpec
588
) -> bool: