[Feat][Core] Add disk offloading support to SimpleCPUOffloadConnector (#49644)

Signed-off-by: Guanyi Chen <939416532@qq.com>
Co-authored-by: Claude <noreply@anthropic.com>

vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py

@@ -41,6 +41,15 @@ logger = init_logger(__name__)
 # Default CPU capacity: 8 GB
 DEFAULT_CPU_CAPACITY_BYTES = 8 * (1024**3)
 
+VALID_KV_OFFLOAD_BACKENDS = ("cpu", "disk")
+# Keys that only apply to the disk backend, warned about under "cpu".
+_DISK_ONLY_KEYS = (
+    "disk_path",
+    "disk_capacity_bytes",
+    "disk_buffer_slots",
+    "use_page_cache",
+)
+
 
 class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
     """CPU KV cache offloading with custom kernel transfers and BlockPool LRU."""
@@ -82,6 +91,36 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
 
         lazy_offload = bool(extra_config.get("lazy_offload", False))
 
+        kv_offload_backend = str(extra_config.get("kv_offload_backend", "cpu"))
+        if kv_offload_backend not in VALID_KV_OFFLOAD_BACKENDS:
+            raise ValueError(
+                f"Unknown kv_offload_backend {kv_offload_backend!r}; "
+                f"expected one of {VALID_KV_OFFLOAD_BACKENDS}"
+            )
+        disk_mode = kv_offload_backend == "disk"
+
+        disk_path = extra_config.get("disk_path", None) or None
+        disk_capacity_bytes = int(
+            extra_config.get("disk_capacity_bytes", 100 * (1024**3))
+        )
+        disk_buffer_slots = max(1, int(extra_config.get("disk_buffer_slots", 2)))
+        use_page_cache = bool(extra_config.get("use_page_cache", False))
+
+        if disk_mode:
+            if disk_path is None:
+                raise ValueError(
+                    'kv_offload_backend="disk" requires disk_path to be set.'
+                )
+        else:
+            ignored = [k for k in _DISK_ONLY_KEYS if k in extra_config]
+            if ignored:
+                logger.warning(
+                    'kv_offload_backend is "cpu", ignoring disk-only config: %s. '
+                    'Set kv_offload_backend="disk" to enable the disk backend.',
+                    ", ".join(ignored),
+                )
+            disk_path = None
+
         self.scheduler_manager: SimpleCPUOffloadScheduler | None = None
         self.worker_handler: SimpleCPUOffloadWorker | None = None
 
@@ -94,11 +133,13 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
 
         logger.info(
             "SimpleCPUOffloadConnector: role=%s, "
-            "per_rank=%.2f GB, world_size=%d, mode=%s",
+            "per_rank=%.2f GB, world_size=%d, mode=%s, backend=%s, disk=%s",
             role.name,
             cpu_capacity_per_rank / (1024**3),
             world_size,
             "lazy" if lazy_offload else "eager",
+            kv_offload_backend,
+            disk_path or "none",
         )
 
         if role == KVConnectorRole.SCHEDULER:
@@ -115,10 +156,18 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
                 scheduler_block_size=scheduler_block_size,
                 hash_block_size=hash_block_size,
                 lazy_offload=lazy_offload,
+                disk_capacity_bytes=disk_capacity_bytes if disk_mode else 0,
             )
         elif role == KVConnectorRole.WORKER:
             self.worker_handler = SimpleCPUOffloadWorker(
-                vllm_config, kv_cache_config, cpu_capacity_per_rank
+                vllm_config,
+                kv_cache_config,
+                cpu_capacity_per_rank,
+                kv_offload_backend=kv_offload_backend,
+                disk_path=disk_path,
+                disk_capacity_bytes=disk_capacity_bytes,
+                disk_buffer_slots=disk_buffer_slots,
+                use_page_cache=use_page_cache,
             )
 
     # --- Worker-side methods ---

vllm/v1/simple_kv_offload/disk_backend.py

