[BugFix] Revert "[KV Offload] Use background thread for mmap / cpu_tensors pinning" (#46958)
Signed-off-by: <> Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>
Varun Sundar Rabindranath committed
Jun 29, 2026 at 10:54 UTC
36bbecd6436d0dd4c7a27fbb09a787e00534d647
1 file changed
+38
-89
vllm/v1/kv_offload/cpu/gpu_worker.py
+38
-89
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import functools
4
-import threading
4
import time
5
from collections import deque
6
from dataclasses import dataclass
@@ -121,6 +120,36 @@ def compute_sub_block_ptrs(
120
output[:] = flat[skip_count : skip_count + num_sub_blocks]
121
122
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
+
153
def _new_descriptor_buffers(
154
num_copy_ops: int,
155
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -150,8 +179,6 @@ class SingleDirectionOffloadingHandler:
179
kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]],
180
gpu_to_cpu: bool,
181
mmap_region: SharedOffloadRegion | None = None,
153
- pin_thread: threading.Thread | None = None,
154
- manually_pinned_tensors: list[torch.Tensor] | None = None,
182
):
183
"""
184
Initialize a SingleDirectionOffloadingHandler.
@@ -199,8 +226,6 @@ class SingleDirectionOffloadingHandler:
226
227
# mmap_region to clean up on shutdown (gpu_to_cpu handler owns it)
228
self._mmap_region = mmap_region
202
- self._pin_thread = pin_thread
203
- self._manually_pinned_tensors = manually_pinned_tensors
229
# job_id -> event
230
self._transfer_events: dict[int, torch.Event] = {}
231
# queue of transfers (job_id, stream, event)
@@ -433,23 +458,8 @@ class SingleDirectionOffloadingHandler:
458
self._stream_pool.clear()
459
self._event_pool.clear()
460
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
-
461
self.src_tensors.clear()
462
self.dst_tensors.clear()
452
-
463
if self._mmap_region is not None:
464
self._mmap_region.cleanup()
465
self._mmap_region = None
@@ -471,14 +481,12 @@ class CPUOffloadingWorker(OffloadingWorker):
481
mmap_region: SharedOffloadRegion | None = None,
482
):
483
pin_memory = PIN_MEMORY
474
- self.pin_thread: threading.Thread | None = None
475
- self._manually_pinned_tensors: list[torch.Tensor] = []
476
-
484
logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors))
478
- self._mmap_region = mmap_region
485
+ if mmap_region is not None and pin_memory:
486
+ pin_mmap_region(mmap_region)
487
488
gpu_tensors: list[torch.Tensor] = []
481
- self.cpu_tensors: list[torch.Tensor] = []
489
+ cpu_tensors: list[torch.Tensor] = []
490
for kv_cache_tensor in kv_caches.tensors:
491
gpu_page_size_bytes = kv_cache_tensor.page_size_bytes
492
gpu_tensor = kv_cache_tensor.tensor.view(torch.int8).view(
@@ -494,13 +502,10 @@ class CPUOffloadingWorker(OffloadingWorker):
502
(num_cpu_blocks, cpu_page_size_bytes),
503
dtype=torch.int8,
504
device="cpu",
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(),
505
+ pin_memory=pin_memory,
506
)
507
logger.debug(
503
- "torch.zeros tensor %d×%d (%.2f GB): %.3f s",
508
+ "torch.zeros pinned tensor %d×%d (%.2f GB): %.3f s",
509
num_cpu_blocks,
510
cpu_page_size_bytes,
511
num_cpu_blocks * cpu_page_size_bytes / 1e9,
@@ -508,81 +513,25 @@ class CPUOffloadingWorker(OffloadingWorker):
513
)
514
515
gpu_tensors.append(gpu_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...")
516
+ cpu_tensors.append(cpu_tensor)
517
518
self._store_handler = SingleDirectionOffloadingHandler(
519
gpu_tensors=gpu_tensors,
530
- cpu_tensors=self.cpu_tensors,
520
+ cpu_tensors=cpu_tensors,
521
block_size_factor=block_size_factor,
522
kv_cache_groups_data_refs=kv_caches.group_data_refs,
523
gpu_to_cpu=True,
524
mmap_region=mmap_region,
535
- pin_thread=self.pin_thread,
536
- manually_pinned_tensors=self._manually_pinned_tensors,
525
)
526
527
self._load_handler = SingleDirectionOffloadingHandler(
528
gpu_tensors=gpu_tensors,
541
- cpu_tensors=self.cpu_tensors,
529
+ cpu_tensors=cpu_tensors,
530
block_size_factor=block_size_factor,
531
kv_cache_groups_data_refs=kv_caches.group_data_refs,
532
gpu_to_cpu=False,
533
)
534
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
-
535
def submit_store(
536
self, job_id: int, src_spec: GPULoadStoreSpec, dst_spec: LoadStoreSpec
537
) -> bool: