perf(connector): optimize prefix load staging pipeline

daser/connector/daser_connector.py

@@ -2,7 +2,6 @@
 
 # Standard
 import asyncio
-from concurrent.futures import ThreadPoolExecutor
 import threading
 from typing import TYPE_CHECKING, Any
 
@@ -40,7 +39,7 @@
 from daser.connector.staging import (
     copy_staging_to_kv_cache as _copy_staging_to_kv_cache,
 )
-from daser.connector.worker import WorkerConnectorMixin
+from daser.connector.worker import _LOAD_REQUEST_MAX_INFLIGHT, WorkerConnectorMixin
 from daser.logging import init_logger
 
 logger = init_logger(__name__)
@@ -124,6 +123,13 @@ def __init__(
             self._transfer_ready = False
             self._transfer_mode = str(extra.get("transfer_mode", "iouring"))
             self._ipc_load_async = IPCClientAsync(self._socket_path)
+            self._ipc_load_async_pool = [
+                self._ipc_load_async,
+                *[
+                    IPCClientAsync(self._socket_path)
+                    for _ in range(max(0, _LOAD_REQUEST_MAX_INFLIGHT - 1))
+                ],
+            ]
             self._ipc_store_async = IPCClientAsync(self._socket_path)
             self._kv_caches: dict[str, torch.Tensor] = {}
             self._layer_names: list[str] = []
@@ -133,15 +139,15 @@ def __init__(
             self._pending_save_staging_bytes = 0
             self._store_staging_bytes = 0
             self._pending_store_staging_limit_bytes = 0
-            self._staging_pool = None
+            self._store_staging_pool = None
             self._pending_commits: set[str] = set()
             self._pending_finished_saves: dict[str, Any] = {}
             self._pending_loads: dict[str, Any] = {}
             self._invalid_load_block_ids: set[int] = set()
-            self._load_executor = ThreadPoolExecutor(
-                max_workers=1,
-                thread_name_prefix="daser-load",
-            )
+            self._load_request_queue = None
+            self._load_request_dispatcher = None
+            self._load_request_queue_lock = threading.Lock()
+            self._load_request_dispatcher_future = None
             self._load_loop = asyncio.new_event_loop()
             self._store_loop = asyncio.new_event_loop()
             self._load_thread = threading.Thread(

daser/connector/ipc_client.py

@@ -547,3 +547,81 @@ async def transfer_load_cuda(
                 "spans": spans,
             }
         )
+
+    async def register_load_staging_cuda(
+        self,
+        buffer_index: int,
+        cuda_ipc_handle: bytes,
+        allocation_bytes: int,
+        device_id: int,
+        device_ptr: int,
+        allocation_base_ptr: int,
+        allocation_offset: int,
+        producer_pid: int,
+    ) -> None:
+        """Register one fixed load staging CUDA allocation with the server.
+
+        Args:
+            buffer_index: Worker-local fixed staging buffer index.
+            cuda_ipc_handle: exported CUDA IPC memory handle.
+            allocation_bytes: byte size of the CUDA allocation to map.
+            device_id: CUDA device ordinal for the exported allocation.
+            device_ptr: raw device pointer for same-process server harnesses.
+            allocation_base_ptr: base pointer of the CUDA allocation owning
+                ``device_ptr``.
+            allocation_offset: byte offset of ``device_ptr`` from
+                ``allocation_base_ptr``.
+            producer_pid: process ID that exported the pointer.
+
+        Async/thread-safety:
+            Serializes with other calls on the dedicated async client
+            connection. Intended for worker initialization before hot-path
+            cache-hit loads.
+        """
+        await self.call(
+            {
+                "op": "register_load_staging",
+                "payload": {
+                    "buffer_index": int(buffer_index),
+                    "cuda_ipc_handle": cuda_ipc_handle,
+                    "allocation_bytes": int(allocation_bytes),
+                    "device_id": int(device_id),
+                    "device_ptr": int(device_ptr),
+                    "allocation_base_ptr": int(allocation_base_ptr),
+                    "allocation_offset": int(allocation_offset),
+                    "producer_pid": int(producer_pid),
+                },
+            }
+        )
+
+    async def transfer_load_registered_cuda(
+        self,
+        buffer_index: int,
+        nbytes: int,
+        spans: list[dict[str, int]],
+    ) -> dict[str, Any]:
+        """Load into a previously registered fixed CUDA staging buffer.
+
+        Args:
+            buffer_index: Worker-local fixed staging buffer index registered
+                through ``register_load_staging_cuda``.
+            nbytes: logical bytes to write for this transfer.
+            spans: byte spans containing target_offset, nbytes, and file_offset.
+
+        Returns:
+            Server response including transferred bytes and timing counters.
+
+        Async/thread-safety:
+            Runs on the worker load event loop and avoids per-load CUDA IPC
+            handle export/open payloads on the hot path.
+        """
+        return await self.call(
+            {
+                "op": "transfer_load",
+                "payload": {
+                    "load_staging_buffer_index": int(buffer_index),
+                    "nbytes": int(nbytes),
+                },
+                "spans": spans,
+            }
+        )