@@ -0,0 +1,336 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Disk I/O backend for GPU<->NVMe block transfers via pinned staging buffers.
+
+Uses separate IO threads for store and load so that loads (latency-critical)
+never block behind stores (background work). Each thread owns its own pinned
+staging buffers to avoid contention.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import queue
+import threading
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.v1.simple_kv_offload.cuda_mem_ops import (
+    CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
+    CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
+    BatchMemcpyParams,
+    build_params,
+    copy_blocks,
+    pin_tensor,
+)
+
+logger = init_logger(__name__)
+
+O_DIRECT = getattr(os, "O_DIRECT", 0)
+_ALIGNMENT = 4096
+
+
+def _alloc_aligned(num_slots: int, bpb: int) -> torch.Tensor:
+    """Allocate a staging buffer whose base address is O_DIRECT aligned.
+
+    The CPU allocator only guarantees 64-byte alignment, so over-allocate by
+    one alignment unit and return an aligned view. The view keeps the backing
+    storage alive.
+    """
+    nbytes = num_slots * bpb
+    raw = torch.zeros(nbytes + _ALIGNMENT, dtype=torch.int8, device="cpu")
+    offset = -raw.data_ptr() % _ALIGNMENT
+    return raw[offset : offset + nbytes].view(num_slots, bpb)
+
+
+class DiskBackend:
+    """Async disk offload backend with pipelined GPU DMA and interleaved IO.
+
+    Architecture:
+    - Separate coordinator threads for store and load (never block each other)
+    - Interleaved pipeline: DMA slot N while preadv/pwritev slot N-1
+    - O_DIRECT by default; page cache is opt-in via use_page_cache
+
+    Same launch_copy interface as DmaCopyBackend so the worker can swap
+    backends without changing calling code.
+    """
+
+    def __init__(self) -> None:
+        self._store_params: BatchMemcpyParams | None = None
+        self._load_params: BatchMemcpyParams | None = None
+        self._load_stream: torch.cuda.Stream | None = None
+        self._store_stream: torch.cuda.Stream | None = None
+        self._store_queue: queue.SimpleQueue = queue.SimpleQueue()
+        self._load_queue: queue.SimpleQueue = queue.SimpleQueue()
+        self._store_thread: threading.Thread | None = None
+        self._load_thread: threading.Thread | None = None
+        self._shutdown: bool = False
+        self._fd: int = -1
+        self._disk_path: str = ""
+        self._total_block_bytes: int = 0
+        self._store_buffer_caches: dict[str, torch.Tensor] = {}
+        self._load_buffer_caches: dict[str, torch.Tensor] = {}
+        self._store_slot_views: list[list[memoryview]] = []
+        self._load_slot_views: list[list[memoryview]] = []
+        self._per_tensor_bpb: list[int] = []
+        self._tensor_names: list[str] = []
+
+    def init(
+        self,
+        gpu_caches: dict[str, torch.Tensor],
+        device: torch.device,
+        load_stream: torch.cuda.Stream,
+        store_stream: torch.cuda.Stream,
+        disk_path: str,
+        num_disk_slots: int,
+        total_block_bytes: int,
+        num_buffer_slots: int = 2,
+        use_page_cache: bool = False,
+    ) -> None:
+        self._load_stream = load_stream
+        self._store_stream = store_stream
+        self._total_block_bytes = total_block_bytes
+        self._num_buffer_slots = num_buffer_slots
+        self._tensor_names = list(gpu_caches.keys())
+        self._per_tensor_bpb = [
+            t.stride(0) * t.element_size() for t in gpu_caches.values()
+        ]
+
+        assert total_block_bytes % _ALIGNMENT == 0, (
+            f"total_block_bytes={total_block_bytes} not aligned to {_ALIGNMENT}"
+        )
+
+        # Separate buffer pools for store and load threads
+        self._store_buffer_caches = {}
+        self._load_buffer_caches = {}
+        for name, gpu_t in gpu_caches.items():
+            bpb = gpu_t.stride(0) * gpu_t.element_size()
+            store_buf = _alloc_aligned(num_buffer_slots, bpb)
+            pin_tensor(store_buf)
+            self._store_buffer_caches[name] = store_buf
+            load_buf = _alloc_aligned(num_buffer_slots, bpb)
+            pin_tensor(load_buf)
+            self._load_buffer_caches[name] = load_buf
+
+        # Pre-built iovec views per slot (avoid per-transfer .numpy() calls)
+        self._store_slot_views = [
+            [
+                memoryview(self._store_buffer_caches[name][slot].numpy())
+                for name in self._tensor_names
+            ]
+            for slot in range(num_buffer_slots)
+        ]
+        self._load_slot_views = [
+            [
+                memoryview(self._load_buffer_caches[name][slot].numpy())
+                for name in self._tensor_names
+            ]
+            for slot in range(num_buffer_slots)
+        ]
+
+        self._store_params = build_params(
+            gpu_caches,
+            self._store_buffer_caches,
+            store_stream,
+            src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
+        )
+        self._load_params = build_params(
+            self._load_buffer_caches,
+            gpu_caches,
+            load_stream,
+            src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
+        )
+
+        os.makedirs(os.path.dirname(disk_path) or ".", exist_ok=True)
+        # Slot contents never outlive the process, so unlink then O_EXCL rather
+        # than reopening: a pre-existing file would otherwise keep its own
+        # (possibly world-readable) mode, and blocks may encode user prompts.
+        with contextlib.suppress(FileNotFoundError):
+            os.unlink(disk_path)
+        # O_DIRECT by default: page cache would consume the very host DRAM this
+        # backend exists to conserve, and doubles the copy on the store path.
+        flags = os.O_RDWR | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW
+        if not use_page_cache:
+            flags |= O_DIRECT
+        self._fd = os.open(disk_path, flags, 0o600)
+        self._disk_path = disk_path
+        os.ftruncate(self._fd, num_disk_slots * total_block_bytes)
+
+        logger.info(
+            "DiskBackend: path=%s, slots=%d, total=%.2f GB, buf=%dx%d bytes"
+            " (page_cache=%s)",
+            disk_path,
+            num_disk_slots,
+            (num_disk_slots * total_block_bytes) / (1024**3),
+            num_buffer_slots,
+            total_block_bytes,
+            use_page_cache,
+        )
+
+        self._store_thread = threading.Thread(
+            target=self._store_loop,
+            args=(device, store_stream),
+            daemon=True,
+        )
+        self._load_thread = threading.Thread(
+            target=self._load_loop,
+            args=(device, load_stream),
+            daemon=True,
+        )
+        self._store_thread.start()
+        self._load_thread.start()
+
+    def launch_copy(
+        self,
+        src_blocks: list[int],
+        dst_blocks: list[int],
+        is_store: bool,
+        event_idx: int,
+        events_list: list[tuple[int, torch.Event]],
+        wait_event: torch.Event | None = None,
+    ) -> None:
+        q = self._store_queue if is_store else self._load_queue
+        q.put((src_blocks, dst_blocks, event_idx, events_list, wait_event))
+
+    def shutdown(self) -> None:
+        if self._shutdown:
+            return
+        self._shutdown = True
+        self._store_queue.put(None)
+        self._load_queue.put(None)
+        if self._store_thread is not None:
+            self._store_thread.join(timeout=10.0)
+        if self._load_thread is not None:
+            self._load_thread.join(timeout=10.0)
+        if self._fd < 0:
+            return
+        # Slot contents can encode user prompts, so drop the name now rather
+        # than leaving them readable until the next run overwrites the file.
+        # Unlinking only removes the directory entry: any thread still holding
+        # the fd keeps writing to the (now anonymous) inode, which the kernel
+        # frees once the last fd goes away.
+        with contextlib.suppress(OSError):
+            os.unlink(self._disk_path)
+        # Closing under a still-running IO thread would let the fd number be
+        # reused by an unrelated open(), turning its next pwritev into a write
+        # into that file. Leaking one fd for the remaining process lifetime is
+        # the cheaper failure.
+        if any(
+            t is not None and t.is_alive()
+            for t in (self._store_thread, self._load_thread)
+        ):
+            logger.warning(
+                "IO thread still running after shutdown timeout; leaking fd %d",
+                self._fd,
+            )
+            return
+        os.close(self._fd)
+        self._fd = -1
+
+    def _store_loop(
+        self,
+        device: torch.device,
+        stream: torch.cuda.Stream,
+    ) -> None:
+        current_platform.set_device(device)
+        while True:
+            item = self._store_queue.get()
+            if item is None:
+                return
+            (src_blocks, dst_blocks, event_idx, events_list, wait_event) = item
+            if wait_event is not None:
+                stream.wait_event(wait_event)
+            self._do_store(src_blocks, dst_blocks, stream)
+            event = torch.Event()
+            event.record(stream)
+            events_list.append((event_idx, event))
+
+    def _writev_slot(self, buf_slot: int, file_offset: int) -> None:
+        written = os.pwritev(self._fd, self._store_slot_views[buf_slot], file_offset)
+        if written < self._total_block_bytes:
+            raise OSError(
+                f"Short write: expected {self._total_block_bytes} bytes, "
+                f"wrote {written}"
+            )
+
+    def _readv_slot(self, buf_slot: int, file_offset: int) -> None:
+        bytes_read = os.preadv(self._fd, self._load_slot_views[buf_slot], file_offset)
+        if bytes_read < self._total_block_bytes:
+            raise OSError(
+                f"Short read: expected {self._total_block_bytes} bytes, "
+                f"read {bytes_read}"
+            )
+
+    def _load_loop(
+        self,
+        device: torch.device,
+        stream: torch.cuda.Stream,
+    ) -> None:
+        current_platform.set_device(device)
+        while True:
+            item = self._load_queue.get()
+            if item is None:
+                return
+            (src_blocks, dst_blocks, event_idx, events_list, wait_event) = item
+            if wait_event is not None:
+                stream.wait_event(wait_event)
+            self._do_load(src_blocks, dst_blocks, stream)
+            event = torch.Event()
+            event.record(stream)
+            events_list.append((event_idx, event))
+
+    def _do_store(
+        self,
+        gpu_blocks: list[int],
+        disk_slots: list[int],
+        stream: torch.cuda.Stream,
+    ) -> None:
+        """GPU -> buffer (DMA) -> disk (pwritev), interleaved double-buffer."""
+        assert self._store_params is not None
+        n = self._num_buffer_slots
+        # (DMA event, file offset) of the block already staged in each slot.
+        pending: list[tuple[torch.Event, int] | None] = [None] * n
+
+        for i, (gpu_blk, disk_slot) in enumerate(zip(gpu_blocks, disk_slots)):
+            buf_slot = i % n
+            prev = pending[buf_slot]
+            if prev is not None:
+                prev[0].synchronize()
+                self._writev_slot(buf_slot, prev[1])
+
+            copy_blocks([gpu_blk], [buf_slot], self._store_params)
+            ev = torch.Event()
+            ev.record(stream)
+            pending[buf_slot] = (ev, disk_slot * self._total_block_bytes)
+
+        for slot, last in enumerate(pending):
+            if last is not None:
+                last[0].synchronize()
+                self._writev_slot(slot, last[1])
+
+    def _do_load(
+        self,
+        disk_slots: list[int],
+        gpu_blocks: list[int],
+        stream: torch.cuda.Stream,
+    ) -> None:
+        """Disk (preadv) -> buffer -> GPU (DMA), interleaved double-buffer."""
+        assert self._load_params is not None
+        n = self._num_buffer_slots
+        prev_dma_events: list[torch.Event | None] = [None] * n
+
+        for i, (disk_slot, gpu_blk) in enumerate(zip(disk_slots, gpu_blocks)):
+            buf_slot = i % n
+            prev = prev_dma_events[buf_slot]
+            if prev is not None:
+                prev.synchronize()
+
+            self._readv_slot(buf_slot, disk_slot * self._total_block_bytes)
+
+            copy_blocks([buf_slot], [gpu_blk], self._load_params)
+            ev = torch.Event()
+            ev.record(stream)
+            prev_dma_events[buf_slot] = ev

vllm/v1/simple_kv_offload/manager.py

@@ -75,9 +75,14 @@ class SimpleCPUOffloadScheduler:
         scheduler_block_size: int,
         hash_block_size: int,
         lazy_offload: bool = False,
+        disk_capacity_bytes: int = 0,
     ):
         self.vllm_config = vllm_config
         self.kv_cache_config = kv_cache_config
+        # When disk mode is active, the offload pool size is disk-based.
+        offload_capacity = (
+            disk_capacity_bytes if disk_capacity_bytes > 0 else cpu_capacity_bytes
+        )
         self.enable_kv_cache_events = (
             vllm_config.kv_events_config is not None
             and vllm_config.kv_events_config.enable_kv_cache_events
@@ -90,7 +95,7 @@ class SimpleCPUOffloadScheduler:
         # Derive a CPU KVCacheConfig from the GPU config and build a coordinator
         assert kv_cache_config is not None
         self.cpu_kv_cache_config = self._derive_cpu_config(
-            kv_cache_config, cpu_capacity_bytes
+            kv_cache_config, offload_capacity
         )
         self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks
         # Find the full attention kv group for prefix cache matching.
@@ -111,10 +116,12 @@ class SimpleCPUOffloadScheduler:
         assert self.block_size % self.fa_block_size == 0
 
         logger.info(
-            "SimpleCPUOffloadScheduler: Allocating %d CPU blocks (%.2f GB, mode=%s)",
+            "SimpleCPUOffloadScheduler: Allocating %d offload blocks "
+            "(%.2f GB, mode=%s, backend=%s)",
             self.num_cpu_blocks,
-            cpu_capacity_bytes / (1024**3),
+            offload_capacity / (1024**3),
             "lazy" if lazy_offload else "eager",
+            "disk" if disk_capacity_bytes > 0 else "cpu",
         )
 
         spec_config = vllm_config.speculative_config

vllm/v1/simple_kv_offload/worker.py

@@ -11,6 +11,7 @@ from vllm.logger import init_logger
 from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend
 from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor
+from vllm.v1.simple_kv_offload.disk_backend import DiskBackend
 from vllm.v1.simple_kv_offload.metadata import (
     SimpleCPUOffloadMetadata,
     SimpleCPUOffloadWorkerMetadata,
@@ -30,10 +31,20 @@ class SimpleCPUOffloadWorker:
         vllm_config: VllmConfig,
         kv_cache_config: "KVCacheConfig | None",
         cpu_capacity_bytes: int,
+        kv_offload_backend: str = "cpu",
+        disk_path: str | None = None,
+        disk_capacity_bytes: int = 0,
+        disk_buffer_slots: int = 2,
+        use_page_cache: bool = False,
     ):
         self.vllm_config = vllm_config
         self.kv_cache_config = kv_cache_config
         self.cpu_capacity_bytes = cpu_capacity_bytes
+        self.disk_path = disk_path
+        self.disk_capacity_bytes = disk_capacity_bytes
+        self.disk_buffer_slots = disk_buffer_slots
+        self.use_page_cache = use_page_cache
+        self.disk_mode = kv_offload_backend == "disk"
 
         self.gpu_kv_caches: dict[str, torch.Tensor] | None = None
         self.cpu_kv_caches: dict[str, torch.Tensor] | None = None
@@ -44,7 +55,7 @@ class SimpleCPUOffloadWorker:
         self.load_stream: torch.cuda.Stream | None = None
         self.store_stream: torch.cuda.Stream | None = None
 
-        self._backend = DmaCopyBackend()
+        self._backend: DmaCopyBackend | DiskBackend | None = None
 
         # Ordered (event_idx, Event). Events pre-allocated on main thread.
         self._load_events: list[tuple[int, torch.Event]] = []
@@ -145,9 +156,57 @@ class SimpleCPUOffloadWorker:
 
         self.num_cpu_blocks = max(1, self.cpu_capacity_bytes // total_bytes_per_block)
 
+        # Use lowest priority so KV cache I/O yields to compute streams.
+        low_pri, _ = torch.cuda.Stream.priority_range()
+        self.load_stream = torch.cuda.Stream(priority=low_pri)
+        self.store_stream = torch.cuda.Stream(priority=low_pri)
+
+        self.gpu_kv_caches = unique_gpu_caches
+
+        if self.disk_mode:
+            self._init_disk_mode(unique_gpu_caches, total_bytes_per_block, self.device)
+        else:
+            self._init_cpu_mode(unique_gpu_caches, total_bytes_per_block, self.device)
+
+    def _init_disk_mode(
+        self,
+        unique_gpu_caches: dict[str, torch.Tensor],
+        total_bytes_per_block: int,
+        device: torch.device,
+    ) -> None:
+        num_disk_slots = max(1, self.disk_capacity_bytes // total_bytes_per_block)
+        self.num_cpu_blocks = num_disk_slots
+
         logger.info(
-            "SimpleCPUOffloadWorker: %d unique GPU KV tensors, "
-            "allocating %d CPU blocks (%.2f GB)",
+            "SimpleCPUOffloadWorker [DISK]: %d tensors, %d disk slots (%.2f GB)",
+            len(unique_gpu_caches),
+            num_disk_slots,
+            (num_disk_slots * total_bytes_per_block) / (1024**3),
+        )
+
+        assert self.disk_path is not None
+        rank_path = f"{self.disk_path}.rank_{device.index or 0}"
+        self._backend = DiskBackend()
+        self._backend.init(
+            unique_gpu_caches,
+            device,
+            self.load_stream,
+            self.store_stream,
+            rank_path,
+            num_disk_slots,
+            total_bytes_per_block,
+            self.disk_buffer_slots,
+            self.use_page_cache,
+        )
+
+    def _init_cpu_mode(
+        self,
+        unique_gpu_caches: dict[str, torch.Tensor],
+        total_bytes_per_block: int,
+        device: torch.device,
+    ) -> None:
+        logger.info(
+            "SimpleCPUOffloadWorker [CPU]: %d tensors, %d CPU blocks (%.2f GB)",
             len(unique_gpu_caches),
             self.num_cpu_blocks,
             (self.num_cpu_blocks * total_bytes_per_block) / (1024**3),
@@ -159,7 +218,6 @@ class SimpleCPUOffloadWorker:
                 "Pinned memory not available. CPU offload performance may be degraded."
             )
 
-        self.gpu_kv_caches = unique_gpu_caches
         self.cpu_kv_caches = {}
         for name, gpu_tensor in unique_gpu_caches.items():
             cpu_shape = (self.num_cpu_blocks,) + gpu_tensor.shape[1:]
@@ -171,16 +229,11 @@ class SimpleCPUOffloadWorker:
                 pin_tensor(tensor)
             self.cpu_kv_caches[name] = tensor
 
-        # Use lowest priority so KV cache I/O yields to compute streams.
-        low_pri, _ = torch.cuda.Stream.priority_range()
-        self.load_stream = torch.cuda.Stream(priority=low_pri)
-        self.store_stream = torch.cuda.Stream(priority=low_pri)
-
-        # Initialize copy backend with caches and streams.
+        self._backend = DmaCopyBackend()
         self._backend.init(
-            self.gpu_kv_caches,
+            unique_gpu_caches,
             self.cpu_kv_caches,
-            self.device,
+            device,
             self.load_stream,
             self.store_stream,
         )
@@ -224,8 +277,10 @@ class SimpleCPUOffloadWorker:
         # (1) Submit transfers
         metadata = self._connector_metadata
         if metadata is not None:
+            backend = self._backend
+            assert backend is not None
             if metadata.load_cpu_blocks:
-                self._backend.launch_copy(
+                backend.launch_copy(
                     metadata.load_cpu_blocks,
                     metadata.load_gpu_blocks,
                     is_store=False,
@@ -236,7 +291,7 @@ class SimpleCPUOffloadWorker:
                 if self._store_compute_done is None:
                     self._store_compute_done = torch.Event()
                 self._store_compute_done.record(torch.cuda.current_stream())
-                self._backend.launch_copy(
+                backend.launch_copy(
                     metadata.store_gpu_blocks,
                     metadata.store_cpu_blocks,
                     is_store=True,