From b38a02e0fccbe8eeebbec7d22164f29d1a9b8db6 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Thu, 25 Jun 2026 22:10:45 +0800 Subject: [PATCH 01/11] perf(connector): add fixed load staging pool - Add a fixed-size staging pool for load buffers. - Cover reuse and oversized-request behavior with focused tests. --- daser/connector/staging.py | 78 ++++++++++++++++++++++++++++ tests/connector/test_gds_transfer.py | 40 ++++++++++++++ 2 files changed, 118 insertions(+) diff --git a/daser/connector/staging.py b/daser/connector/staging.py index cf664bd..c5217e0 100644 --- a/daser/connector/staging.py +++ b/daser/connector/staging.py @@ -166,6 +166,84 @@ def release(self, lease: CudaStagingLease) -> None: self._free.append(lease.tensor) +class FixedCudaStagingPool: + """Fixed-size worker-side CUDA staging buffers for load pipelining. + + Args: + device: Device on which staging tensors are allocated. + buffer_bytes: Size of each fixed staging buffer. + depth: Number of fixed buffers to allocate. + + Async/thread-safety: + The pool is owned by one vLLM worker load thread. It never allocates + after construction, so hot-path cache-hit loads cannot trigger + unexpected CUDA OOM from staging growth. + """ + + def __init__( + self, + device: torch.device, + buffer_bytes: int, + depth: int, + ) -> None: + if buffer_bytes <= 0: + raise ValueError("buffer_bytes must be positive") + if depth <= 0: + raise ValueError("depth must be positive") + self._device = device + self._buffer_bytes = buffer_bytes + self._free: list[torch.Tensor] = [ + torch.empty(buffer_bytes, dtype=torch.uint8, device=device) + for _ in range(depth) + ] + + @property + def buffer_bytes(self) -> int: + """Return the fixed capacity of each staging buffer.""" + return self._buffer_bytes + + @property + def available(self) -> int: + """Return the number of currently free staging buffers.""" + return len(self._free) + + def acquire(self, nbytes: int) -> CudaStagingLease: + """Lease one preallocated staging buffer. + + Args: + nbytes: Logical transfer byte count. + + Returns: + Lease whose view is limited to ``nbytes``. + + Raises: + ValueError: If ``nbytes`` exceeds the fixed buffer size. + RuntimeError: If all fixed buffers are currently in use. + """ + if nbytes < 0: + raise ValueError("nbytes must be non-negative") + if nbytes > self._buffer_bytes: + raise ValueError( + f"staging request {nbytes} exceeds fixed staging buffer " + f"{self._buffer_bytes}" + ) + if not self._free: + raise RuntimeError("no fixed staging buffers available") + return CudaStagingLease( + pool=self, + tensor=self._free.pop(0), + nbytes=nbytes, + ) + + def release(self, lease: CudaStagingLease) -> None: + """Return a fixed staging lease to the free list. + + Args: + lease: Lease previously returned by ``acquire``. + """ + self._free.append(lease.tensor) + + @dataclass(frozen=True) class StagedStoreBatch: """Worker-owned CUDA staging batch ready for server-side transfer. diff --git a/tests/connector/test_gds_transfer.py b/tests/connector/test_gds_transfer.py index 4834fbf..4eb9324 100644 --- a/tests/connector/test_gds_transfer.py +++ b/tests/connector/test_gds_transfer.py @@ -84,3 +84,43 @@ def test_missing_file_raises(tmp_path): """FileNotFoundError raised when store file does not exist.""" with pytest.raises(FileNotFoundError): GDSTransferLayer(str(tmp_path / "nonexistent.store")) + + +def test_fixed_staging_pool_reuses_two_preallocated_buffers() -> None: + """Fixed load staging pool reuses bounded preallocated buffers.""" + # First Party + from daser.connector.staging import FixedCudaStagingPool + + pool = FixedCudaStagingPool( + device=torch.device("cpu"), + buffer_bytes=16, + depth=2, + ) + + first = pool.acquire(8) + second = pool.acquire(16) + + with pytest.raises(RuntimeError, match="no fixed staging buffers available"): + pool.acquire(1) + + first.release() + third = pool.acquire(4) + + assert third.tensor.data_ptr() == first.tensor.data_ptr() + second.release() + third.release() + + +def test_fixed_staging_pool_rejects_oversized_request() -> None: + """Fixed load staging pool rejects requests larger than one buffer.""" + # First Party + from daser.connector.staging import FixedCudaStagingPool + + pool = FixedCudaStagingPool( + device=torch.device("cpu"), + buffer_bytes=16, + depth=2, + ) + + with pytest.raises(ValueError, match="exceeds fixed staging buffer"): + pool.acquire(17) From ab4306d8b227fe66045b6ec5f0e2e0a4e1e93943 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Thu, 25 Jun 2026 22:12:00 +0800 Subject: [PATCH 02/11] perf(connector): add depth-two load pipeline helper - Add FIFO helper for fixed-depth load batch submission. - Cover depth-two submission and buffer reuse order. --- daser/connector/worker.py | 48 ++++++++++++++++++++++++++++ tests/connector/test_gds_transfer.py | 34 ++++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/daser/connector/worker.py b/daser/connector/worker.py index 6defaca..d9448ad 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -4,6 +4,8 @@ # Standard import asyncio +from collections import deque +from collections.abc import Callable, Sequence from dataclasses import dataclass import os import time @@ -81,6 +83,52 @@ _ROPE_WARMUP_BLOCKS = 1 +def _run_depth_two_pipeline( + batches: Sequence[Any], + submit: Callable[[Any, int], Any], + consume: Callable[[Any], int], + buffer_indices: Sequence[int], +) -> int: + """Run a fixed-depth FIFO submit/consume pipeline. + + Args: + batches: Batch payloads to submit. + submit: Function called with ``(batch, buffer_index)``. It returns an + opaque in-flight state. + consume: Function called with an in-flight state. It must return the + reusable buffer index. + buffer_indices: Fixed buffer identifiers available to the pipeline. + + Returns: + Maximum number of in-flight batches observed. + + Async/thread-safety: + Pure synchronous helper. Callers own any synchronization required by + the state returned from ``submit`` before ``consume`` releases a buffer. + """ + if not buffer_indices: + raise ValueError("buffer_indices must not be empty") + next_batch = 0 + max_inflight = 0 + in_flight: deque[Any] = deque() + + while next_batch < len(batches) and len(in_flight) < len(buffer_indices): + buffer_index = buffer_indices[len(in_flight)] + in_flight.append(submit(batches[next_batch], buffer_index)) + next_batch += 1 + max_inflight = max(max_inflight, len(in_flight)) + + while in_flight: + state = in_flight.popleft() + reusable_buffer = consume(state) + if next_batch < len(batches): + in_flight.append(submit(batches[next_batch], reusable_buffer)) + next_batch += 1 + max_inflight = max(max_inflight, len(in_flight)) + + return max_inflight + + @dataclass class _DeferredFinishedSave: """Store work held until vLLM reports a request as finished.""" diff --git a/tests/connector/test_gds_transfer.py b/tests/connector/test_gds_transfer.py index 4eb9324..1a714ce 100644 --- a/tests/connector/test_gds_transfer.py +++ b/tests/connector/test_gds_transfer.py @@ -124,3 +124,37 @@ def test_fixed_staging_pool_rejects_oversized_request() -> None: with pytest.raises(ValueError, match="exceeds fixed staging buffer"): pool.acquire(17) + + +def test_depth_two_pipeline_submits_next_batch_after_oldest_consumes() -> None: + """Depth-two load pipeline reuses the oldest consumed buffer first.""" + # First Party + from daser.connector.worker import _run_depth_two_pipeline + + events: list[str] = [] + + def submit(batch: str, buffer_index: int) -> tuple[str, int]: + events.append(f"submit:{batch}:buf{buffer_index}") + return batch, buffer_index + + def consume(state: tuple[str, int]) -> int: + batch, buffer_index = state + events.append(f"consume:{batch}:buf{buffer_index}") + return buffer_index + + max_inflight = _run_depth_two_pipeline( + batches=["b0", "b1", "b2"], + submit=submit, + consume=consume, + buffer_indices=[0, 1], + ) + + assert max_inflight == 2 + assert events == [ + "submit:b0:buf0", + "submit:b1:buf1", + "consume:b0:buf0", + "submit:b2:buf0", + "consume:b1:buf1", + "consume:b2:buf0", + ] From 0a1654cf047a710ed5053d8a8ef489c967c71c03 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Thu, 25 Jun 2026 22:18:23 +0800 Subject: [PATCH 03/11] perf(connector): pipeline async load batches - Use two fixed load staging buffers for worker cache-hit loads. - Overlap server load submission with prior batch restore. - Preserve FIFO restore ordering and final CUDA synchronization. --- daser/connector/staging.py | 46 +++- daser/connector/worker.py | 332 +++++++++++++++++++----- tests/connector/test_daser_connector.py | 76 ++++++ 3 files changed, 381 insertions(+), 73 deletions(-) diff --git a/daser/connector/staging.py b/daser/connector/staging.py index c5217e0..b549e3d 100644 --- a/daser/connector/staging.py +++ b/daser/connector/staging.py @@ -192,10 +192,11 @@ def __init__( raise ValueError("depth must be positive") self._device = device self._buffer_bytes = buffer_bytes - self._free: list[torch.Tensor] = [ + self._buffers: list[torch.Tensor] = [ torch.empty(buffer_bytes, dtype=torch.uint8, device=device) for _ in range(depth) ] + self._free_indices: list[int] = list(range(depth)) @property def buffer_bytes(self) -> int: @@ -205,7 +206,7 @@ def buffer_bytes(self) -> int: @property def available(self) -> int: """Return the number of currently free staging buffers.""" - return len(self._free) + return len(self._free_indices) def acquire(self, nbytes: int) -> CudaStagingLease: """Lease one preallocated staging buffer. @@ -227,11 +228,40 @@ def acquire(self, nbytes: int) -> CudaStagingLease: f"staging request {nbytes} exceeds fixed staging buffer " f"{self._buffer_bytes}" ) - if not self._free: + if not self._free_indices: raise RuntimeError("no fixed staging buffers available") + return self.acquire_index(self._free_indices[0], nbytes) + + def acquire_index(self, index: int, nbytes: int) -> CudaStagingLease: + """Lease a specific preallocated staging buffer. + + Args: + index: Fixed buffer index to lease. + nbytes: Logical transfer byte count. + + Returns: + Lease whose view is limited to ``nbytes``. + + Raises: + ValueError: If ``index`` is invalid or ``nbytes`` exceeds the fixed + buffer size. + RuntimeError: If the requested buffer is currently in use. + """ + if index < 0 or index >= len(self._buffers): + raise ValueError(f"fixed staging buffer index out of range: {index}") + if nbytes < 0: + raise ValueError("nbytes must be non-negative") + if nbytes > self._buffer_bytes: + raise ValueError( + f"staging request {nbytes} exceeds fixed staging buffer " + f"{self._buffer_bytes}" + ) + if index not in self._free_indices: + raise RuntimeError(f"fixed staging buffer {index} is not available") + self._free_indices.remove(index) return CudaStagingLease( pool=self, - tensor=self._free.pop(0), + tensor=self._buffers[index], nbytes=nbytes, ) @@ -241,7 +271,13 @@ def release(self, lease: CudaStagingLease) -> None: Args: lease: Lease previously returned by ``acquire``. """ - self._free.append(lease.tensor) + for index, tensor in enumerate(self._buffers): + if tensor is lease.tensor: + if index not in self._free_indices: + self._free_indices.append(index) + self._free_indices.sort() + return + raise ValueError("lease does not belong to this fixed staging pool") @dataclass(frozen=True) diff --git a/daser/connector/worker.py b/daser/connector/worker.py index d9448ad..b1dd917 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -35,6 +35,7 @@ FUSED_RESTORE_MIN_SLOTS, CudaStagingLease, CudaStagingPool, + FixedCudaStagingPool, StagedStoreBatch, ) from daser.connector.staging import ( @@ -81,6 +82,8 @@ logger = init_logger(__name__) _ROPE_WARMUP_BLOCKS = 1 +_LOAD_PIPELINE_DEPTH = 2 +_LoadBatch = tuple[int, list[dict[str, int]], list[tuple[int, int, Any]]] def _run_depth_two_pipeline( @@ -181,6 +184,46 @@ def release(self) -> None: self.lease = None +@dataclass +class _InflightLoadBatch: + """One submitted load batch and its fixed staging lease. + + Attributes: + buffer_index: Fixed load staging buffer index. + total_bytes: Logical bytes in this transfer batch. + per_req_ranges: Ranges used to restore staging bytes into vLLM KV cache. + staging_lease: Fixed staging lease retained until restore completes. + future: Future returned by the load event loop. + submitted_at: Wall-clock timestamp used for wait timing. + """ + + buffer_index: int + total_bytes: int + per_req_ranges: list[tuple[int, int, Any]] + staging_lease: CudaStagingLease + future: Any + submitted_at: float + + +@dataclass +class _ConsumedLoadBatch: + """Timing and accounting data for one restored load batch.""" + + buffer_index: int + bytes: int + copies: int + copy_runs: int + ipc_ms: float + wait_ms: float + copy_ms: float + transfer_open_ms: float + transfer_load_ms: float + transfer_sync_ms: float + l1_hits: int + l1_misses: int + l2_reads: int + + class _ImmediateLoadError: """Completed future used when a load cannot be submitted. @@ -407,11 +450,18 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: initial_bytes=self._store_staging_bytes, max_buffer_bytes=self._store_staging_bytes, ) + self._load_staging_pool = FixedCudaStagingPool( + device=sample.device, + buffer_bytes=self._store_staging_bytes, + depth=_LOAD_PIPELINE_DEPTH, + ) logger.info( - "[CONNECTOR] preallocated staging buffer=%d cap=%d pending=%d", + "[CONNECTOR] preallocated staging buffer=%d cap=%d pending=%d " + "load_pipeline_depth=%d", self._store_staging_bytes, self._store_staging_bytes, self._pending_store_staging_limit_bytes, + _LOAD_PIPELINE_DEPTH, ) if sample.dim() >= 5: _warm_rope_apply_backends( @@ -481,11 +531,18 @@ def register_cross_layers_kv_cache( initial_bytes=self._store_staging_bytes, max_buffer_bytes=self._store_staging_bytes, ) + self._load_staging_pool = FixedCudaStagingPool( + device=kv_cache.device, + buffer_bytes=self._store_staging_bytes, + depth=_LOAD_PIPELINE_DEPTH, + ) logger.info( - "[CONNECTOR] register_cross_layers_kv_cache: layers=%d shape=%s dtype=%s", + "[CONNECTOR] register_cross_layers_kv_cache: layers=%d shape=%s " + "dtype=%s load_pipeline_depth=%d", len(self._layer_names), tuple(kv_cache.shape), kv_cache.dtype, + _LOAD_PIPELINE_DEPTH, ) _warm_rope_apply_backends( device=kv_cache.device, @@ -629,10 +686,11 @@ def _load_kv_specs( RPCs to the dedicated load asyncio loop and waits inside the background thread, not on vLLM's worker thread. """ + load_staging_pool = self._ensure_load_staging_pool(sample_tensor) load_batches = _build_load_read_batches( reqs_to_load, self._slot_size, - max_batch_bytes=(self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES), + max_batch_bytes=load_staging_pool.buffer_bytes, ) if not load_batches: return @@ -649,76 +707,48 @@ def _load_kv_specs( total_transfer_open_ms = 0.0 total_transfer_load_ms = 0.0 total_transfer_sync_ms = 0.0 - for total_bytes, spans, per_req_ranges in load_batches: - total_bytes_loaded += total_bytes - staging_lease = self._acquire_staging(total_bytes, sample_tensor.device) - staging = staging_lease.view - try: - cp_staging = cupy.asarray(staging) - cuda_handle = export_cuda_ipc_handle(cp_staging) - device_id = cuda_array_device_id(cp_staging) - device_ptr = cuda_array_pointer(cp_staging) - ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) - - ipc_start = time.perf_counter() - load_response = self._submit_load_coroutine( - self._transfer_load_cuda( - cuda_ipc_handle=cuda_handle, - nbytes=total_bytes, - device_id=device_id, - device_ptr=device_ptr, - allocation_base_ptr=ipc_base_ptr, - allocation_offset=ipc_offset, - producer_pid=os.getpid(), - spans=spans, - ) - ).result(timeout=120.0) - total_ipc_ms += (time.perf_counter() - ipc_start) * 1000 - total_transfer_open_ms += float( - load_response.get("transfer_open_ms", 0.0) - ) - total_transfer_load_ms += float( - load_response.get("transfer_load_ms", 0.0) - ) - total_transfer_sync_ms += float( - load_response.get("transfer_sync_ms", 0.0) - ) - stats = load_response.get("transfer_stats_delta", {}) - if isinstance(stats, dict): - total_l1_hits += int(stats.get("l1_hits", 0)) - total_l1_misses += int(stats.get("l1_misses", 0)) - total_l2_reads += int(stats.get("l2_reads", 0)) - - copy_runs = _build_load_copy_runs(per_req_ranges) - total_copy_runs += len(copy_runs) - copy_start = time.perf_counter() - for run in copy_runs: - total_copies += _copy_staging_to_kv_cache( - staging=staging[run.start : run.end], - kv_caches=self._kv_caches, - layer_names=self._layer_names, - block_ids=run.block_ids, - slot_size=self._slot_size, - load_key_scale=self._load_key_scale, - load_value_scale=self._load_value_scale, - pos_offset=run.pos_offset, - rope_delta_scale=self._rope_delta_scale, - rope_base=self._rope_base, - rope_rotary_dim=self._rope_rotary_dim, - rope_is_neox_style=self._rope_is_neox_style, - ) - total_copy_ms += (time.perf_counter() - copy_start) * 1000 - sync_start = time.perf_counter() - _synchronize_cuda_tensor(sample_tensor) - total_sync_ms += (time.perf_counter() - sync_start) * 1000 - finally: - staging_lease.release() + consumed_batches: list[_ConsumedLoadBatch] = [] + + def consume_load_batch(state: _InflightLoadBatch) -> int: + consumed = self._consume_loaded_batch(state, sample_tensor) + consumed_batches.append(consumed) + return consumed.buffer_index + + max_inflight_batches = _run_depth_two_pipeline( + batches=load_batches, + submit=lambda batch, buffer_index: self._submit_load_batch( + batch, + buffer_index, + sample_tensor, + ), + consume=consume_load_batch, + buffer_indices=range(_LOAD_PIPELINE_DEPTH), + ) + + for consumed in consumed_batches: + total_bytes_loaded += consumed.bytes + total_copies += consumed.copies + total_copy_runs += consumed.copy_runs + total_ipc_ms += consumed.ipc_ms + total_copy_ms += consumed.copy_ms + total_transfer_open_ms += consumed.transfer_open_ms + total_transfer_load_ms += consumed.transfer_load_ms + total_transfer_sync_ms += consumed.transfer_sync_ms + total_l1_hits += consumed.l1_hits + total_l1_misses += consumed.l1_misses + total_l2_reads += consumed.l2_reads + + sync_start = time.perf_counter() + _synchronize_cuda_tensor(sample_tensor) + total_sync_ms += (time.perf_counter() - sync_start) * 1000 logger.debug( "[CONNECTOR] start_load_kv timing: reqs=%d batches=%d bytes=%d " "copy_runs=%d gpu_copies=%d ipc_ms=%.3f copy_ms=%.3f " "sync_ms=%.3f transfer_open_ms=%.3f transfer_load_ms=%.3f " - "transfer_sync_ms=%.3f l1_hits=%d l1_misses=%d l2_reads=%d", + "transfer_sync_ms=%.3f l1_hits=%d l1_misses=%d l2_reads=%d " + "load_pipeline_depth=%d max_inflight_batches=%d " + "load_buffer_bytes=%d dynamic_load_allocations=%d", len(reqs_to_load), len(load_batches), total_bytes_loaded, @@ -733,8 +763,174 @@ def _load_kv_specs( total_l1_hits, total_l1_misses, total_l2_reads, + _LOAD_PIPELINE_DEPTH, + max_inflight_batches, + load_staging_pool.buffer_bytes, + 0, ) + def _ensure_load_staging_pool( + self, + sample_tensor: torch.Tensor, + ) -> FixedCudaStagingPool: + """Return the fixed load staging pool, creating it before traffic if needed. + + Args: + sample_tensor: Representative KV cache tensor used for device + placement when the pool was not initialized during registration. + + Returns: + Fixed load staging pool with two preallocated buffers. + + Async/thread-safety: + Called from the worker load executor. Normal production flow creates + the pool during KV-cache registration; this fallback keeps tests and + deferred initialization paths explicit while still allocating only + once before batch submission. + """ + pool = getattr(self, "_load_staging_pool", None) + if pool is not None: + return pool + buffer_bytes = max( + self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, + self._slot_size, + ) + pool = FixedCudaStagingPool( + device=sample_tensor.device, + buffer_bytes=buffer_bytes, + depth=_LOAD_PIPELINE_DEPTH, + ) + self._load_staging_pool = pool + return pool + + def _submit_load_batch( + self, + batch: _LoadBatch, + buffer_index: int, + sample_tensor: torch.Tensor, + ) -> _InflightLoadBatch: + """Submit one load batch into a fixed staging buffer. + + Args: + batch: Tuple from ``build_load_read_batches``. + buffer_index: Fixed load staging buffer to use. + sample_tensor: Representative KV cache tensor for device context. + + Returns: + In-flight batch state consumed by ``_consume_loaded_batch``. + + Async/thread-safety: + Runs on the connector load executor. The returned state owns the + fixed staging lease until consumption finishes. + """ + del sample_tensor + total_bytes, spans, per_req_ranges = batch + pool = self._ensure_load_staging_pool(next(iter(self._kv_caches.values()))) + staging_lease = pool.acquire_index(buffer_index, total_bytes) + staging = staging_lease.view + cp_staging = cupy.asarray(staging) + cuda_handle = export_cuda_ipc_handle(cp_staging) + device_id = cuda_array_device_id(cp_staging) + device_ptr = cuda_array_pointer(cp_staging) + ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) + + submitted_at = time.perf_counter() + future = self._submit_load_coroutine( + self._transfer_load_cuda( + cuda_ipc_handle=cuda_handle, + nbytes=total_bytes, + device_id=device_id, + device_ptr=device_ptr, + allocation_base_ptr=ipc_base_ptr, + allocation_offset=ipc_offset, + producer_pid=os.getpid(), + spans=spans, + ) + ) + return _InflightLoadBatch( + buffer_index=buffer_index, + total_bytes=total_bytes, + per_req_ranges=per_req_ranges, + staging_lease=staging_lease, + future=future, + submitted_at=submitted_at, + ) + + def _consume_loaded_batch( + self, + state: _InflightLoadBatch, + sample_tensor: torch.Tensor, + ) -> _ConsumedLoadBatch: + """Wait for one submitted load batch and restore it into vLLM KV cache. + + Args: + state: In-flight load batch returned by ``_submit_load_batch``. + sample_tensor: Representative KV cache tensor for synchronization. + + Returns: + Timing and accounting data for the restored batch. + + Async/thread-safety: + Runs on the connector load executor. It releases the fixed staging + lease only after restore kernels that read the staging view have + been synchronized. + """ + try: + wait_start = time.perf_counter() + load_response = state.future.result(timeout=120.0) + wait_ms = (time.perf_counter() - wait_start) * 1000 + ipc_ms = (time.perf_counter() - state.submitted_at) * 1000 + transfer_open_ms = float(load_response.get("transfer_open_ms", 0.0)) + transfer_load_ms = float(load_response.get("transfer_load_ms", 0.0)) + transfer_sync_ms = float(load_response.get("transfer_sync_ms", 0.0)) + stats = load_response.get("transfer_stats_delta", {}) + l1_hits = 0 + l1_misses = 0 + l2_reads = 0 + if isinstance(stats, dict): + l1_hits = int(stats.get("l1_hits", 0)) + l1_misses = int(stats.get("l1_misses", 0)) + l2_reads = int(stats.get("l2_reads", 0)) + + copy_runs = _build_load_copy_runs(state.per_req_ranges) + copy_start = time.perf_counter() + copies = 0 + staging = state.staging_lease.view + for run in copy_runs: + copies += _copy_staging_to_kv_cache( + staging=staging[run.start : run.end], + kv_caches=self._kv_caches, + layer_names=self._layer_names, + block_ids=run.block_ids, + slot_size=self._slot_size, + load_key_scale=self._load_key_scale, + load_value_scale=self._load_value_scale, + pos_offset=run.pos_offset, + rope_delta_scale=self._rope_delta_scale, + rope_base=self._rope_base, + rope_rotary_dim=self._rope_rotary_dim, + rope_is_neox_style=self._rope_is_neox_style, + ) + copy_ms = (time.perf_counter() - copy_start) * 1000 + _synchronize_cuda_tensor(sample_tensor) + return _ConsumedLoadBatch( + buffer_index=state.buffer_index, + bytes=state.total_bytes, + copies=copies, + copy_runs=len(copy_runs), + ipc_ms=ipc_ms, + wait_ms=wait_ms, + copy_ms=copy_ms, + transfer_open_ms=transfer_open_ms, + transfer_load_ms=transfer_load_ms, + transfer_sync_ms=transfer_sync_ms, + l1_hits=l1_hits, + l1_misses=l1_misses, + l2_reads=l2_reads, + ) + finally: + state.staging_lease.release() + def wait_for_layer_load(self, layer_name: str) -> None: """No-op because async loads complete before vLLM resumes requests. diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index 79003f6..bd72b35 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -2,6 +2,7 @@ # Standard import asyncio +from types import SimpleNamespace # Third Party import cupy @@ -842,6 +843,81 @@ def fake_submit_load_coroutine(coro): assert "start_load_kv timing" not in caplog.text +def test_worker_load_pipeline_submits_two_batches_before_consuming(monkeypatch) -> None: + """Worker load path submits two fixed-buffer batches before consuming.""" + + class Probe(WorkerConnectorMixin): + def __init__(self) -> None: + self._slot_size = 4 + self._store_staging_bytes = 8 + self._kv_caches = {"layer.0": torch.empty(4, 1, 1, 4, dtype=torch.uint8)} + self._layer_names = ["layer.0"] + self._load_key_scale = 1.0 + self._load_value_scale = 1.0 + self._rope_delta_scale = 1.0 + self._rope_base = 10000.0 + self._rope_rotary_dim = 0 + self._rope_is_neox_style = True + self._load_staging_pool = None + + def load_specs(self, reqs_to_load: dict[str, ReqLoadSpec]) -> None: + """Expose load execution through a public test helper.""" + self._load_kv_specs(reqs_to_load, next(iter(self._kv_caches.values()))) + + probe = Probe() + reqs_to_load = { + "req": ReqLoadSpec( + chunk_key="hit", + start_slot=0, + num_slots=3, + block_ids=[0, 1, 2], + file_offset=0, + token_count=12, + ) + } + events: list[str] = [] + + def fake_submit(batch, buffer_index: int, sample_tensor: torch.Tensor): + del sample_tensor + total_bytes, _spans, _per_req_ranges = batch + events.append(f"submit:{total_bytes}:buf{buffer_index}") + return ("state", total_bytes, buffer_index) + + def fake_consume(state, sample_tensor: torch.Tensor): + del sample_tensor + _tag, total_bytes, buffer_index = state + events.append(f"consume:{total_bytes}:buf{buffer_index}") + return SimpleNamespace( + buffer_index=buffer_index, + bytes=total_bytes, + copies=1, + copy_runs=1, + ipc_ms=0.0, + copy_ms=0.0, + transfer_open_ms=0.0, + transfer_load_ms=0.0, + transfer_sync_ms=0.0, + l1_hits=0, + l1_misses=0, + l2_reads=0, + ) + + monkeypatch.setattr(probe, "_submit_load_batch", fake_submit) + monkeypatch.setattr(probe, "_consume_loaded_batch", fake_consume) + monkeypatch.setattr( + "daser.connector.worker._synchronize_cuda_tensor", + lambda _: None, + ) + + probe.load_specs(reqs_to_load) + + assert events[:3] == [ + "submit:8:buf0", + "submit:4:buf1", + "consume:8:buf0", + ] + + def test_start_load_kv_releases_waiting_request_when_load_cannot_start() -> None: """Worker should not leave async-hit requests waiting forever.""" connector = _WorkerProbe("") From c33ae6f8d3ba9a25dae81cd7a9a82c7b10539318 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Fri, 26 Jun 2026 16:46:38 +0800 Subject: [PATCH 04/11] perf(connector): register load staging mappings - Add IPC support for registering fixed load staging buffers by index. - Use registered staging buffer indexes on the worker load hot path with a CUDA IPC fallback. - Cover registered staging RPCs and server mapping reuse in tests. --- daser/connector/ipc_client.py | 78 +++++++++++++++ daser/connector/staging.py | 18 ++++ daser/connector/worker.py | 86 +++++++++++++++-- daser/server/ipc/server.py | 148 +++++++++++++++++++++-------- tests/connector/test_ipc_client.py | 58 +++++++++++ tests/server/test_ipc_server.py | 109 +++++++++++++++++++++ 6 files changed, 448 insertions(+), 49 deletions(-) diff --git a/daser/connector/ipc_client.py b/daser/connector/ipc_client.py index 0e514b9..15bd404 100644 --- a/daser/connector/ipc_client.py +++ b/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, + } + ) diff --git a/daser/connector/staging.py b/daser/connector/staging.py index b549e3d..8514862 100644 --- a/daser/connector/staging.py +++ b/daser/connector/staging.py @@ -208,6 +208,24 @@ def available(self) -> int: """Return the number of currently free staging buffers.""" return len(self._free_indices) + def buffer(self, index: int) -> torch.Tensor: + """Return a fixed staging backing tensor by index. + + Args: + index: Fixed buffer index to inspect. + + Returns: + The full preallocated backing tensor for the requested index. + + Async/thread-safety: + Intended for one-time CUDA IPC registration before load traffic. + Callers must not mutate the returned tensor independently of the + pool lease protocol. + """ + if index < 0 or index >= len(self._buffers): + raise ValueError(f"fixed staging buffer index out of range: {index}") + return self._buffers[index] + def acquire(self, nbytes: int) -> CudaStagingLease: """Lease one preallocated staging buffer. diff --git a/daser/connector/worker.py b/daser/connector/worker.py index b1dd917..ef334d2 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -455,6 +455,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: buffer_bytes=self._store_staging_bytes, depth=_LOAD_PIPELINE_DEPTH, ) + self._load_staging_registered = False logger.info( "[CONNECTOR] preallocated staging buffer=%d cap=%d pending=%d " "load_pipeline_depth=%d", @@ -536,6 +537,7 @@ def register_cross_layers_kv_cache( buffer_bytes=self._store_staging_bytes, depth=_LOAD_PIPELINE_DEPTH, ) + self._load_staging_registered = False logger.info( "[CONNECTOR] register_cross_layers_kv_cache: layers=%d shape=%s " "dtype=%s load_pipeline_depth=%d", @@ -828,15 +830,21 @@ def _submit_load_batch( pool = self._ensure_load_staging_pool(next(iter(self._kv_caches.values()))) staging_lease = pool.acquire_index(buffer_index, total_bytes) staging = staging_lease.view - cp_staging = cupy.asarray(staging) - cuda_handle = export_cuda_ipc_handle(cp_staging) - device_id = cuda_array_device_id(cp_staging) - device_ptr = cuda_array_pointer(cp_staging) - ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) submitted_at = time.perf_counter() - future = self._submit_load_coroutine( - self._transfer_load_cuda( + if getattr(self, "_load_staging_registered", False): + transfer_coro = self._transfer_load_registered_cuda( + buffer_index=buffer_index, + nbytes=total_bytes, + spans=spans, + ) + else: + cp_staging = cupy.asarray(staging) + cuda_handle = export_cuda_ipc_handle(cp_staging) + device_id = cuda_array_device_id(cp_staging) + device_ptr = cuda_array_pointer(cp_staging) + ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) + transfer_coro = self._transfer_load_cuda( cuda_ipc_handle=cuda_handle, nbytes=total_bytes, device_id=device_id, @@ -846,7 +854,7 @@ def _submit_load_batch( producer_pid=os.getpid(), spans=spans, ) - ) + future = self._submit_load_coroutine(transfer_coro) return _InflightLoadBatch( buffer_index=buffer_index, total_bytes=total_bytes, @@ -1186,6 +1194,52 @@ def _init_server_transfer(self) -> None: self._submit_store_coroutine(self._ipc_store_async.init_transfer()).result( timeout=120.0 ) + self._register_load_staging_buffers() + + def _register_load_staging_buffers(self) -> None: + """Register fixed load staging buffers with the server. + + Async/thread-safety: + Called after server transfer initialization and before request + traffic. Registration failures are logged and leave the worker on + the compatible per-load CUDA IPC payload path. + """ + if getattr(self, "_load_staging_registered", False): + return + pool = getattr(self, "_load_staging_pool", None) + if pool is None or getattr(self, "_ipc_load_async", None) is None: + return + try: + for buffer_index in range(_LOAD_PIPELINE_DEPTH): + tensor = pool.buffer(buffer_index) + cp_tensor = cupy.asarray(tensor) + device_ptr = cuda_array_pointer(cp_tensor) + ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) + self._submit_load_coroutine( + self._ipc_load_async.register_load_staging_cuda( + buffer_index=buffer_index, + cuda_ipc_handle=export_cuda_ipc_handle(cp_tensor), + allocation_bytes=int(tensor.numel()), + device_id=cuda_array_device_id(cp_tensor), + device_ptr=device_ptr, + allocation_base_ptr=ipc_base_ptr, + allocation_offset=ipc_offset, + producer_pid=os.getpid(), + ) + ).result(timeout=120.0) + except Exception as exc: # noqa: BLE001 + logger.warning( + "[CONNECTOR] registered load staging unavailable; falling back " + "to per-load CUDA IPC payloads: %s", + exc, + ) + self._load_staging_registered = False + return + self._load_staging_registered = True + logger.info( + "[CONNECTOR] registered %d fixed load staging buffers", + _LOAD_PIPELINE_DEPTH, + ) def _reap_save_futures(self, block: bool) -> None: """Collect completed background save tasks. @@ -1490,6 +1544,22 @@ async def _transfer_load_cuda(self, **kwargs: Any) -> dict[str, Any]: """ return await self._ipc_load_async.transfer_load_cuda(**kwargs) + async def _transfer_load_registered_cuda(self, **kwargs: Any) -> dict[str, Any]: + """Load through a pre-registered fixed CUDA staging buffer. + + Args: + **kwargs: forwarded registered-buffer transfer fields. + + Returns: + Server load response with timing counters. + + Async/thread-safety: + Runs on the worker load event loop. The server has already opened + the CUDA IPC mapping during initialization, so this hot-path call + only identifies the staging buffer index and logical byte range. + """ + return await self._ipc_load_async.transfer_load_registered_cuda(**kwargs) + async def _transfer_store_cuda(self, **kwargs: Any) -> list[str]: """Store through the dedicated worker store IPC client. diff --git a/daser/server/ipc/server.py b/daser/server/ipc/server.py index e73a6ae..95102db 100644 --- a/daser/server/ipc/server.py +++ b/daser/server/ipc/server.py @@ -126,6 +126,7 @@ def __init__( self._cuda_ipc_cache: OrderedDict[ tuple[int, int, int, int | None], "_CachedCudaArray" ] = OrderedDict() + self._load_staging_buffers: dict[int, _CachedCudaArray] = {} self._op_handlers: dict[ str, Callable[[dict[str, Any]], Awaitable[dict[str, Any]]] ] = { @@ -143,6 +144,7 @@ def __init__( "init_transfer": self._op_init_transfer, "transfer_store": self._transfer_store, "transfer_load": self._transfer_load, + "register_load_staging": self._register_load_staging, "evict_chunk": self._op_evict_chunk, "release_chunk_writer": self._op_release_chunk_writer, } @@ -484,6 +486,9 @@ async def _transfer_load(self, msg: dict[str, Any]) -> dict[str, Any]: load_ms = 0.0 sync_ms = 0.0 close_ms = 0.0 + close_one_shot_buffer = isinstance(buffer, _UncachedCudaArray) and ( + "load_staging_buffer_index" not in payload + ) try: before = asdict(transfer.stats) started = time.perf_counter() @@ -532,7 +537,7 @@ async def _transfer_load(self, msg: dict[str, Any]) -> dict[str, Any]: ) return response finally: - if isinstance(buffer, _UncachedCudaArray): + if close_one_shot_buffer: close = getattr(buffer, "close", None) close_start = time.perf_counter() close() @@ -543,6 +548,34 @@ async def _transfer_load(self, msg: dict[str, Any]) -> dict[str, Any]: close_ms, ) + async def _register_load_staging(self, msg: dict[str, Any]) -> dict[str, Any]: + """Register one fixed CUDA load staging buffer for indexed reuse. + + Args: + msg: IPC request whose payload describes a worker-owned fixed CUDA + staging allocation. + + Returns: + ``{"ok": True}`` after the CUDA IPC mapping is open and cached. + + Async/thread-safety: + Runs on the IPC event loop during worker initialization. Replacing + an existing index closes the old mapping only after the new mapping + is available. + """ + payload = msg.get("payload", {}) + buffer_index = int(payload["buffer_index"]) + opened = self._open_cuda_ipc_payload( + payload=payload, + nbytes_key="allocation_bytes", + cache_mapping=False, + ) + previous = self._load_staging_buffers.get(buffer_index) + self._load_staging_buffers[buffer_index] = opened + if previous is not None: + previous.close() + return {"ok": True} + def _record_transfer_metrics( self, op: str, @@ -669,49 +702,79 @@ def _payload_buffer(self, payload: dict[str, Any]) -> Any: """Return a byte-addressable buffer for an IPC transfer payload.""" if "data" in payload: return bytearray(payload["data"]) + if "load_staging_buffer_index" in payload: + buffer_index = int(payload["load_staging_buffer_index"]) + try: + return self._load_staging_buffers[buffer_index] + except KeyError as exc: + raise ValueError( + f"unknown load staging buffer index: {buffer_index}" + ) from exc if "cuda_ipc_handle" in payload: - local_ptr = None - producer_pid = int(payload.get("producer_pid", -1)) - device_ptr = int(payload["device_ptr"]) - nbytes = int(payload["nbytes"]) - device_id = int(payload["device_id"]) if "device_id" in payload else None - allocation_offset = int(payload.get("allocation_offset", 0)) - allocation_base_ptr = int( - payload.get("allocation_base_ptr", device_ptr - allocation_offset) + return self._open_cuda_ipc_payload(payload=payload, nbytes_key="nbytes") + raise ValueError("transfer payload requires data or cuda_ipc_handle") + + def _open_cuda_ipc_payload( + self, + *, + payload: dict[str, Any], + nbytes_key: str, + cache_mapping: bool = True, + ) -> "_CachedCudaArray": + """Open or reuse a CUDA IPC payload mapping. + + Args: + payload: CUDA IPC payload with handle, device, allocation base, and + offset fields. + nbytes_key: Payload field that identifies the mapping size. + cache_mapping: If True, reuse the general CUDA IPC LRU cache. Fixed + registered staging buffers pass False because their lifetime is + owned by ``_load_staging_buffers`` and must not be LRU-evicted. + + Returns: + Cached or one-shot CUDA array wrapper for the mapped allocation. + """ + local_ptr = None + producer_pid = int(payload.get("producer_pid", -1)) + device_ptr = int(payload["device_ptr"]) + nbytes = int(payload[nbytes_key]) + device_id = int(payload["device_id"]) if "device_id" in payload else None + allocation_offset = int(payload.get("allocation_offset", 0)) + allocation_base_ptr = int( + payload.get("allocation_base_ptr", device_ptr - allocation_offset) + ) + if producer_pid == os.getpid(): + local_ptr = allocation_base_ptr + if local_ptr is None and cache_mapping: + key = ( + producer_pid, + allocation_base_ptr, + nbytes + allocation_offset, + device_id, ) - if producer_pid == os.getpid(): - local_ptr = allocation_base_ptr - if local_ptr is None: - key = ( - producer_pid, - allocation_base_ptr, - nbytes + allocation_offset, - device_id, + cached = self._cuda_ipc_cache.get(key) + if cached is None: + self._evict_cuda_ipc_cache_if_needed() + opened = open_cuda_ipc_buffer( + handle=payload["cuda_ipc_handle"], + nbytes=nbytes, + device_id=device_id, + local_ptr=None, + allocation_offset=allocation_offset, ) - cached = self._cuda_ipc_cache.get(key) - if cached is None: - self._evict_cuda_ipc_cache_if_needed() - opened = open_cuda_ipc_buffer( - handle=payload["cuda_ipc_handle"], - nbytes=nbytes, - device_id=device_id, - local_ptr=None, - allocation_offset=allocation_offset, - ) - cached = _CachedCudaArray(opened) - self._cuda_ipc_cache[key] = cached - else: - self._cuda_ipc_cache.move_to_end(key) - return cached - opened = open_cuda_ipc_buffer( - handle=payload["cuda_ipc_handle"], - nbytes=nbytes, - device_id=device_id, - local_ptr=local_ptr, - allocation_offset=allocation_offset, - ) - return _UncachedCudaArray(opened) - raise ValueError("transfer payload requires data or cuda_ipc_handle") + cached = _CachedCudaArray(opened) + self._cuda_ipc_cache[key] = cached + else: + self._cuda_ipc_cache.move_to_end(key) + return cached + opened = open_cuda_ipc_buffer( + handle=payload["cuda_ipc_handle"], + nbytes=nbytes, + device_id=device_id, + local_ptr=local_ptr, + allocation_offset=allocation_offset, + ) + return _UncachedCudaArray(opened) def _evict_cuda_ipc_cache_if_needed(self) -> None: """Evict one cached CUDA IPC mapping when the cache is full.""" @@ -722,6 +785,9 @@ def _evict_cuda_ipc_cache_if_needed(self) -> None: def _close_cuda_ipc_cache(self) -> None: """Close all cached CUDA IPC mappings.""" + for cached in self._load_staging_buffers.values(): + cached.close() + self._load_staging_buffers.clear() for cached in self._cuda_ipc_cache.values(): cached.close() self._cuda_ipc_cache.clear() diff --git a/tests/connector/test_ipc_client.py b/tests/connector/test_ipc_client.py index c76b4ba..1922677 100644 --- a/tests/connector/test_ipc_client.py +++ b/tests/connector/test_ipc_client.py @@ -370,6 +370,64 @@ async def fake_call(self, payload: dict) -> dict: assert recorded[1]["payload"]["allocation_offset"] == 2272 +@pytest.mark.asyncio +async def test_async_client_registers_and_loads_registered_staging_buffer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Async client can register fixed load staging and load by buffer index.""" + + recorded: list[dict] = [] + + async def fake_call(self, payload: dict) -> dict: + del self + recorded.append(payload) + return {"ok": True, "bytes": 16} + + monkeypatch.setattr(IPCClientAsync, "call", fake_call) + client = IPCClientAsync("/tmp/daser.sock") + + await client.register_load_staging_cuda( + buffer_index=1, + cuda_ipc_handle=b"h" * 64, + allocation_bytes=4096, + device_id=0, + device_ptr=223456, + allocation_base_ptr=221184, + allocation_offset=2272, + producer_pid=43, + ) + response = await client.transfer_load_registered_cuda( + buffer_index=1, + nbytes=16, + spans=[{"target_offset": 0, "nbytes": 16, "file_offset": 0}], + ) + + assert response == {"ok": True, "bytes": 16} + assert recorded == [ + { + "op": "register_load_staging", + "payload": { + "buffer_index": 1, + "cuda_ipc_handle": b"h" * 64, + "allocation_bytes": 4096, + "device_id": 0, + "device_ptr": 223456, + "allocation_base_ptr": 221184, + "allocation_offset": 2272, + "producer_pid": 43, + }, + }, + { + "op": "transfer_load", + "payload": { + "load_staging_buffer_index": 1, + "nbytes": 16, + }, + "spans": [{"target_offset": 0, "nbytes": 16, "file_offset": 0}], + }, + ] + + @pytest.mark.asyncio async def test_clients_transfer_drain(tmp_path): """Sync and async clients can drain server-owned transfer work.""" diff --git a/tests/server/test_ipc_server.py b/tests/server/test_ipc_server.py index 907528b..e783ecf 100644 --- a/tests/server/test_ipc_server.py +++ b/tests/server/test_ipc_server.py @@ -565,6 +565,115 @@ def fake_ensure_transfer(_server: IPCServer) -> FakeTransfer: assert opened_buffers[0].closed == 1 +@pytest.mark.asyncio +async def test_registered_load_staging_reuses_preopened_cuda_mapping( + tmp_path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Registered load staging buffers are opened once and reused by index.""" + + class FakeOpened: + def __init__(self) -> None: + self.array = bytearray(1024) + self.closed = 0 + + def close(self) -> None: + self.closed += 1 + + opened_buffers: list[FakeOpened] = [] + open_calls: list[dict[str, Any]] = [] + loaded_buffers: list[Any] = [] + + def fake_open_cuda_ipc_buffer(**kwargs: Any) -> FakeOpened: + open_calls.append(kwargs) + opened = FakeOpened() + opened_buffers.append(opened) + return opened + + class FakeTransfer(TransferLayer): + async def load_bytes(self, dst: Any, file_offset: int, nbytes: int) -> int: + return 0 + + async def store_bytes(self, src: Any, file_offset: int, nbytes: int) -> int: + return 0 + + async def load_bytes_grouped( + self, + dst: Any, + spans: list[dict[str, int]], + ) -> int: + loaded_buffers.append(dst) + return sum(int(span["nbytes"]) for span in spans) + + def close(self) -> None: + pass + + def fake_ensure_transfer(_server: IPCServer) -> FakeTransfer: + return FakeTransfer() + + monkeypatch.setattr( + "daser.server.ipc.server.open_cuda_ipc_buffer", + fake_open_cuda_ipc_buffer, + ) + monkeypatch.setattr( + "daser.server.ipc.server._CachedCudaArray.synchronize", + lambda _self: None, + ) + monkeypatch.setattr(IPCServer, "_ensure_transfer", fake_ensure_transfer) + + core = make_core() + server = IPCServer(str(tmp_path / "test.sock"), core, make_runtime_config(tmp_path)) + + await server.start() + try: + registered = await _send_recv( + str(tmp_path / "test.sock"), + { + "op": "register_load_staging", + "payload": { + "buffer_index": 1, + "cuda_ipc_handle": b"h" * 64, + "allocation_bytes": 1024, + "device_id": 0, + "device_ptr": 123456, + "allocation_base_ptr": 122880, + "allocation_offset": 576, + "producer_pid": 42, + }, + }, + ) + loaded = await _send_recv( + str(tmp_path / "test.sock"), + { + "op": "transfer_load", + "payload": { + "load_staging_buffer_index": 1, + "nbytes": 128, + }, + "spans": [{"target_offset": 0, "nbytes": 128, "file_offset": 0}], + }, + ) + finally: + await server.stop() + + assert registered == {"ok": True} + assert loaded["ok"] is True + assert loaded["bytes"] == 128 + assert len(opened_buffers) == 1 + assert len(loaded_buffers) == 1 + assert loaded_buffers[0][0:128] == opened_buffers[0].array[0:128] + assert open_calls == [ + { + "handle": b"h" * 64, + "nbytes": 1024, + "device_id": 0, + "local_ptr": None, + "allocation_offset": 576, + } + ] + assert opened_buffers[0].closed == 1 + + @pytest.mark.asyncio async def test_stop_accepting_closes_listener_before_transfer( tmp_path, From 03214e506c6fc79e68586edb1993b437194ae91c Mon Sep 17 00:00:00 2001 From: GentleCold Date: Fri, 26 Jun 2026 18:16:03 +0800 Subject: [PATCH 05/11] refactor(connector): clarify store staging pool - Rename the reusable staging pool to store-only terminology. - Keep load restore paths on the fixed staging pool. - Update connector tests to use the store staging pool name. --- daser/connector/daser_connector.py | 2 +- daser/connector/staging.py | 28 ++++++++++++++++++++----- daser/connector/worker.py | 12 +++++------ tests/connector/test_daser_connector.py | 14 ++++++------- 4 files changed, 37 insertions(+), 19 deletions(-) diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index 46dc178..232d568 100644 --- a/daser/connector/daser_connector.py +++ b/daser/connector/daser_connector.py @@ -133,7 +133,7 @@ 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] = {} diff --git a/daser/connector/staging.py b/daser/connector/staging.py index 8514862..574e7e3 100644 --- a/daser/connector/staging.py +++ b/daser/connector/staging.py @@ -4,6 +4,7 @@ # Standard from dataclasses import dataclass, replace +from typing import Protocol # Third Party import torch @@ -36,9 +37,25 @@ ] = {} +class _CudaStagingLeaseOwner(Protocol): + """Pool protocol required by ``CudaStagingLease``. + + Async/thread-safety: + Implementations define their own ownership model; release is invoked + by the worker thread after CUDA IPC users are finished with the lease. + """ + + def release(self, lease: "CudaStagingLease") -> None: + """Return a lease to its owning pool. + + Args: + lease: Lease previously returned by that pool. + """ + + @dataclass class CudaStagingLease: - """One logical staging allocation leased from ``CudaStagingPool``. + """One logical staging allocation leased from a worker-side staging pool. Args: pool: Owning pool that will receive the allocation on release. @@ -50,7 +67,7 @@ class CudaStagingLease: completion. It must not be reused while an async store future owns it. """ - pool: "CudaStagingPool" + pool: _CudaStagingLeaseOwner tensor: torch.Tensor nbytes: int _released: bool = False @@ -79,8 +96,8 @@ def release(self) -> None: self.pool.release(self) -class CudaStagingPool: - """Reusable worker-side GPU staging buffers for CUDA IPC transfer. +class StoreCudaStagingPool: + """Reusable worker-side GPU staging buffers for store CUDA IPC transfer. Args: device: Device on which staging tensors are allocated. @@ -90,7 +107,8 @@ class CudaStagingPool: Async/thread-safety: The pool is owned by one vLLM worker thread. It does not use locks; the worker tracks async store futures and releases leases only after those - futures complete. + futures complete. Load/prefix restore traffic uses + ``FixedCudaStagingPool`` instead. """ def __init__( diff --git a/daser/connector/worker.py b/daser/connector/worker.py index ef334d2..71781d5 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -34,9 +34,9 @@ DEFAULT_STORE_STAGING_BYTES, FUSED_RESTORE_MIN_SLOTS, CudaStagingLease, - CudaStagingPool, FixedCudaStagingPool, StagedStoreBatch, + StoreCudaStagingPool, ) from daser.connector.staging import ( build_load_copy_runs as _build_load_copy_runs, @@ -445,7 +445,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, self._slot_size, ) - self._staging_pool = CudaStagingPool( + self._store_staging_pool = StoreCudaStagingPool( device=sample.device, initial_bytes=self._store_staging_bytes, max_buffer_bytes=self._store_staging_bytes, @@ -527,7 +527,7 @@ def register_cross_layers_kv_cache( self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, self._slot_size, ) - self._staging_pool = CudaStagingPool( + self._store_staging_pool = StoreCudaStagingPool( device=kv_cache.device, initial_bytes=self._store_staging_bytes, max_buffer_bytes=self._store_staging_bytes, @@ -1331,18 +1331,18 @@ def _acquire_staging( Called from the worker thread. Store-path callers must retain the lease until the background server transfer completes. """ - pool = getattr(self, "_staging_pool", None) + pool = getattr(self, "_store_staging_pool", None) if pool is None: max_bytes = max( nbytes, self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, ) - pool = CudaStagingPool( + pool = StoreCudaStagingPool( device=device, initial_bytes=0, max_buffer_bytes=max_bytes, ) - self._staging_pool = pool + self._store_staging_pool = pool return pool.acquire(nbytes) def _stage_store_batch( diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index bd72b35..18e1ced 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -37,7 +37,7 @@ DEFAULT_ROPE_DELTA_SCALE, DEFAULT_STORE_STAGING_BYTES, MIN_STORE_STAGING_BYTES, - CudaStagingPool, + StoreCudaStagingPool, ) from daser.connector.staging import ( apply_rope_delta_to_key_block as _apply_rope_delta_to_key_block, @@ -2201,7 +2201,7 @@ def __init__(self) -> None: self._slot_size = self.slot_size self._store_staging_bytes = 4096 self._pending_store_staging_limit_bytes = 8192 - self._staging_pool = None + self._store_staging_pool = None self._pending_save_staging_bytes = 0 self._save_futures = [] self._rope_rotary_dim = 8 @@ -2242,7 +2242,7 @@ def __init__(self) -> None: self._slot_size = self.slot_size self._store_staging_bytes = 4096 self._pending_store_staging_limit_bytes = 8192 - self._staging_pool = None + self._store_staging_pool = None self._pending_save_staging_bytes = 0 self._save_futures = [] self._rope_rotary_dim = 8 @@ -2286,7 +2286,7 @@ def __init__(self) -> None: self._slot_size = self.slot_size self._store_staging_bytes = 4096 self._pending_store_staging_limit_bytes = 8192 - self._staging_pool = None + self._store_staging_pool = None self._pending_save_staging_bytes = 0 self._save_futures = [] self._rope_rotary_dim = 8 @@ -3724,9 +3724,9 @@ def __init__(self, total_memory: int) -> None: assert tight_pending == (8 << 30) // 5 -def test_cuda_staging_pool_reuses_preallocated_buffer(): - """Staging pool reuses its init-time allocation after release.""" - pool = CudaStagingPool( +def test_store_cuda_staging_pool_reuses_preallocated_buffer(): + """Store staging pool reuses its init-time allocation after release.""" + pool = StoreCudaStagingPool( device=torch.device("cpu"), initial_bytes=128, max_buffer_bytes=256, From ebd0315b0c8daa6245b7719e46d9ee7cea066825 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Sat, 27 Jun 2026 02:51:29 +0800 Subject: [PATCH 06/11] perf(connector): use fixed staging pools - Reuse one fixed CUDA staging pool implementation for store and load paths. - Block store staging acquisition on completed store futures instead of allocating dynamically. - Consume load pipeline batches in completion order with fixed in-flight depth. --- daser/connector/staging.py | 112 +++--------------- daser/connector/worker.py | 149 +++++++++++++++++++++--- tests/connector/test_daser_connector.py | 32 +++-- tests/connector/test_gds_transfer.py | 45 +++++-- 4 files changed, 216 insertions(+), 122 deletions(-) diff --git a/daser/connector/staging.py b/daser/connector/staging.py index 574e7e3..d906cec 100644 --- a/daser/connector/staging.py +++ b/daser/connector/staging.py @@ -3,6 +3,7 @@ from __future__ import annotations # Standard +from collections.abc import Callable from dataclasses import dataclass, replace from typing import Protocol @@ -96,96 +97,8 @@ def release(self) -> None: self.pool.release(self) -class StoreCudaStagingPool: - """Reusable worker-side GPU staging buffers for store CUDA IPC transfer. - - Args: - device: Device on which staging tensors are allocated. - initial_bytes: Size of the buffer allocated during initialization. - max_buffer_bytes: Maximum size of one staging buffer. - - Async/thread-safety: - The pool is owned by one vLLM worker thread. It does not use locks; the - worker tracks async store futures and releases leases only after those - futures complete. Load/prefix restore traffic uses - ``FixedCudaStagingPool`` instead. - """ - - def __init__( - self, - device: torch.device, - initial_bytes: int, - max_buffer_bytes: int, - ) -> None: - if initial_bytes < 0: - raise ValueError("initial_bytes must be non-negative") - if max_buffer_bytes <= 0: - raise ValueError("max_buffer_bytes must be positive") - self._device = device - self._max_buffer_bytes = max_buffer_bytes - self._free: list[torch.Tensor] = [] - if initial_bytes: - self._free.append( - torch.empty(initial_bytes, dtype=torch.uint8, device=device) - ) - - @property - def max_buffer_bytes(self) -> int: - """Return the maximum bytes allowed for a single staging buffer.""" - return self._max_buffer_bytes - - def acquire(self, nbytes: int) -> CudaStagingLease: - """Lease a staging tensor with at least ``nbytes`` capacity. - - Args: - nbytes: Logical transfer byte count. - - Returns: - Lease whose ``view`` is sized to ``nbytes``. - - Raises: - ValueError: If ``nbytes`` exceeds the configured single-buffer cap. - - Async/thread-safety: - Synchronous worker-thread allocation path. Reuses preallocated - buffers first; only grows when the current workload exceeds the - initialization size. - """ - if nbytes < 0: - raise ValueError("nbytes must be non-negative") - if nbytes > self._max_buffer_bytes: - raise ValueError( - f"staging request {nbytes} exceeds cap {self._max_buffer_bytes}" - ) - selected_idx = -1 - selected_size = 0 - for idx, candidate in enumerate(self._free): - capacity = int(candidate.numel()) - if capacity >= nbytes and (selected_idx < 0 or capacity < selected_size): - selected_idx = idx - selected_size = capacity - if selected_idx >= 0: - tensor = self._free.pop(selected_idx) - else: - capacity = max(nbytes, min(self._max_buffer_bytes, nbytes)) - tensor = torch.empty(capacity, dtype=torch.uint8, device=self._device) - return CudaStagingLease(pool=self, tensor=tensor, nbytes=nbytes) - - def release(self, lease: CudaStagingLease) -> None: - """Return a completed staging lease to the free list. - - Args: - lease: Lease previously returned by ``acquire``. - - Async/thread-safety: - Caller must ensure no CUDA operation or server CUDA IPC mapping is - still using the tensor. - """ - self._free.append(lease.tensor) - - class FixedCudaStagingPool: - """Fixed-size worker-side CUDA staging buffers for load pipelining. + """Fixed-size worker-side CUDA staging buffers. Args: device: Device on which staging tensors are allocated. @@ -193,9 +106,10 @@ class FixedCudaStagingPool: depth: Number of fixed buffers to allocate. Async/thread-safety: - The pool is owned by one vLLM worker load thread. It never allocates - after construction, so hot-path cache-hit loads cannot trigger - unexpected CUDA OOM from staging growth. + The pool is owned by one vLLM worker thread. It never allocates after + construction, so load and store staging cannot grow into unexpected + CUDA OOM. Callers that need blocking semantics can pass a release + callback to ``acquire`` when all buffers are currently leased. """ def __init__( @@ -244,11 +158,18 @@ def buffer(self, index: int) -> torch.Tensor: raise ValueError(f"fixed staging buffer index out of range: {index}") return self._buffers[index] - def acquire(self, nbytes: int) -> CudaStagingLease: + def acquire( + self, + nbytes: int, + wait_for_release: Callable[[], None] | None = None, + ) -> CudaStagingLease: """Lease one preallocated staging buffer. Args: nbytes: Logical transfer byte count. + wait_for_release: Optional callback invoked once when no buffer is + currently free. Store callers use this to wait for an older + asynchronous store to finish and release its lease. Returns: Lease whose view is limited to ``nbytes``. @@ -264,6 +185,8 @@ def acquire(self, nbytes: int) -> CudaStagingLease: f"staging request {nbytes} exceeds fixed staging buffer " f"{self._buffer_bytes}" ) + if not self._free_indices and wait_for_release is not None: + wait_for_release() if not self._free_indices: raise RuntimeError("no fixed staging buffers available") return self.acquire_index(self._free_indices[0], nbytes) @@ -316,6 +239,9 @@ def release(self, lease: CudaStagingLease) -> None: raise ValueError("lease does not belong to this fixed staging pool") +StoreCudaStagingPool = FixedCudaStagingPool + + @dataclass(frozen=True) class StagedStoreBatch: """Worker-owned CUDA staging batch ready for server-side transfer. diff --git a/daser/connector/worker.py b/daser/connector/worker.py index 71781d5..df0e84b 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -6,6 +6,8 @@ import asyncio from collections import deque from collections.abc import Callable, Sequence +from concurrent.futures import FIRST_COMPLETED +from concurrent.futures import wait as wait_futures from dataclasses import dataclass import os import time @@ -83,31 +85,38 @@ _ROPE_WARMUP_BLOCKS = 1 _LOAD_PIPELINE_DEPTH = 2 +_MIN_STORE_STAGING_POOL_DEPTH = 1 _LoadBatch = tuple[int, list[dict[str, int]], list[tuple[int, int, Any]]] -def _run_depth_two_pipeline( +def _run_fixed_depth_pipeline( batches: Sequence[Any], submit: Callable[[Any, int], Any], + is_complete: Callable[[Any], bool], consume: Callable[[Any], int], buffer_indices: Sequence[int], + wait_for_completion: Callable[[Sequence[Any]], Any] | None = None, ) -> int: - """Run a fixed-depth FIFO submit/consume pipeline. + """Run a fixed-depth completion-driven submit/consume pipeline. Args: batches: Batch payloads to submit. submit: Function called with ``(batch, buffer_index)``. It returns an opaque in-flight state. + is_complete: Predicate that returns True when an in-flight state can be + consumed without blocking on its load future. consume: Function called with an in-flight state. It must return the reusable buffer index. buffer_indices: Fixed buffer identifiers available to the pipeline. + wait_for_completion: Optional blocking wait that returns one completed + in-flight state when none are immediately complete. Returns: Maximum number of in-flight batches observed. Async/thread-safety: - Pure synchronous helper. Callers own any synchronization required by - the state returned from ``submit`` before ``consume`` releases a buffer. + Pure synchronous helper. Callers own any synchronization required to + wait for completion and to consume a completed state safely. """ if not buffer_indices: raise ValueError("buffer_indices must not be empty") @@ -122,7 +131,15 @@ def _run_depth_two_pipeline( max_inflight = max(max_inflight, len(in_flight)) while in_flight: - state = in_flight.popleft() + ready_state = next((state for state in in_flight if is_complete(state)), None) + if ready_state is None: + ready_state = ( + wait_for_completion(tuple(in_flight)) + if wait_for_completion is not None + else in_flight[0] + ) + in_flight.remove(ready_state) + state = ready_state reusable_buffer = consume(state) if next_batch < len(batches): in_flight.append(submit(batches[next_batch], reusable_buffer)) @@ -132,6 +149,58 @@ def _run_depth_two_pipeline( return max_inflight +def _run_depth_two_pipeline( + batches: Sequence[Any], + submit: Callable[[Any, int], Any], + consume: Callable[[Any], int], + buffer_indices: Sequence[int], +) -> int: + """Run a fixed-depth pipeline, consuming states immediately. + + Args: + batches: Batch payloads to submit. + submit: Function called with ``(batch, buffer_index)``. It returns an + opaque in-flight state. + consume: Function called with an in-flight state. It must return the + reusable buffer index. + buffer_indices: Fixed buffer identifiers available to the pipeline. + + Returns: + Maximum number of in-flight batches observed. + + Async/thread-safety: + Compatibility wrapper for tests and callers whose states are always + ready when selected. + """ + return _run_fixed_depth_pipeline( + batches=batches, + submit=submit, + is_complete=lambda _state: True, + consume=consume, + buffer_indices=buffer_indices, + ) + + +def _store_staging_pool_depth(buffer_bytes: int, pending_limit_bytes: int) -> int: + """Return fixed store staging pool depth for the configured byte budget. + + Args: + buffer_bytes: Capacity of one fixed staging buffer. + pending_limit_bytes: Total pending store staging byte budget. + + Returns: + Number of fixed store staging buffers to preallocate. + + Async/thread-safety: + Pure helper used during worker-side pool initialization. + """ + if buffer_bytes <= 0: + raise ValueError("buffer_bytes must be positive") + if pending_limit_bytes <= 0: + return _MIN_STORE_STAGING_POOL_DEPTH + return max(_MIN_STORE_STAGING_POOL_DEPTH, pending_limit_bytes // buffer_bytes) + + @dataclass class _DeferredFinishedSave: """Store work held until vLLM reports a request as finished.""" @@ -447,8 +516,11 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: ) self._store_staging_pool = StoreCudaStagingPool( device=sample.device, - initial_bytes=self._store_staging_bytes, - max_buffer_bytes=self._store_staging_bytes, + buffer_bytes=self._store_staging_bytes, + depth=_store_staging_pool_depth( + self._store_staging_bytes, + self._pending_store_staging_limit_bytes, + ), ) self._load_staging_pool = FixedCudaStagingPool( device=sample.device, @@ -529,8 +601,11 @@ def register_cross_layers_kv_cache( ) self._store_staging_pool = StoreCudaStagingPool( device=kv_cache.device, - initial_bytes=self._store_staging_bytes, - max_buffer_bytes=self._store_staging_bytes, + buffer_bytes=self._store_staging_bytes, + depth=_store_staging_pool_depth( + self._store_staging_bytes, + self._pending_store_staging_limit_bytes, + ), ) self._load_staging_pool = FixedCudaStagingPool( device=kv_cache.device, @@ -716,15 +791,29 @@ def consume_load_batch(state: _InflightLoadBatch) -> int: consumed_batches.append(consumed) return consumed.buffer_index - max_inflight_batches = _run_depth_two_pipeline( + def wait_for_load_completion( + states: Sequence[_InflightLoadBatch], + ) -> _InflightLoadBatch: + """Wait until any in-flight load future completes.""" + future_to_state = {state.future: state for state in states} + done, _pending = wait_futures( + future_to_state.keys(), + return_when=FIRST_COMPLETED, + ) + first_done = next(iter(done)) + return future_to_state[first_done] + + max_inflight_batches = _run_fixed_depth_pipeline( batches=load_batches, submit=lambda batch, buffer_index: self._submit_load_batch( batch, buffer_index, sample_tensor, ), + is_complete=lambda state: state.future.done(), consume=consume_load_batch, buffer_indices=range(_LOAD_PIPELINE_DEPTH), + wait_for_completion=wait_for_load_completion, ) for consumed in consumed_batches: @@ -1313,6 +1402,33 @@ def _wait_for_save_staging_capacity(self, nbytes: int) -> None: record.release() self._reap_save_futures(block=False) + def _wait_for_store_staging_release(self, nbytes: int) -> None: + """Wait until one store staging lease can return to the fixed pool. + + Args: + nbytes: Size of the staging lease requested by the caller. + + Async/thread-safety: + Called from the worker thread when the fixed store staging pool is + exhausted. It first applies byte-budget backpressure, then waits + for one oldest store future if no lease was released yet. + """ + pool = getattr(self, "_store_staging_pool", None) + before = pool.available if pool is not None else 0 + self._wait_for_save_staging_capacity(nbytes) + if pool is None or pool.available > before or not self._save_futures: + return + record = self._save_futures.pop(0) + try: + record.future.result(timeout=120.0) + finally: + self._pending_save_staging_bytes = max( + 0, + self._pending_save_staging_bytes - record.staging_bytes, + ) + record.release() + self._reap_save_futures(block=False) + def _acquire_staging( self, nbytes: int, @@ -1339,11 +1455,18 @@ def _acquire_staging( ) pool = StoreCudaStagingPool( device=device, - initial_bytes=0, - max_buffer_bytes=max_bytes, + buffer_bytes=max_bytes, + depth=_store_staging_pool_depth( + max_bytes, + self._pending_store_staging_limit_bytes + or DEFAULT_PENDING_STORE_STAGING_BYTES, + ), ) self._store_staging_pool = pool - return pool.acquire(nbytes) + return pool.acquire( + nbytes, + wait_for_release=lambda: self._wait_for_store_staging_release(nbytes), + ) def _stage_store_batch( self, diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index 18e1ced..741ccd5 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -805,6 +805,9 @@ def release(self) -> None: return class _Future: + def done(self) -> bool: + return True + def result(self, timeout: float): del timeout return { @@ -843,8 +846,8 @@ def fake_submit_load_coroutine(coro): assert "start_load_kv timing" not in caplog.text -def test_worker_load_pipeline_submits_two_batches_before_consuming(monkeypatch) -> None: - """Worker load path submits two fixed-buffer batches before consuming.""" +def test_worker_load_pipeline_consumes_completed_batch_first(monkeypatch) -> None: + """Worker load path consumes whichever fixed-buffer batch completes first.""" class Probe(WorkerConnectorMixin): def __init__(self) -> None: @@ -864,6 +867,13 @@ def load_specs(self, reqs_to_load: dict[str, ReqLoadSpec]) -> None: """Expose load execution through a public test helper.""" self._load_kv_specs(reqs_to_load, next(iter(self._kv_caches.values()))) + class _FakeFuture: + def __init__(self, total_bytes: int) -> None: + self.total_bytes = total_bytes + + def done(self) -> bool: + return completed_by_bytes[self.total_bytes] + probe = Probe() reqs_to_load = { "req": ReqLoadSpec( @@ -876,17 +886,25 @@ def load_specs(self, reqs_to_load: dict[str, ReqLoadSpec]) -> None: ) } events: list[str] = [] + completed_by_bytes = {8: False, 4: True} def fake_submit(batch, buffer_index: int, sample_tensor: torch.Tensor): del sample_tensor total_bytes, _spans, _per_req_ranges = batch events.append(f"submit:{total_bytes}:buf{buffer_index}") - return ("state", total_bytes, buffer_index) + return SimpleNamespace( + total_bytes=total_bytes, + buffer_index=buffer_index, + future=_FakeFuture(total_bytes), + ) def fake_consume(state, sample_tensor: torch.Tensor): del sample_tensor - _tag, total_bytes, buffer_index = state + total_bytes = state.total_bytes + buffer_index = state.buffer_index events.append(f"consume:{total_bytes}:buf{buffer_index}") + if total_bytes == 4: + completed_by_bytes[8] = True return SimpleNamespace( buffer_index=buffer_index, bytes=total_bytes, @@ -914,7 +932,7 @@ def fake_consume(state, sample_tensor: torch.Tensor): assert events[:3] == [ "submit:8:buf0", "submit:4:buf1", - "consume:8:buf0", + "consume:4:buf1", ] @@ -3728,8 +3746,8 @@ def test_store_cuda_staging_pool_reuses_preallocated_buffer(): """Store staging pool reuses its init-time allocation after release.""" pool = StoreCudaStagingPool( device=torch.device("cpu"), - initial_bytes=128, - max_buffer_bytes=256, + buffer_bytes=128, + depth=1, ) lease = pool.acquire(64) diff --git a/tests/connector/test_gds_transfer.py b/tests/connector/test_gds_transfer.py index 1a714ce..9ab1110 100644 --- a/tests/connector/test_gds_transfer.py +++ b/tests/connector/test_gds_transfer.py @@ -87,7 +87,7 @@ def test_missing_file_raises(tmp_path): def test_fixed_staging_pool_reuses_two_preallocated_buffers() -> None: - """Fixed load staging pool reuses bounded preallocated buffers.""" + """Fixed staging pool reuses bounded preallocated buffers.""" # First Party from daser.connector.staging import FixedCudaStagingPool @@ -111,8 +111,33 @@ def test_fixed_staging_pool_reuses_two_preallocated_buffers() -> None: third.release() +def test_fixed_staging_pool_can_block_until_buffer_release() -> None: + """Fixed staging pool can wait for a callback to release capacity.""" + # First Party + from daser.connector.staging import FixedCudaStagingPool + + pool = FixedCudaStagingPool( + device=torch.device("cpu"), + buffer_bytes=16, + depth=1, + ) + first = pool.acquire(8) + waits = 0 + + def release_first() -> None: + nonlocal waits + waits += 1 + first.release() + + second = pool.acquire(4, wait_for_release=release_first) + + assert waits == 1 + assert second.tensor.data_ptr() == first.tensor.data_ptr() + second.release() + + def test_fixed_staging_pool_rejects_oversized_request() -> None: - """Fixed load staging pool rejects requests larger than one buffer.""" + """Fixed staging pool rejects requests larger than one buffer.""" # First Party from daser.connector.staging import FixedCudaStagingPool @@ -126,12 +151,13 @@ def test_fixed_staging_pool_rejects_oversized_request() -> None: pool.acquire(17) -def test_depth_two_pipeline_submits_next_batch_after_oldest_consumes() -> None: - """Depth-two load pipeline reuses the oldest consumed buffer first.""" +def test_depth_two_pipeline_consumes_completed_batch_first() -> None: + """Depth-two load pipeline reuses whichever buffer completes first.""" # First Party - from daser.connector.worker import _run_depth_two_pipeline + from daser.connector.worker import _run_fixed_depth_pipeline events: list[str] = [] + ready: dict[str, bool] = {"b0": False, "b1": True, "b2": True} def submit(batch: str, buffer_index: int) -> tuple[str, int]: events.append(f"submit:{batch}:buf{buffer_index}") @@ -142,9 +168,10 @@ def consume(state: tuple[str, int]) -> int: events.append(f"consume:{batch}:buf{buffer_index}") return buffer_index - max_inflight = _run_depth_two_pipeline( + max_inflight = _run_fixed_depth_pipeline( batches=["b0", "b1", "b2"], submit=submit, + is_complete=lambda state: ready[state[0]], consume=consume, buffer_indices=[0, 1], ) @@ -153,8 +180,8 @@ def consume(state: tuple[str, int]) -> int: assert events == [ "submit:b0:buf0", "submit:b1:buf1", - "consume:b0:buf0", - "submit:b2:buf0", "consume:b1:buf1", - "consume:b2:buf0", + "submit:b2:buf1", + "consume:b2:buf1", + "consume:b0:buf0", ] From 13ef6b2134d6b28346719e4cf79a6bffae7ada5e Mon Sep 17 00:00:00 2001 From: GentleCold Date: Sat, 27 Jun 2026 04:44:22 +0800 Subject: [PATCH 07/11] perf(connector): parallelize tiered load path - add fixed load IPC clients so staged load batches use independent server connections - consume io_uring L2 miss reads as they complete and copy each batch to the destination immediately - cover buffer-scoped load clients and completion-driven L2 miss promotion --- daser/connector/daser_connector.py | 9 +- daser/connector/worker.py | 71 ++++++++++-- daser/transfer/iouring/layer.py | 72 +++++++------ tests/connector/test_daser_connector.py | 35 ++++++ .../transfer/test_tiered_iouring_transfer.py | 101 ++++++++++++++++++ 5 files changed, 246 insertions(+), 42 deletions(-) diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index 232d568..406d414 100644 --- a/daser/connector/daser_connector.py +++ b/daser/connector/daser_connector.py @@ -40,7 +40,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_PIPELINE_DEPTH, WorkerConnectorMixin from daser.logging import init_logger logger = init_logger(__name__) @@ -124,6 +124,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_PIPELINE_DEPTH - 1)) + ], + ] self._ipc_store_async = IPCClientAsync(self._socket_path) self._kv_caches: dict[str, torch.Tensor] = {} self._layer_names: list[str] = [] diff --git a/daser/connector/worker.py b/daser/connector/worker.py index df0e84b..d525b39 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -934,6 +934,7 @@ def _submit_load_batch( device_ptr = cuda_array_pointer(cp_staging) ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) transfer_coro = self._transfer_load_cuda( + buffer_index=buffer_index, cuda_ipc_handle=cuda_handle, nbytes=total_bytes, device_id=device_id, @@ -1189,10 +1190,16 @@ def shutdown(self) -> None: load_executor = getattr(self, "_load_executor", None) if load_executor is not None: load_executor.shutdown(wait=True) - load_client = getattr(self, "_ipc_load_async", None) + load_clients = list( + dict.fromkeys( + getattr(self, "_ipc_load_async_pool", []) + or [getattr(self, "_ipc_load_async", None)] + ) + ) store_client = getattr(self, "_ipc_store_async", None) - if load_client is not None: - self._submit_load_coroutine(load_client.close()).result(timeout=10.0) + for load_client in load_clients: + if load_client is not None: + self._submit_load_coroutine(load_client.close()).result(timeout=10.0) if store_client is not None: self._submit_store_coroutine(store_client.close()).result(timeout=10.0) self._load_loop.call_soon_threadsafe(self._load_loop.stop) @@ -1277,14 +1284,53 @@ def _init_server_transfer(self) -> None: and getattr(self, "_store_loop", None) is not None ): return - self._submit_load_coroutine(self._ipc_load_async.init_transfer()).result( - timeout=120.0 - ) + for load_client in self._load_ipc_clients(): + self._submit_load_coroutine(load_client.init_transfer()).result( + timeout=120.0 + ) self._submit_store_coroutine(self._ipc_store_async.init_transfer()).result( timeout=120.0 ) self._register_load_staging_buffers() + def _load_ipc_clients(self) -> list[Any]: + """Return fixed load IPC clients used for parallel load RPCs. + + Returns: + Load IPC clients. The list falls back to the legacy single client + when the connector was constructed by tests or older harnesses. + + Async/thread-safety: + Called on the vLLM worker thread during initialization and from the + load event loop when selecting the client for a submitted batch. + """ + clients = getattr(self, "_ipc_load_async_pool", None) + if clients: + return list(clients) + client = getattr(self, "_ipc_load_async", None) + return [client] if client is not None else [] + + def _load_ipc_client_for_buffer(self, buffer_index: int | None = None) -> Any: + """Return the load IPC client assigned to a staging buffer. + + Args: + buffer_index: Fixed staging buffer index for the transfer. + + Returns: + Async IPC client for the selected load lane. + + Async/thread-safety: + Pure selection helper. Each returned client owns its own IPC socket, + so separate fixed staging buffers can have concurrent server RPCs + instead of serializing on one client lock. + """ + clients = self._load_ipc_clients() + if not clients: + raise RuntimeError("load IPC client is not initialized") + if buffer_index is None: + return clients[0] + return clients[int(buffer_index) % len(clients)] + def _register_load_staging_buffers(self) -> None: """Register fixed load staging buffers with the server. @@ -1296,7 +1342,7 @@ def _register_load_staging_buffers(self) -> None: if getattr(self, "_load_staging_registered", False): return pool = getattr(self, "_load_staging_pool", None) - if pool is None or getattr(self, "_ipc_load_async", None) is None: + if pool is None or not self._load_ipc_clients(): return try: for buffer_index in range(_LOAD_PIPELINE_DEPTH): @@ -1304,8 +1350,9 @@ def _register_load_staging_buffers(self) -> None: cp_tensor = cupy.asarray(tensor) device_ptr = cuda_array_pointer(cp_tensor) ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) + load_client = self._load_ipc_client_for_buffer(buffer_index) self._submit_load_coroutine( - self._ipc_load_async.register_load_staging_cuda( + load_client.register_load_staging_cuda( buffer_index=buffer_index, cuda_ipc_handle=export_cuda_ipc_handle(cp_tensor), allocation_bytes=int(tensor.numel()), @@ -1665,7 +1712,9 @@ async def _transfer_load_cuda(self, **kwargs: Any) -> dict[str, Any]: Runs on the worker load event loop. A dedicated client keeps cache-hit loads from queueing behind store RPCs. """ - return await self._ipc_load_async.transfer_load_cuda(**kwargs) + buffer_index = kwargs.pop("buffer_index", None) + client = self._load_ipc_client_for_buffer(buffer_index) + return await client.transfer_load_cuda(**kwargs) async def _transfer_load_registered_cuda(self, **kwargs: Any) -> dict[str, Any]: """Load through a pre-registered fixed CUDA staging buffer. @@ -1681,7 +1730,9 @@ async def _transfer_load_registered_cuda(self, **kwargs: Any) -> dict[str, Any]: the CUDA IPC mapping during initialization, so this hot-path call only identifies the staging buffer index and logical byte range. """ - return await self._ipc_load_async.transfer_load_registered_cuda(**kwargs) + buffer_index = int(kwargs.get("buffer_index", 0)) + client = self._load_ipc_client_for_buffer(buffer_index) + return await client.transfer_load_registered_cuda(**kwargs) async def _transfer_store_cuda(self, **kwargs: Any) -> list[str]: """Store through the dedicated worker store IPC client. diff --git a/daser/transfer/iouring/layer.py b/daser/transfer/iouring/layer.py index 95a82a7..0d79fd4 100644 --- a/daser/transfer/iouring/layer.py +++ b/daser/transfer/iouring/layer.py @@ -532,51 +532,61 @@ async def _load_l2_miss_batch( if self._l2 is None: raise RuntimeError("L2 reads are disabled when skip_l2 is true") reads: list[tuple[dict[str, int], PinnedMemorySlice]] = [] + future_to_read: dict[ + asyncio.Future[int], + tuple[dict[str, int], PinnedMemorySlice], + ] = {} + promoted: set[int] = set() try: for span in misses: nbytes = int(span["nbytes"]) key = (int(span["file_offset"]), nbytes) pinned = await self._reserve_l1_buffer(key, nbytes) reads.append((span, pinned)) - - await asyncio.gather( - *( - loop.run_in_executor( - self._l2.executor, - self._read_l2_into, - int(span["file_offset"]), - pinned, - self._next_uring(), - ) - for span, pinned in reads + future = loop.run_in_executor( + self._l2.executor, + self._read_l2_into, + int(span["file_offset"]), + pinned, + self._next_uring(), ) - ) + future_to_read[future] = (span, pinned) - self._copy_grouped_to_dst( - dst, - [ - ( - int(span["target_offset"]), - pinned, - 0, - int(span["nbytes"]), + pending: set[asyncio.Future[int]] = set(future_to_read) + while pending: + done, pending = await asyncio.wait( + pending, + return_when=asyncio.FIRST_COMPLETED, + ) + for future in done: + future.result() + span, pinned = future_to_read[future] + self._copy_grouped_to_dst( + dst, + [ + ( + int(span["target_offset"]), + pinned, + 0, + int(span["nbytes"]), + ) + ], ) - for span, pinned in reads - ], - ) - async with self._lock: - self._raise_l2_error_locked() - for span, pinned in reads: - nbytes = int(span["nbytes"]) - key = (int(span["file_offset"]), nbytes) - self._stats.l2_reads += 1 - self._l1.put(key, pinned) + async with self._lock: + self._raise_l2_error_locked() + nbytes = int(span["nbytes"]) + key = (int(span["file_offset"]), nbytes) + self._stats.l2_reads += 1 + self._l1.put(key, pinned) + promoted.add(id(pinned)) except BaseException: + if future_to_read: + await asyncio.gather(*future_to_read, return_exceptions=True) live_buffers = set() async with self._lock: live_buffers = self._l1.resident_slice_ids() for _span, pinned in reads: - if id(pinned) not in live_buffers: + if id(pinned) not in live_buffers and id(pinned) not in promoted: pinned.close() raise diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index 741ccd5..21595ad 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -936,6 +936,41 @@ def fake_consume(state, sample_tensor: torch.Tensor): ] +def test_worker_load_batches_use_buffer_scoped_ipc_clients() -> None: + """Parallel load batches should not serialize on one async IPC client lock.""" + + class _Client: + def __init__(self, name: str) -> None: + self.name = name + self.calls: list[dict[str, object]] = [] + + async def transfer_load_registered_cuda(self, **kwargs): + self.calls.append(kwargs) + return {"client": self.name} + + class Probe(WorkerConnectorMixin): + def __init__(self) -> None: + self._ipc_load_async = _Client("fallback") + self._ipc_load_async_pool = [_Client("load0"), _Client("load1")] + + probe = Probe() + + coro0 = probe._transfer_load_registered_cuda( # noqa: SLF001 + buffer_index=0, + nbytes=4, + spans=[], + ) + coro1 = probe._transfer_load_registered_cuda( # noqa: SLF001 + buffer_index=1, + nbytes=4, + spans=[], + ) + + assert asyncio.run(coro0) == {"client": "load0"} + assert asyncio.run(coro1) == {"client": "load1"} + assert probe._ipc_load_async.calls == [] # noqa: SLF001 + + def test_start_load_kv_releases_waiting_request_when_load_cannot_start() -> None: """Worker should not leave async-hit requests waiting forever.""" connector = _WorkerProbe("") diff --git a/tests/transfer/test_tiered_iouring_transfer.py b/tests/transfer/test_tiered_iouring_transfer.py index 812d0fe..adecced 100644 --- a/tests/transfer/test_tiered_iouring_transfer.py +++ b/tests/transfer/test_tiered_iouring_transfer.py @@ -94,6 +94,49 @@ def _read_l2_into( return super()._read_l2_into(file_offset, dst, uring) +class DelayedL2ReadProbe(TieredIOUringTransferLayer): + """Test transfer layer that pauses selected L2 reads.""" + + def __init__( + self, + path: str, + l1_bytes: int, + l2_bytes: int, + delayed_offsets: set[int], + ) -> None: + super().__init__( + path=path, + l1_bytes=l1_bytes, + l2_bytes=l2_bytes, + ) + self.delayed_offsets = delayed_offsets + self.release_read = threading.Event() + self.read_started = threading.Event() + self.copy_offsets: list[int] = [] + + def _read_l2_into( + self, + file_offset: int, + dst: object, + uring: NativeIOUring, + ) -> int: + """Pause configured reads before delegating to the real L2 reader.""" + if file_offset in self.delayed_offsets: + self.delayed_offsets.remove(file_offset) + self.read_started.set() + self.release_read.wait(timeout=5.0) + return super()._read_l2_into(file_offset, dst, uring) + + def _copy_grouped_to_dst( + self, + dst: object, + chunks: list[tuple[int, object, int, int]], + ) -> None: + """Record destination copy offsets before delegating.""" + self.copy_offsets.extend(int(chunk[0]) for chunk in chunks) + super()._copy_grouped_to_dst(dst, chunks) + + def _run(coro: object) -> object: """Run a coroutine on the current test event loop.""" return asyncio.get_event_loop().run_until_complete(coro) @@ -598,6 +641,64 @@ async def load_block(i: int) -> tuple[int, bytes]: _run(scenario()) +def test_iouring_l2_miss_batch_copies_each_read_as_it_completes(tmp_path) -> None: + """A completed L2 miss should copy to destination before slower reads finish.""" + + async def scenario() -> None: + path = str(tmp_path / "daser.store") + writer = TieredIOUringTransferLayer( + path=path, + l1_bytes=ALIGNMENT, + l2_bytes=ALIGNMENT * 2, + ) + try: + await writer.store_bytes(_block(b"a"), file_offset=0, nbytes=ALIGNMENT) + await writer.store_bytes( + _block(b"b"), + file_offset=ALIGNMENT, + nbytes=ALIGNMENT, + ) + await writer.drain() + finally: + writer.close() + + layer = DelayedL2ReadProbe( + path=path, + l1_bytes=ALIGNMENT * 2, + l2_bytes=ALIGNMENT * 2, + delayed_offsets={0}, + ) + try: + dst = bytearray(ALIGNMENT * 2) + task = asyncio.create_task( + layer.load_bytes_grouped( + dst, + [ + {"target_offset": 0, "file_offset": 0, "nbytes": ALIGNMENT}, + { + "target_offset": ALIGNMENT, + "file_offset": ALIGNMENT, + "nbytes": ALIGNMENT, + }, + ], + ) + ) + assert await asyncio.to_thread(layer.read_started.wait, timeout=2.0) + deadline = time.perf_counter() + 2.0 + while not layer.copy_offsets and time.perf_counter() < deadline: + await asyncio.sleep(0.01) + assert layer.copy_offsets == [ALIGNMENT] + layer.release_read.set() + loaded = await asyncio.wait_for(task, timeout=2.0) + assert loaded == ALIGNMENT * 2 + assert bytes(dst) == bytes(_block(b"a") + _block(b"b")) + finally: + layer.release_read.set() + layer.close() + + _run(scenario()) + + def test_iouring_store_returns_after_l1_before_l2_flush(tmp_path) -> None: """Store returns once L1 is readable while L2 persistence continues.""" From 0de4a04f8e0091b6f082c879de255867ae06ed61 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Sat, 27 Jun 2026 12:32:50 +0800 Subject: [PATCH 08/11] perf(connector): release loaded requests independently - track request-level load futures while preserving grouped load pipeline submission - mark each request complete after all batches containing it restore into KV cache - cover independent request load completion and request-id annotated load ranges --- daser/connector/staging.py | 39 ++- daser/connector/worker.py | 337 ++++++++++++++++-------- tests/connector/test_daser_connector.py | 130 +++++++++ 3 files changed, 392 insertions(+), 114 deletions(-) diff --git a/daser/connector/staging.py b/daser/connector/staging.py index d906cec..61db637 100644 --- a/daser/connector/staging.py +++ b/daser/connector/staging.py @@ -5,7 +5,7 @@ # Standard from collections.abc import Callable from dataclasses import dataclass, replace -from typing import Protocol +from typing import Any, Protocol # Third Party import torch @@ -900,12 +900,15 @@ def copy_cross_layer_kv_cache_to_staging( def build_load_read_plan( reqs_to_load: dict[str, ReqLoadSpec], slot_size: int, -) -> tuple[int, list[dict[str, int]], list[tuple[int, int, ReqLoadSpec]]]: + include_req_ids: bool = False, +) -> tuple[int, list[dict[str, int]], list[Any]]: """Build a combined transfer-load plan for one forward step. Args: reqs_to_load: request ID to load spec from scheduler metadata. slot_size: bytes per vLLM KV slot. + include_req_ids: when True, include request IDs in per-request ranges + for worker-side request completion tracking. Returns: ``(total_bytes, spans, per_req_ranges)`` where spans target one @@ -914,9 +917,9 @@ def build_load_read_plan( """ total_bytes = 0 spans: list[dict[str, int]] = [] - per_req_ranges: list[tuple[int, int, ReqLoadSpec]] = [] + per_req_ranges: list[Any] = [] source_ranges: dict[tuple[str, int, int, int, int], tuple[int, int]] = {} - for spec in reqs_to_load.values(): + for req_id, spec in reqs_to_load.items(): num_slots = len(spec.block_ids) if num_slots == 0: continue @@ -937,7 +940,10 @@ def build_load_read_plan( total_bytes = end else: start, end = existing - per_req_ranges.append((start, end, spec)) + if include_req_ids: + per_req_ranges.append((start, end, req_id, spec)) + else: + per_req_ranges.append((start, end, spec)) return total_bytes, spans, per_req_ranges @@ -945,13 +951,16 @@ def build_load_read_batches( reqs_to_load: dict[str, ReqLoadSpec], slot_size: int, max_batch_bytes: int, -) -> list[tuple[int, list[dict[str, int]], list[tuple[int, int, ReqLoadSpec]]]]: + include_req_ids: bool = False, +) -> list[tuple[int, list[dict[str, int]], list[Any]]]: """Build bounded load staging plans for one forward step. Args: reqs_to_load: request ID to load spec from scheduler metadata. slot_size: bytes per vLLM KV slot. max_batch_bytes: Maximum staging bytes for one transfer batch. + include_req_ids: when True, include request IDs in per-request ranges + for worker-side request completion tracking. Returns: List of ``build_load_read_plan``-style tuples. Individual requests are @@ -965,9 +974,7 @@ def build_load_read_batches( if max_batch_bytes <= 0: raise ValueError("max_batch_bytes must be positive") max_slots = max(1, max_batch_bytes // slot_size) - batches: list[ - tuple[int, list[dict[str, int]], list[tuple[int, int, ReqLoadSpec]]] - ] = [] + batches: list[tuple[int, list[dict[str, int]], list[Any]]] = [] current: dict[str, ReqLoadSpec] = {} current_slots = 0 synthetic_id = 0 @@ -975,7 +982,13 @@ def build_load_read_batches( def flush() -> None: nonlocal current, current_slots if current: - batches.append(build_load_read_plan(current, slot_size)) + batches.append( + build_load_read_plan( + current, + slot_size, + include_req_ids=include_req_ids, + ) + ) current = {} current_slots = 0 @@ -1047,7 +1060,11 @@ def flush() -> None: run_pos_offset = 0 run_block_ids = [] - for start, end, spec in per_req_ranges: + for item in per_req_ranges: + if len(item) == 3: + start, end, spec = item + else: + start, end, _req_id, spec = item if not spec.block_ids: continue if run_start >= 0 and start == run_end and spec.pos_offset == run_pos_offset: diff --git a/daser/connector/worker.py b/daser/connector/worker.py index d525b39..707c752 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -10,6 +10,7 @@ from concurrent.futures import wait as wait_futures from dataclasses import dataclass import os +import threading import time from typing import TYPE_CHECKING, Any @@ -86,7 +87,7 @@ _ROPE_WARMUP_BLOCKS = 1 _LOAD_PIPELINE_DEPTH = 2 _MIN_STORE_STAGING_POOL_DEPTH = 1 -_LoadBatch = tuple[int, list[dict[str, int]], list[tuple[int, int, Any]]] +_LoadBatch = tuple[int, list[dict[str, int]], list[Any]] def _run_fixed_depth_pipeline( @@ -268,7 +269,7 @@ class _InflightLoadBatch: buffer_index: int total_bytes: int - per_req_ranges: list[tuple[int, int, Any]] + per_req_ranges: list[Any] staging_lease: CudaStagingLease future: Any submitted_at: float @@ -321,6 +322,48 @@ def result(self, timeout: float | None = None) -> None: raise RuntimeError(self._message) +class _RequestLoadFuture: + """Small future used to release request loads independently. + + Async/thread-safety: + The connector load executor marks completion and the vLLM worker thread + polls ``done``/``result`` from ``get_finished``. ``threading.Event`` + provides cross-thread visibility for the result state. + """ + + def __init__(self) -> None: + self._event = threading.Event() + self._error: BaseException | None = None + + def done(self) -> bool: + """Return whether this request load has completed.""" + return self._event.is_set() + + def result(self, timeout: float | None = None) -> None: + """Wait for completion and raise a stored load error, if any. + + Args: + timeout: Optional timeout in seconds. + """ + if not self._event.wait(timeout): + raise TimeoutError("request load did not complete before timeout") + if self._error is not None: + raise self._error + + def set_result(self) -> None: + """Mark this request load as successful.""" + self._event.set() + + def set_exception(self, error: BaseException) -> None: + """Mark this request load as failed. + + Args: + error: Exception to re-raise from ``result``. + """ + self._error = error + self._event.set() + + def _cuda_allocation_base_and_offset(device_ptr: int) -> tuple[int, int]: """Return CUDA allocation base pointer and byte offset for ``device_ptr``. @@ -694,22 +737,28 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non self._mark_load_start_failed(reqs_to_load, "no registered KV cache tensor") return - block_ids = [ - block_id for spec in reqs_to_load.values() for block_id in spec.block_ids - ] - base_req_ids = {base_req_id(req_id) for req_id in reqs_to_load} - future = self._load_executor.submit( + request_futures = { + base_req_id(req_id): _RequestLoadFuture() for req_id in reqs_to_load + } + self._load_executor.submit( self._load_kv_specs, reqs_to_load, sample_tensor, + request_futures, ) pending_loads = getattr(self, "_pending_loads", None) if pending_loads is None: pending_loads = {} self._pending_loads = pending_loads - for req_id in base_req_ids: + for req_id, request_future in request_futures.items(): + block_ids = [ + block_id + for spec_req_id, spec in reqs_to_load.items() + if base_req_id(spec_req_id) == req_id + for block_id in spec.block_ids + ] pending_loads[req_id] = _PendingLoad( - future=future, + future=request_future, block_ids=list(block_ids), lease=None, ) @@ -750,6 +799,7 @@ def _load_kv_specs( self, reqs_to_load: dict[str, Any], sample_tensor: torch.Tensor, + request_futures: dict[str, _RequestLoadFuture] | None = None, ) -> None: """Load cross-layer KV specs into vLLM cache on a background thread. @@ -757,108 +807,126 @@ def _load_kv_specs( reqs_to_load: Worker load specs from scheduler metadata. sample_tensor: Representative KV cache tensor used for device and synchronization. + request_futures: Optional base request ID to completion future map. + When provided, each request future is marked as soon as all + batches containing that request have been restored. Async/thread-safety: Runs on the connector load executor. It submits server transfer RPCs to the dedicated load asyncio loop and waits inside the background thread, not on vLLM's worker thread. """ - load_staging_pool = self._ensure_load_staging_pool(sample_tensor) - load_batches = _build_load_read_batches( - reqs_to_load, - self._slot_size, - max_batch_bytes=load_staging_pool.buffer_bytes, - ) - if not load_batches: - return - - total_copies = 0 - total_copy_runs = 0 - total_bytes_loaded = 0 - total_ipc_ms = 0.0 - total_copy_ms = 0.0 - total_sync_ms = 0.0 - total_l1_hits = 0 - total_l1_misses = 0 - total_l2_reads = 0 - total_transfer_open_ms = 0.0 - total_transfer_load_ms = 0.0 - total_transfer_sync_ms = 0.0 - consumed_batches: list[_ConsumedLoadBatch] = [] - - def consume_load_batch(state: _InflightLoadBatch) -> int: - consumed = self._consume_loaded_batch(state, sample_tensor) - consumed_batches.append(consumed) - return consumed.buffer_index - - def wait_for_load_completion( - states: Sequence[_InflightLoadBatch], - ) -> _InflightLoadBatch: - """Wait until any in-flight load future completes.""" - future_to_state = {state.future: state for state in states} - done, _pending = wait_futures( - future_to_state.keys(), - return_when=FIRST_COMPLETED, + try: + load_staging_pool = self._ensure_load_staging_pool(sample_tensor) + load_batches = _build_load_read_batches( + reqs_to_load, + self._slot_size, + max_batch_bytes=load_staging_pool.buffer_bytes, + include_req_ids=request_futures is not None, + ) + if not load_batches: + for request_future in (request_futures or {}).values(): + request_future.set_result() + return + remaining_batches = self._request_load_batch_counts(load_batches) + + total_copies = 0 + total_copy_runs = 0 + total_bytes_loaded = 0 + total_ipc_ms = 0.0 + total_copy_ms = 0.0 + total_sync_ms = 0.0 + total_l1_hits = 0 + total_l1_misses = 0 + total_l2_reads = 0 + total_transfer_open_ms = 0.0 + total_transfer_load_ms = 0.0 + total_transfer_sync_ms = 0.0 + consumed_batches: list[_ConsumedLoadBatch] = [] + + def consume_load_batch(state: _InflightLoadBatch) -> int: + consumed = self._consume_loaded_batch(state, sample_tensor) + consumed_batches.append(consumed) + self._mark_loaded_batch_requests( + state, + request_futures, + remaining_batches, + ) + return consumed.buffer_index + + def wait_for_load_completion( + states: Sequence[_InflightLoadBatch], + ) -> _InflightLoadBatch: + """Wait until any in-flight load future completes.""" + future_to_state = {state.future: state for state in states} + done, _pending = wait_futures( + future_to_state.keys(), + return_when=FIRST_COMPLETED, + ) + first_done = next(iter(done)) + return future_to_state[first_done] + + max_inflight_batches = _run_fixed_depth_pipeline( + batches=load_batches, + submit=lambda batch, buffer_index: self._submit_load_batch( + batch, + buffer_index, + sample_tensor, + ), + is_complete=lambda state: state.future.done(), + consume=consume_load_batch, + buffer_indices=range(_LOAD_PIPELINE_DEPTH), + wait_for_completion=wait_for_load_completion, ) - first_done = next(iter(done)) - return future_to_state[first_done] - - max_inflight_batches = _run_fixed_depth_pipeline( - batches=load_batches, - submit=lambda batch, buffer_index: self._submit_load_batch( - batch, - buffer_index, - sample_tensor, - ), - is_complete=lambda state: state.future.done(), - consume=consume_load_batch, - buffer_indices=range(_LOAD_PIPELINE_DEPTH), - wait_for_completion=wait_for_load_completion, - ) - - for consumed in consumed_batches: - total_bytes_loaded += consumed.bytes - total_copies += consumed.copies - total_copy_runs += consumed.copy_runs - total_ipc_ms += consumed.ipc_ms - total_copy_ms += consumed.copy_ms - total_transfer_open_ms += consumed.transfer_open_ms - total_transfer_load_ms += consumed.transfer_load_ms - total_transfer_sync_ms += consumed.transfer_sync_ms - total_l1_hits += consumed.l1_hits - total_l1_misses += consumed.l1_misses - total_l2_reads += consumed.l2_reads - - sync_start = time.perf_counter() - _synchronize_cuda_tensor(sample_tensor) - total_sync_ms += (time.perf_counter() - sync_start) * 1000 - logger.debug( - "[CONNECTOR] start_load_kv timing: reqs=%d batches=%d bytes=%d " - "copy_runs=%d gpu_copies=%d ipc_ms=%.3f copy_ms=%.3f " - "sync_ms=%.3f transfer_open_ms=%.3f transfer_load_ms=%.3f " - "transfer_sync_ms=%.3f l1_hits=%d l1_misses=%d l2_reads=%d " - "load_pipeline_depth=%d max_inflight_batches=%d " - "load_buffer_bytes=%d dynamic_load_allocations=%d", - len(reqs_to_load), - len(load_batches), - total_bytes_loaded, - total_copy_runs, - total_copies, - total_ipc_ms, - total_copy_ms, - total_sync_ms, - total_transfer_open_ms, - total_transfer_load_ms, - total_transfer_sync_ms, - total_l1_hits, - total_l1_misses, - total_l2_reads, - _LOAD_PIPELINE_DEPTH, - max_inflight_batches, - load_staging_pool.buffer_bytes, - 0, - ) + for consumed in consumed_batches: + total_bytes_loaded += consumed.bytes + total_copies += consumed.copies + total_copy_runs += consumed.copy_runs + total_ipc_ms += consumed.ipc_ms + total_copy_ms += consumed.copy_ms + total_transfer_open_ms += consumed.transfer_open_ms + total_transfer_load_ms += consumed.transfer_load_ms + total_transfer_sync_ms += consumed.transfer_sync_ms + total_l1_hits += consumed.l1_hits + total_l1_misses += consumed.l1_misses + total_l2_reads += consumed.l2_reads + + sync_start = time.perf_counter() + _synchronize_cuda_tensor(sample_tensor) + total_sync_ms += (time.perf_counter() - sync_start) * 1000 + + logger.debug( + "[CONNECTOR] start_load_kv timing: reqs=%d batches=%d bytes=%d " + "copy_runs=%d gpu_copies=%d ipc_ms=%.3f copy_ms=%.3f " + "sync_ms=%.3f transfer_open_ms=%.3f transfer_load_ms=%.3f " + "transfer_sync_ms=%.3f l1_hits=%d l1_misses=%d l2_reads=%d " + "load_pipeline_depth=%d max_inflight_batches=%d " + "load_buffer_bytes=%d dynamic_load_allocations=%d", + len(reqs_to_load), + len(load_batches), + total_bytes_loaded, + total_copy_runs, + total_copies, + total_ipc_ms, + total_copy_ms, + total_sync_ms, + total_transfer_open_ms, + total_transfer_load_ms, + total_transfer_sync_ms, + total_l1_hits, + total_l1_misses, + total_l2_reads, + _LOAD_PIPELINE_DEPTH, + max_inflight_batches, + load_staging_pool.buffer_bytes, + 0, + ) + except BaseException as exc: + for request_future in (request_futures or {}).values(): + if not request_future.done(): + request_future.set_exception(exc) + raise def _ensure_load_staging_pool( self, @@ -894,6 +962,69 @@ def _ensure_load_staging_pool( self._load_staging_pool = pool return pool + def _request_load_batch_counts( + self, + load_batches: Sequence[_LoadBatch], + ) -> dict[str, int]: + """Return how many load batches each request must wait for. + + Args: + load_batches: Bounded load batches whose ranges may include request + IDs when request-level completion is enabled. + + Returns: + Base request ID to number of batches containing that request. + + Async/thread-safety: + Pure CPU helper used on the connector load executor before the + completion-driven load pipeline starts. + """ + counts: dict[str, int] = {} + for _total_bytes, _spans, per_req_ranges in load_batches: + batch_req_ids: set[str] = set() + for item in per_req_ranges: + if len(item) < 4: + continue + _start, _end, req_id, _spec = item + batch_req_ids.add(base_req_id(str(req_id))) + for req_id in batch_req_ids: + counts[req_id] = counts.get(req_id, 0) + 1 + return counts + + def _mark_loaded_batch_requests( + self, + state: _InflightLoadBatch, + request_futures: dict[str, _RequestLoadFuture] | None, + remaining_batches: dict[str, int], + ) -> None: + """Mark request futures whose final load batch has been restored. + + Args: + state: Consumed load batch state. + request_futures: Base request ID to request completion future map. + remaining_batches: Mutable base request ID to outstanding batch + count map. + + Async/thread-safety: + Called on the connector load executor immediately after staging + bytes have been copied into the vLLM KV cache and synchronized. + """ + if not request_futures: + return + batch_req_ids: set[str] = set() + for item in state.per_req_ranges: + if len(item) < 4: + continue + _start, _end, req_id, _spec = item + batch_req_ids.add(base_req_id(str(req_id))) + for req_id in batch_req_ids: + remaining = max(0, remaining_batches.get(req_id, 1) - 1) + remaining_batches[req_id] = remaining + if remaining == 0: + future = request_futures.get(req_id) + if future is not None and not future.done(): + future.set_result() + def _submit_load_batch( self, batch: _LoadBatch, diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index 21595ad..affa133 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -70,6 +70,7 @@ WorkerConnectorMixin, _DeferredFinishedSave, _PendingLoad, + _RequestLoadFuture, ) BLOCK_TOKENS = 4 @@ -971,6 +972,135 @@ def __init__(self) -> None: assert probe._ipc_load_async.calls == [] # noqa: SLF001 +def test_group_load_marks_each_request_done_after_its_batch(monkeypatch) -> None: + """A request should become collectable once its load batch is restored.""" + + class Probe(WorkerConnectorMixin): + def __init__(self) -> None: + self._slot_size = 4 + self._store_staging_bytes = 4 + self._kv_caches = {"layer.0": torch.empty(4, 1, 1, 4, dtype=torch.uint8)} + self._layer_names = ["layer.0"] + self._load_key_scale = 1.0 + self._load_value_scale = 1.0 + self._rope_delta_scale = 1.0 + self._rope_base = 10000.0 + self._rope_rotary_dim = 0 + self._rope_is_neox_style = True + self._load_staging_pool = None + self._pending_loads = {} + self.mid_batch_checks = 0 + + def load_specs(self, reqs_to_load: dict[str, ReqLoadSpec]) -> None: + """Expose load execution through a public test helper.""" + request_futures = {req_id: _RequestLoadFuture() for req_id in reqs_to_load} + for req_id, spec in reqs_to_load.items(): + self._pending_loads[req_id] = _PendingLoad( + future=request_futures[req_id], + block_ids=list(spec.block_ids), + lease=None, + ) + self._load_kv_specs( + reqs_to_load, + next(iter(self._kv_caches.values())), + request_futures, + ) + + def _mark_loaded_batch_requests( + self, + state, + request_futures, + remaining_batches, + ) -> None: + super()._mark_loaded_batch_requests( + state, + request_futures, + remaining_batches, + ) + if state.total_bytes == 4: + self.mid_batch_checks += 1 + assert self._pending_loads["req-b"].future.done() + assert not self._pending_loads["req-a"].future.done() + + class _FakeFuture: + def __init__(self, total_bytes: int) -> None: + self.total_bytes = total_bytes + + def done(self) -> bool: + return completed_by_bytes[self.total_bytes] + + probe = Probe() + reqs_to_load = { + "req-a": ReqLoadSpec("hit-a", 0, 1, [0], 0, 4), + "req-b": ReqLoadSpec("hit-b", 1, 1, [1], 4, 4), + } + completed_by_bytes = {4: True} + consume_count = 0 + + def fake_submit(batch, buffer_index: int, sample_tensor: torch.Tensor): + del sample_tensor + total_bytes, _spans, per_req_ranges = batch + # Make req-b complete first even though it was submitted second. + total = 8 + buffer_index if per_req_ranges[0][3].chunk_key == "hit-a" else 4 + completed_by_bytes.setdefault(total, total == 4) + return SimpleNamespace( + total_bytes=total, + buffer_index=buffer_index, + per_req_ranges=per_req_ranges, + future=_FakeFuture(total), + ) + + def fake_consume(state, sample_tensor: torch.Tensor): + nonlocal consume_count + del sample_tensor + consume_count += 1 + if state.total_bytes == 4: + completed_by_bytes[8] = True + return SimpleNamespace( + buffer_index=state.buffer_index, + bytes=4, + copies=1, + copy_runs=1, + ipc_ms=0.0, + copy_ms=0.0, + transfer_open_ms=0.0, + transfer_load_ms=0.0, + transfer_sync_ms=0.0, + l1_hits=0, + l1_misses=0, + l2_reads=0, + ) + + monkeypatch.setattr(probe, "_submit_load_batch", fake_submit) + monkeypatch.setattr(probe, "_consume_loaded_batch", fake_consume) + monkeypatch.setattr( + "daser.connector.worker._synchronize_cuda_tensor", + lambda _: None, + ) + + probe.load_specs(reqs_to_load) + + assert consume_count == 2 + assert probe.mid_batch_checks == 1 + + +def test_get_finished_releases_each_request_load_independently() -> None: + """Completed request loads should not wait for unrelated pending loads.""" + connector = _AsyncLoadProbe() + done = _AsyncLoadFuture(done=True) + pending = _AsyncLoadFuture(done=False) + connector.seed_load("req-a", done, block_ids=[1]) + connector.seed_load("req-b", pending, block_ids=[2]) + + finished_sending, finished_recving = connector.get_finished(set()) + + assert finished_sending is None + assert finished_recving == {"req-a"} + assert done.result_calls == 1 + assert pending.result_calls == 0 + assert set(connector._pending_loads) == {"req-b"} # noqa: SLF001 + + def test_start_load_kv_releases_waiting_request_when_load_cannot_start() -> None: """Worker should not leave async-hit requests waiting forever.""" connector = _WorkerProbe("") From a918a8d42bf9eb10854bf88bccb3bce2ab5031c3 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Sat, 27 Jun 2026 13:32:37 +0800 Subject: [PATCH 09/11] fix(connector): complete split load request futures - map staging-local split load IDs back to their owning request before updating request-level load futures - cover split staging completion so async loads cannot leave requests waiting forever --- daser/connector/worker.py | 19 +++++++++++++++++-- tests/connector/test_daser_connector.py | 24 ++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/daser/connector/worker.py b/daser/connector/worker.py index 707c752..04aa440 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -202,6 +202,21 @@ def _store_staging_pool_depth(buffer_bytes: int, pending_limit_bytes: int) -> in return max(_MIN_STORE_STAGING_POOL_DEPTH, pending_limit_bytes // buffer_bytes) +def _load_completion_req_id(req_id: str) -> str: + """Return the vLLM request ID that owns a worker-side load range. + + Args: + req_id: Scheduler request ID or staging-local split ID. + + Returns: + Base vLLM request ID used for request-level load completion. + + Async/thread-safety: + Pure string helper. Safe to call from worker and executor threads. + """ + return base_req_id(req_id.split("#", 1)[0]) + + @dataclass class _DeferredFinishedSave: """Store work held until vLLM reports a request as finished.""" @@ -986,7 +1001,7 @@ def _request_load_batch_counts( if len(item) < 4: continue _start, _end, req_id, _spec = item - batch_req_ids.add(base_req_id(str(req_id))) + batch_req_ids.add(_load_completion_req_id(str(req_id))) for req_id in batch_req_ids: counts[req_id] = counts.get(req_id, 0) + 1 return counts @@ -1016,7 +1031,7 @@ def _mark_loaded_batch_requests( if len(item) < 4: continue _start, _end, req_id, _spec = item - batch_req_ids.add(base_req_id(str(req_id))) + batch_req_ids.add(_load_completion_req_id(str(req_id))) for req_id in batch_req_ids: remaining = max(0, remaining_batches.get(req_id, 1) - 1) remaining_batches[req_id] = remaining diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index affa133..b8203e2 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -1331,6 +1331,30 @@ def test_get_finished_reports_completed_async_load() -> None: assert connector._invalid_load_block_ids == set() # noqa: SLF001 +def test_request_load_completion_handles_split_staging_ids() -> None: + """Split staging batches should complete the original request future.""" + future = _RequestLoadFuture() + connector = _AsyncLoadProbe() + state = SimpleNamespace( + per_req_ranges=[ + ( + 0, + 4, + "req#0", + ReqLoadSpec("hit-a", 0, 1, [4], 0, 4), + ) + ] + ) + + connector._mark_loaded_batch_requests( # noqa: SLF001 + state, + {"req": future}, + {"req": 1}, + ) + + assert future.done() + + def test_get_finished_releases_request_after_async_load_failure() -> None: """Failed async loads should report completion and mark invalid blocks.""" connector = _AsyncLoadProbe() From 88b3bf4d27e120d2d6ab24d59e20b8b2097e78bc Mon Sep 17 00:00:00 2001 From: GentleCold Date: Sat, 27 Jun 2026 13:51:28 +0800 Subject: [PATCH 10/11] perf(connector): queue request-level loads - add a fixed-width worker load queue that groups request loads across scheduler windows - drive queued loads through the existing load pipeline with request-level completion futures - cover queued load grouping and pump batching behavior --- daser/connector/daser_connector.py | 3 + daser/connector/worker.py | 187 +++++++++++++++++++++--- tests/connector/test_daser_connector.py | 78 ++++++++++ 3 files changed, 251 insertions(+), 17 deletions(-) diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index 406d414..98cddcb 100644 --- a/daser/connector/daser_connector.py +++ b/daser/connector/daser_connector.py @@ -145,6 +145,9 @@ def __init__( self._pending_finished_saves: dict[str, Any] = {} self._pending_loads: dict[str, Any] = {} self._invalid_load_block_ids: set[int] = set() + self._load_request_queue = None + self._load_request_queue_lock = threading.Lock() + self._load_request_pump_future = None self._load_executor = ThreadPoolExecutor( max_workers=1, thread_name_prefix="daser-load", diff --git a/daser/connector/worker.py b/daser/connector/worker.py index 04aa440..662a680 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -10,6 +10,7 @@ from concurrent.futures import wait as wait_futures from dataclasses import dataclass import os +import queue import threading import time from typing import TYPE_CHECKING, Any @@ -86,6 +87,8 @@ _ROPE_WARMUP_BLOCKS = 1 _LOAD_PIPELINE_DEPTH = 2 +_LOAD_QUEUE_BATCH_SIZE = 2 +_LOAD_QUEUE_COALESCE_TIMEOUT_S = 0.002 _MIN_STORE_STAGING_POOL_DEPTH = 1 _LoadBatch = tuple[int, list[dict[str, int]], list[Any]] @@ -379,6 +382,78 @@ def set_exception(self, error: BaseException) -> None: self._event.set() +@dataclass +class _QueuedLoadRequest: + """One request waiting in the worker-side load queue. + + Attributes: + req_id: Base vLLM request ID used for completion. + spec_id: Scheduler metadata ID for this load spec. + spec: Load spec to restore into the vLLM KV cache. + future: Per-request completion future observed by ``get_finished``. + """ + + req_id: str + spec_id: str + spec: Any + future: _RequestLoadFuture + + +class LoadRequestQueue: + """Fixed-width worker-side queue for request-level KV loads. + + Args: + batch_size: Maximum number of queued requests grouped into one + ``_load_kv_specs`` invocation. + + Async/thread-safety: + ``put`` is called by vLLM worker threads while the pump runs on the + connector load executor. ``queue.Queue`` provides cross-thread + synchronization; ``stop`` unblocks the pump during shutdown. + """ + + def __init__( + self, + batch_size: int, + coalesce_timeout_s: float = _LOAD_QUEUE_COALESCE_TIMEOUT_S, + ) -> None: + self._batch_size = max(1, batch_size) + self._coalesce_timeout_s = max(0.0, coalesce_timeout_s) + self._items: queue.Queue[_QueuedLoadRequest | None] = queue.Queue() + self._stop_requested = threading.Event() + + def put(self, request: _QueuedLoadRequest) -> None: + """Enqueue one request-level load. + + Args: + request: Request load work item. + """ + self._items.put(request) + + def get_batch(self) -> list[_QueuedLoadRequest] | None: + """Return the next grouped request batch, or None after shutdown.""" + first = self._items.get() + if first is None: + return None + batch = [first] + while len(batch) < self._batch_size: + try: + item = self._items.get(timeout=self._coalesce_timeout_s) + except queue.Empty: + break + if item is None: + self._items.put(None) + break + batch.append(item) + return batch + + def stop(self) -> None: + """Request pump shutdown and unblock any pending ``get_batch``.""" + if not self._stop_requested.is_set(): + self._stop_requested.set() + self._items.put(None) + + def _cuda_allocation_base_and_offset(device_ptr: int) -> tuple[int, int]: """Return CUDA allocation base pointer and byte offset for ``device_ptr``. @@ -752,31 +827,28 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non self._mark_load_start_failed(reqs_to_load, "no registered KV cache tensor") return - request_futures = { - base_req_id(req_id): _RequestLoadFuture() for req_id in reqs_to_load - } - self._load_executor.submit( - self._load_kv_specs, - reqs_to_load, - sample_tensor, - request_futures, - ) pending_loads = getattr(self, "_pending_loads", None) if pending_loads is None: pending_loads = {} self._pending_loads = pending_loads - for req_id, request_future in request_futures.items(): - block_ids = [ - block_id - for spec_req_id, spec in reqs_to_load.items() - if base_req_id(spec_req_id) == req_id - for block_id in spec.block_ids - ] + load_queue = self._ensure_load_request_queue() + for spec_id, spec in reqs_to_load.items(): + req_id = base_req_id(spec_id) + request_future = _RequestLoadFuture() pending_loads[req_id] = _PendingLoad( future=request_future, - block_ids=list(block_ids), + block_ids=list(spec.block_ids), lease=None, ) + load_queue.put( + _QueuedLoadRequest( + req_id=req_id, + spec_id=spec_id, + spec=spec, + future=request_future, + ) + ) + self._ensure_load_request_pump(sample_tensor) def _mark_load_start_failed( self, @@ -810,6 +882,78 @@ def _mark_load_start_failed( lease=None, ) + def _ensure_load_request_queue(self) -> LoadRequestQueue: + """Return the worker-side request load queue, creating it if needed. + + Returns: + Queue used to group request loads across scheduler windows. + + Async/thread-safety: + Called on the vLLM worker thread. Attribute installation is guarded + by a small lock because multiple model-runner calls may race during + startup. + """ + load_queue = getattr(self, "_load_request_queue", None) + if load_queue is not None: + return load_queue + lock = getattr(self, "_load_request_queue_lock", None) + if lock is None: + lock = threading.Lock() + self._load_request_queue_lock = lock + with lock: + load_queue = getattr(self, "_load_request_queue", None) + if load_queue is None: + load_queue = LoadRequestQueue(_LOAD_QUEUE_BATCH_SIZE) + self._load_request_queue = load_queue + return load_queue + + def _ensure_load_request_pump(self, sample_tensor: torch.Tensor) -> None: + """Start the persistent load queue pump if it is not running. + + Args: + sample_tensor: Representative KV cache tensor used by load workers. + + Async/thread-safety: + Called on the vLLM worker thread. The pump itself runs on the + connector load executor and serially groups queued request loads + into fixed-size batches. + """ + pump_future = getattr(self, "_load_request_pump_future", None) + if pump_future is not None and not pump_future.done(): + return + lock = getattr(self, "_load_request_queue_lock", None) + if lock is None: + lock = threading.Lock() + self._load_request_queue_lock = lock + with lock: + pump_future = getattr(self, "_load_request_pump_future", None) + if pump_future is not None and not pump_future.done(): + return + self._load_request_pump_future = self._load_executor.submit( + self._run_load_request_queue, + sample_tensor, + ) + + def _run_load_request_queue(self, sample_tensor: torch.Tensor) -> None: + """Continuously group queued request loads and execute them. + + Args: + sample_tensor: Representative KV cache tensor used for load staging. + + Async/thread-safety: + Runs on the connector load executor. It owns all calls to + ``_load_kv_specs`` while active, preserving the existing single + executor-thread CUDA restore assumptions. + """ + load_queue = self._ensure_load_request_queue() + while True: + queued_batch = load_queue.get_batch() + if queued_batch is None: + return + reqs_to_load = {item.spec_id: item.spec for item in queued_batch} + request_futures = {item.req_id: item.future for item in queued_batch} + self._load_kv_specs(reqs_to_load, sample_tensor, request_futures) + def _load_kv_specs( self, reqs_to_load: dict[str, Any], @@ -1333,6 +1477,15 @@ def shutdown(self) -> None: for req_id in list(getattr(self, "_pending_finished_saves", {})): self.get_finished({req_id}) self._reap_save_futures(block=True) + load_queue = getattr(self, "_load_request_queue", None) + if load_queue is not None: + load_queue.stop() + pump_future = getattr(self, "_load_request_pump_future", None) + if pump_future is not None: + try: + pump_future.result(timeout=120.0) + except Exception: # noqa: BLE001 + pass load_executor = getattr(self, "_load_executor", None) if load_executor is not None: load_executor.shutdown(wait=True) diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index b8203e2..5a40490 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -67,6 +67,7 @@ synchronize_cuda_tensor as _synchronize_cuda_tensor, ) from daser.connector.worker import ( + LoadRequestQueue, WorkerConnectorMixin, _DeferredFinishedSave, _PendingLoad, @@ -160,6 +161,17 @@ def result(self, timeout: float | None = None) -> None: def _refresh_runtime_config(self) -> None: return + def _ensure_load_request_pump(self, sample_tensor: torch.Tensor) -> None: + """Drain one queued batch synchronously for worker probe tests.""" + queued_batch = self._load_request_queue.get_batch() + if queued_batch is None: + return + self._load_kv_specs( + {item.spec_id: item.spec for item in queued_batch}, + sample_tensor, + {item.req_id: item.future for item in queued_batch}, + ) + @property def transfer_ready(self): return self._transfer_ready @@ -1355,6 +1367,72 @@ def test_request_load_completion_handles_split_staging_ids() -> None: assert future.done() +def test_load_request_queue_groups_requests_across_puts() -> None: + """Worker-side load queue should group requests across scheduler windows.""" + queue = LoadRequestQueue(batch_size=2, coalesce_timeout_s=0.0) + first = SimpleNamespace(req_id="req-a") + second = SimpleNamespace(req_id="req-b") + + queue.put(first) + queue.put(second) + + batch = queue.get_batch() + + assert [item.req_id for item in batch or []] == ["req-a", "req-b"] + + +def test_load_request_queue_pump_batches_two_requests() -> None: + """Persistent load pump should pass two queued requests to one load call.""" + + class Probe(_AsyncLoadProbe): + def __init__(self) -> None: + super().__init__() + self._load_request_queue = LoadRequestQueue( + batch_size=2, + coalesce_timeout_s=0.0, + ) + self.load_calls: list[list[str]] = [] + + def _load_kv_specs( + self, + reqs_to_load, + sample_tensor, + request_futures=None, + ) -> None: + del sample_tensor + self.load_calls.append(list(reqs_to_load)) + for future in (request_futures or {}).values(): + future.set_result() + self._load_request_queue.stop() + + probe = Probe() + sample = torch.empty(1) + first_future = _RequestLoadFuture() + second_future = _RequestLoadFuture() + probe._load_request_queue.put( # noqa: SLF001 + SimpleNamespace( + req_id="req-a", + spec_id="req-a", + spec=ReqLoadSpec("hit-a", 0, 1, [1], 0, 4), + future=first_future, + ) + ) + probe._load_request_queue.put( # noqa: SLF001 + SimpleNamespace( + req_id="req-b", + spec_id="req-b", + spec=ReqLoadSpec("hit-b", 1, 1, [2], 4, 4), + future=second_future, + ) + ) + + probe._run_load_request_queue(sample) # noqa: SLF001 + + assert probe.load_calls == [["req-a", "req-b"]] + assert first_future.done() + assert second_future.done() + + def test_get_finished_releases_request_after_async_load_failure() -> None: """Failed async loads should report completion and mark invalid blocks.""" connector = _AsyncLoadProbe() From 5f8d5851844b3a172e0fff140d8b9413f44a413b Mon Sep 17 00:00:00 2001 From: GentleCold Date: Tue, 30 Jun 2026 05:06:25 +0800 Subject: [PATCH 11/11] perf(connector): dispatch request loads by inflight slots - replace grouped load queue pumping with request-level max-inflight dispatch - bound active request loads by fixed staging pool depth and configured inflight limit - expand fixed load staging registration and IPC lanes for request-level dispatch - cover immediate dispatch and staging-depth limiting in connector tests --- daser/connector/daser_connector.py | 12 +- daser/connector/staging.py | 5 + daser/connector/worker.py | 477 +++++++++++++++++++----- tests/connector/test_daser_connector.py | 158 ++++---- 4 files changed, 487 insertions(+), 165 deletions(-) diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index 98cddcb..659780b 100644 --- a/daser/connector/daser_connector.py +++ b/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 _LOAD_PIPELINE_DEPTH, WorkerConnectorMixin +from daser.connector.worker import _LOAD_REQUEST_MAX_INFLIGHT, WorkerConnectorMixin from daser.logging import init_logger logger = init_logger(__name__) @@ -128,7 +127,7 @@ def __init__( self._ipc_load_async, *[ IPCClientAsync(self._socket_path) - for _ in range(max(0, _LOAD_PIPELINE_DEPTH - 1)) + for _ in range(max(0, _LOAD_REQUEST_MAX_INFLIGHT - 1)) ], ] self._ipc_store_async = IPCClientAsync(self._socket_path) @@ -146,12 +145,9 @@ def __init__( self._pending_loads: dict[str, Any] = {} self._invalid_load_block_ids: set[int] = set() self._load_request_queue = None + self._load_request_dispatcher = None self._load_request_queue_lock = threading.Lock() - self._load_request_pump_future = None - self._load_executor = ThreadPoolExecutor( - max_workers=1, - thread_name_prefix="daser-load", - ) + 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( diff --git a/daser/connector/staging.py b/daser/connector/staging.py index 61db637..4394394 100644 --- a/daser/connector/staging.py +++ b/daser/connector/staging.py @@ -140,6 +140,11 @@ def available(self) -> int: """Return the number of currently free staging buffers.""" return len(self._free_indices) + @property + def depth(self) -> int: + """Return the number of fixed staging buffers in this pool.""" + return len(self._buffers) + def buffer(self, index: int) -> torch.Tensor: """Return a fixed staging backing tensor by index. diff --git a/daser/connector/worker.py b/daser/connector/worker.py index 662a680..42e420a 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -87,8 +87,9 @@ _ROPE_WARMUP_BLOCKS = 1 _LOAD_PIPELINE_DEPTH = 2 -_LOAD_QUEUE_BATCH_SIZE = 2 -_LOAD_QUEUE_COALESCE_TIMEOUT_S = 0.002 +_LOAD_REQUEST_MAX_INFLIGHT = 8 +_LOAD_DISPATCH_WAIT_TIMEOUT_S = 0.001 +_LOAD_STAGING_RESERVE_BYTES = 1 << 30 _MIN_STORE_STAGING_POOL_DEPTH = 1 _LoadBatch = tuple[int, list[dict[str, int]], list[Any]] @@ -205,6 +206,40 @@ def _store_staging_pool_depth(buffer_bytes: int, pending_limit_bytes: int) -> in return max(_MIN_STORE_STAGING_POOL_DEPTH, pending_limit_bytes // buffer_bytes) +def _load_staging_pool_depth( + buffer_bytes: int, + pending_limit_bytes: int, + device: torch.device, +) -> int: + """Return fixed load staging depth under memory and inflight constraints. + + Args: + buffer_bytes: Capacity of one fixed staging buffer. + pending_limit_bytes: Existing staging byte budget from store limits. + device: CUDA device used for staging allocation. + + Returns: + Number of fixed load staging buffers to preallocate. + + Async/thread-safety: + Pure helper except for querying CUDA free memory. Called during worker + initialization before request traffic. + """ + depth = min( + _LOAD_REQUEST_MAX_INFLIGHT, + _store_staging_pool_depth(buffer_bytes, pending_limit_bytes), + ) + if device.type != "cuda": + return max(1, depth) + try: + free_bytes, _total_bytes = torch.cuda.mem_get_info(device) + except Exception: # noqa: BLE001 + return max(1, depth) + usable_bytes = max(0, int(free_bytes) - _LOAD_STAGING_RESERVE_BYTES) + memory_depth = max(1, usable_bytes // buffer_bytes) + return max(1, min(depth, memory_depth)) + + def _load_completion_req_id(req_id: str) -> str: """Return the vLLM request ID that owns a worker-side load range. @@ -293,6 +328,29 @@ class _InflightLoadBatch: submitted_at: float +@dataclass +class _InflightRequestLoad: + """One request-level load submitted by the dispatcher. + + Attributes: + item: Original queued request item. + buffer_index: Fixed staging buffer owned by this request state. + batches: Bounded read batches for this request. + next_batch: Index of the next unsubmitted read batch. + remaining_batches: Number of read batches not yet consumed. + active: Current in-flight read batch for this request, if any. + completed: Consumed batch timing rows for logging. + """ + + item: "_QueuedLoadRequest" + buffer_index: int + batches: list[_LoadBatch] + next_batch: int + remaining_batches: int + active: _InflightLoadBatch | None + completed: list["_ConsumedLoadBatch"] + + @dataclass class _ConsumedLoadBatch: """Timing and accounting data for one restored load batch.""" @@ -399,59 +457,98 @@ class _QueuedLoadRequest: future: _RequestLoadFuture -class LoadRequestQueue: - """Fixed-width worker-side queue for request-level KV loads. +class LoadRequestDispatcher: + """Bound request-level load concurrency by max in-flight and staging depth. Args: - batch_size: Maximum number of queued requests grouped into one - ``_load_kv_specs`` invocation. + max_inflight: Maximum request loads allowed in flight. + staging_depth: Number of fixed staging buffers available to request + loads. Async/thread-safety: - ``put`` is called by vLLM worker threads while the pump runs on the - connector load executor. ``queue.Queue`` provides cross-thread - synchronization; ``stop`` unblocks the pump during shutdown. + The dispatcher object is owned by the connector load asyncio loop. Test + helpers may call its pure synchronous scheduling helpers directly. """ def __init__( self, - batch_size: int, - coalesce_timeout_s: float = _LOAD_QUEUE_COALESCE_TIMEOUT_S, + max_inflight: int, + staging_depth: int, ) -> None: - self._batch_size = max(1, batch_size) - self._coalesce_timeout_s = max(0.0, coalesce_timeout_s) - self._items: queue.Queue[_QueuedLoadRequest | None] = queue.Queue() - self._stop_requested = threading.Event() + self._effective_inflight = max(1, min(max_inflight, staging_depth)) + self._free_buffers: deque[int] = deque(range(self._effective_inflight)) - def put(self, request: _QueuedLoadRequest) -> None: - """Enqueue one request-level load. + @property + def effective_inflight(self) -> int: + """Return the active request limit after applying staging depth.""" + return self._effective_inflight + + def submit_ready( + self, + connector: Any, + queued: list[_QueuedLoadRequest], + sample_tensor: torch.Tensor, + ) -> list[_InflightRequestLoad]: + """Submit queued requests while in-flight slots and buffers are free. Args: - request: Request load work item. + connector: Worker connector that owns submit helpers. + queued: Mutable FIFO list of queued request work. + sample_tensor: Representative KV cache tensor. + + Returns: + Newly submitted request states. """ - self._items.put(request) + submitted: list[_InflightRequestLoad] = [] + while queued and self._free_buffers: + buffer_index = self._free_buffers.popleft() + item = queued.pop(0) + state = connector._submit_request_load_for_dispatcher( # noqa: SLF001 + item, + buffer_index, + sample_tensor, + ) + if state.active is None: + self._free_buffers.append(state.buffer_index) + continue + submitted.append(state) + return submitted - def get_batch(self) -> list[_QueuedLoadRequest] | None: - """Return the next grouped request batch, or None after shutdown.""" - first = self._items.get() - if first is None: - return None - batch = [first] - while len(batch) < self._batch_size: - try: - item = self._items.get(timeout=self._coalesce_timeout_s) - except queue.Empty: - break - if item is None: - self._items.put(None) - break - batch.append(item) - return batch + def consume_ready( + self, + connector: Any, + active: list[_InflightRequestLoad], + sample_tensor: torch.Tensor, + ) -> list[_InflightRequestLoad]: + """Consume completed request loads and release their buffers. + + Args: + connector: Worker connector that owns consume helpers. + active: Mutable list of active request load states. + sample_tensor: Representative KV cache tensor. - def stop(self) -> None: - """Request pump shutdown and unblock any pending ``get_batch``.""" - if not self._stop_requested.is_set(): - self._stop_requested.set() - self._items.put(None) + Returns: + States consumed during this call. + """ + consumed: list[_InflightRequestLoad] = [] + for state in list(active): + active_batch = state.active + if active_batch is None: + active.remove(state) + self._free_buffers.append(state.buffer_index) + consumed.append(state) + continue + if not active_batch.future.done(): + continue + reusable_buffer, request_done = connector._consume_dispatcher_load( # noqa: SLF001 + state, + sample_tensor, + ) + if request_done: + active.remove(state) + self._free_buffers.append(reusable_buffer) + consumed.append(state) + return consumed def _cuda_allocation_base_and_offset(device_ptr: int) -> tuple[int, int]: @@ -658,16 +755,21 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: self._load_staging_pool = FixedCudaStagingPool( device=sample.device, buffer_bytes=self._store_staging_bytes, - depth=_LOAD_PIPELINE_DEPTH, + depth=_load_staging_pool_depth( + self._store_staging_bytes, + self._pending_store_staging_limit_bytes, + sample.device, + ), ) self._load_staging_registered = False logger.info( "[CONNECTOR] preallocated staging buffer=%d cap=%d pending=%d " - "load_pipeline_depth=%d", + "load_request_max_inflight=%d load_staging_depth=%d", self._store_staging_bytes, self._store_staging_bytes, self._pending_store_staging_limit_bytes, - _LOAD_PIPELINE_DEPTH, + _LOAD_REQUEST_MAX_INFLIGHT, + self._load_staging_pool.depth, ) if sample.dim() >= 5: _warm_rope_apply_backends( @@ -743,16 +845,21 @@ def register_cross_layers_kv_cache( self._load_staging_pool = FixedCudaStagingPool( device=kv_cache.device, buffer_bytes=self._store_staging_bytes, - depth=_LOAD_PIPELINE_DEPTH, + depth=_load_staging_pool_depth( + self._store_staging_bytes, + self._pending_store_staging_limit_bytes, + kv_cache.device, + ), ) self._load_staging_registered = False logger.info( "[CONNECTOR] register_cross_layers_kv_cache: layers=%d shape=%s " - "dtype=%s load_pipeline_depth=%d", + "dtype=%s load_request_max_inflight=%d load_staging_depth=%d", len(self._layer_names), tuple(kv_cache.shape), kv_cache.dtype, - _LOAD_PIPELINE_DEPTH, + _LOAD_REQUEST_MAX_INFLIGHT, + self._load_staging_pool.depth, ) _warm_rope_apply_backends( device=kv_cache.device, @@ -831,7 +938,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non if pending_loads is None: pending_loads = {} self._pending_loads = pending_loads - load_queue = self._ensure_load_request_queue() + load_queue = self._ensure_load_request_queue(sample_tensor) for spec_id, spec in reqs_to_load.items(): req_id = base_req_id(spec_id) request_future = _RequestLoadFuture() @@ -840,15 +947,16 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non block_ids=list(spec.block_ids), lease=None, ) - load_queue.put( + self._enqueue_load_request( + load_queue, _QueuedLoadRequest( req_id=req_id, spec_id=spec_id, spec=spec, future=request_future, - ) + ), ) - self._ensure_load_request_pump(sample_tensor) + self._ensure_load_request_dispatcher(sample_tensor) def _mark_load_start_failed( self, @@ -882,7 +990,9 @@ def _mark_load_start_failed( lease=None, ) - def _ensure_load_request_queue(self) -> LoadRequestQueue: + def _ensure_load_request_queue( + self, sample_tensor: torch.Tensor | None = None + ) -> Any: """Return the worker-side request load queue, creating it if needed. Returns: @@ -903,56 +1013,155 @@ def _ensure_load_request_queue(self) -> LoadRequestQueue: with lock: load_queue = getattr(self, "_load_request_queue", None) if load_queue is None: - load_queue = LoadRequestQueue(_LOAD_QUEUE_BATCH_SIZE) + if sample_tensor is not None and hasattr(self, "_load_loop"): + future = asyncio.run_coroutine_threadsafe( + self._create_load_request_queue(), + self._load_loop, + ) + load_queue = future.result(timeout=10.0) + else: + load_queue = queue.Queue() self._load_request_queue = load_queue return load_queue - def _ensure_load_request_pump(self, sample_tensor: torch.Tensor) -> None: - """Start the persistent load queue pump if it is not running. + async def _create_load_request_queue(self) -> asyncio.Queue[Any]: + """Create an asyncio request queue on the load event loop.""" + return asyncio.Queue() + + def _enqueue_load_request( + self, + load_queue: Any, + request: _QueuedLoadRequest, + ) -> None: + """Enqueue one request from a worker thread into the load queue. + + Args: + load_queue: Queue created by ``_ensure_load_request_queue``. + request: Request-level load work item. + + Async/thread-safety: + Called on vLLM worker threads. Production queues are asyncio queues + owned by ``_load_loop``; test queues may be synchronous ``queue.Queue`` + instances. + """ + if isinstance(load_queue, asyncio.Queue): + self._load_loop.call_soon_threadsafe(load_queue.put_nowait, request) + else: + load_queue.put(request) + + def _ensure_load_request_dispatcher(self, sample_tensor: torch.Tensor) -> None: + """Start the persistent request load dispatcher if it is not running. Args: sample_tensor: Representative KV cache tensor used by load workers. Async/thread-safety: - Called on the vLLM worker thread. The pump itself runs on the - connector load executor and serially groups queued request loads - into fixed-size batches. + Called on the vLLM worker thread. The dispatcher itself runs on the + connector load asyncio loop and keeps request-level loads in flight + up to the fixed max-inflight/staging-pool limit. """ - pump_future = getattr(self, "_load_request_pump_future", None) - if pump_future is not None and not pump_future.done(): + dispatcher_future = getattr(self, "_load_request_dispatcher_future", None) + if dispatcher_future is not None and not dispatcher_future.done(): return lock = getattr(self, "_load_request_queue_lock", None) if lock is None: lock = threading.Lock() self._load_request_queue_lock = lock with lock: - pump_future = getattr(self, "_load_request_pump_future", None) - if pump_future is not None and not pump_future.done(): + dispatcher_future = getattr(self, "_load_request_dispatcher_future", None) + if dispatcher_future is not None and not dispatcher_future.done(): return - self._load_request_pump_future = self._load_executor.submit( - self._run_load_request_queue, - sample_tensor, + self._load_request_dispatcher_future = asyncio.run_coroutine_threadsafe( + self._run_load_request_dispatcher(sample_tensor), + self._load_loop, ) - def _run_load_request_queue(self, sample_tensor: torch.Tensor) -> None: - """Continuously group queued request loads and execute them. + async def _run_load_request_dispatcher(self, sample_tensor: torch.Tensor) -> None: + """Continuously dispatch request loads as in-flight slots become free. Args: sample_tensor: Representative KV cache tensor used for load staging. Async/thread-safety: - Runs on the connector load executor. It owns all calls to - ``_load_kv_specs`` while active, preserving the existing single - executor-thread CUDA restore assumptions. + Runs on the connector load asyncio loop. The dispatcher never waits + to coalesce requests; it submits each queued request immediately when + both an in-flight slot and a fixed staging buffer are available. """ - load_queue = self._ensure_load_request_queue() + load_queue = self._ensure_load_request_queue(sample_tensor) + load_staging_pool = self._ensure_load_staging_pool(sample_tensor) + dispatcher = LoadRequestDispatcher( + max_inflight=_LOAD_REQUEST_MAX_INFLIGHT, + staging_depth=load_staging_pool.depth, + ) + self._load_request_dispatcher = dispatcher + queued: list[_QueuedLoadRequest] = [] + active: list[_InflightRequestLoad] = [] + try: + while True: + if not queued: + if active: + await self._drain_load_queue(load_queue, queued) + else: + item = await load_queue.get() + if item is None: + return + queued.append(item) + active.extend(dispatcher.submit_ready(self, queued, sample_tensor)) + consumed = dispatcher.consume_ready(self, active, sample_tensor) + if consumed: + continue + if not active: + continue + await self._wait_for_dispatcher_completion(active) + except BaseException as exc: + for state in active: + if not state.item.future.done(): + state.item.future.set_exception(exc) + active_batch = state.active + if active_batch is not None: + active_batch.staging_lease.release() + for item in queued: + if not item.future.done(): + item.future.set_exception(exc) + raise + + async def _drain_load_queue( + self, + load_queue: asyncio.Queue[Any], + queued: list[_QueuedLoadRequest], + ) -> None: + """Move immediately available queue items into a local FIFO list.""" while True: - queued_batch = load_queue.get_batch() - if queued_batch is None: + try: + item = load_queue.get_nowait() + except asyncio.QueueEmpty: return - reqs_to_load = {item.spec_id: item.spec for item in queued_batch} - request_futures = {item.req_id: item.future for item in queued_batch} - self._load_kv_specs(reqs_to_load, sample_tensor, request_futures) + if item is None: + await load_queue.put(None) + return + queued.append(item) + + async def _wait_for_dispatcher_completion( + self, + active: list[_InflightRequestLoad], + ) -> None: + """Wait until any request-level active read future completes.""" + wrapped: dict[asyncio.Future[Any], _InflightRequestLoad] = {} + for state in active: + active_batch = state.active + if active_batch is None: + continue + wrapped[asyncio.wrap_future(active_batch.future)] = state + if not wrapped: + await asyncio.sleep(0) + return + done, _pending = await asyncio.wait( + wrapped.keys(), + timeout=_LOAD_DISPATCH_WAIT_TIMEOUT_S, + return_when=asyncio.FIRST_COMPLETED, + ) + if not done: + await asyncio.sleep(0) def _load_kv_specs( self, @@ -1116,7 +1325,11 @@ def _ensure_load_staging_pool( pool = FixedCudaStagingPool( device=sample_tensor.device, buffer_bytes=buffer_bytes, - depth=_LOAD_PIPELINE_DEPTH, + depth=_load_staging_pool_depth( + buffer_bytes, + getattr(self, "_pending_store_staging_limit_bytes", 0), + sample_tensor.device, + ), ) self._load_staging_pool = pool return pool @@ -1244,6 +1457,100 @@ def _submit_load_batch( submitted_at=submitted_at, ) + def _submit_request_load_for_dispatcher( + self, + item: _QueuedLoadRequest, + buffer_index: int, + sample_tensor: torch.Tensor, + ) -> _InflightRequestLoad: + """Submit the first read batch for one queued request load. + + Args: + item: Request-level queued load. + buffer_index: Fixed staging buffer index reserved by the dispatcher. + sample_tensor: Representative KV cache tensor. + + Returns: + Request state tracked by ``LoadRequestDispatcher``. + + Async/thread-safety: + Runs on the connector load asyncio loop. Each request owns at most one + active staging buffer at a time; multi-segment requests submit the + next segment only after the previous segment has been restored. + """ + load_staging_pool = self._ensure_load_staging_pool(sample_tensor) + load_batches = _build_load_read_batches( + {item.spec_id: item.spec}, + self._slot_size, + max_batch_bytes=load_staging_pool.buffer_bytes, + include_req_ids=True, + ) + if not load_batches: + item.future.set_result() + return _InflightRequestLoad( + item=item, + buffer_index=buffer_index, + batches=[], + next_batch=0, + remaining_batches=0, + active=None, + completed=[], + ) + active_batch = self._submit_load_batch( + load_batches[0], + buffer_index, + sample_tensor, + ) + return _InflightRequestLoad( + item=item, + buffer_index=buffer_index, + batches=load_batches, + next_batch=1, + remaining_batches=len(load_batches), + active=active_batch, + completed=[], + ) + + def _consume_dispatcher_load( + self, + state: _InflightRequestLoad, + sample_tensor: torch.Tensor, + ) -> tuple[int, bool]: + """Consume one completed request read and finish or advance the request. + + Args: + state: Request-level load state with a completed active batch. + sample_tensor: Representative KV cache tensor. + + Returns: + Tuple of released staging buffer index and whether the request is + fully complete. + + Async/thread-safety: + Runs on the connector load asyncio loop after the associated transfer + future is complete. The staging buffer is not reused until restore + kernels have synchronized and the lease has been released. + """ + active_batch = state.active + if active_batch is None: + return state.buffer_index, True + consumed = self._consume_loaded_batch(active_batch, sample_tensor) + state.completed.append(consumed) + state.remaining_batches = max(0, state.remaining_batches - 1) + reusable_buffer = consumed.buffer_index + state.buffer_index = reusable_buffer + if state.remaining_batches == 0: + if not state.item.future.done(): + state.item.future.set_result() + return reusable_buffer, True + state.active = self._submit_load_batch( + state.batches[state.next_batch], + reusable_buffer, + sample_tensor, + ) + state.next_batch += 1 + return reusable_buffer, False + def _consume_loaded_batch( self, state: _InflightLoadBatch, @@ -1479,16 +1786,16 @@ def shutdown(self) -> None: self._reap_save_futures(block=True) load_queue = getattr(self, "_load_request_queue", None) if load_queue is not None: - load_queue.stop() - pump_future = getattr(self, "_load_request_pump_future", None) - if pump_future is not None: + if isinstance(load_queue, asyncio.Queue): + self._load_loop.call_soon_threadsafe(load_queue.put_nowait, None) + else: + load_queue.put(None) + dispatcher_future = getattr(self, "_load_request_dispatcher_future", None) + if dispatcher_future is not None: try: - pump_future.result(timeout=120.0) + dispatcher_future.result(timeout=120.0) except Exception: # noqa: BLE001 pass - load_executor = getattr(self, "_load_executor", None) - if load_executor is not None: - load_executor.shutdown(wait=True) load_clients = list( dict.fromkeys( getattr(self, "_ipc_load_async_pool", []) @@ -1644,7 +1951,7 @@ def _register_load_staging_buffers(self) -> None: if pool is None or not self._load_ipc_clients(): return try: - for buffer_index in range(_LOAD_PIPELINE_DEPTH): + for buffer_index in range(pool.depth): tensor = pool.buffer(buffer_index) cp_tensor = cupy.asarray(tensor) device_ptr = cuda_array_pointer(cp_tensor) @@ -1673,7 +1980,7 @@ def _register_load_staging_buffers(self) -> None: self._load_staging_registered = True logger.info( "[CONNECTOR] registered %d fixed load staging buffers", - _LOAD_PIPELINE_DEPTH, + pool.depth, ) def _reap_save_futures(self, block: bool) -> None: diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index 5a40490..376e781 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -67,9 +67,11 @@ synchronize_cuda_tensor as _synchronize_cuda_tensor, ) from daser.connector.worker import ( - LoadRequestQueue, + LoadRequestDispatcher, WorkerConnectorMixin, _DeferredFinishedSave, + _InflightRequestLoad, + _load_staging_pool_depth, _PendingLoad, _RequestLoadFuture, ) @@ -143,34 +145,20 @@ def __init__(self, store_path: str) -> None: self._save_futures = [] self._pending_save_staging_bytes = 0 - class _InlineExecutor: - def submit(self, fn, *args): - fn(*args) - - class _DoneFuture: - def done(self) -> bool: - return True - - def result(self, timeout: float | None = None) -> None: - del timeout - - return _DoneFuture() - - self._load_executor = _InlineExecutor() - def _refresh_runtime_config(self) -> None: return - def _ensure_load_request_pump(self, sample_tensor: torch.Tensor) -> None: - """Drain one queued batch synchronously for worker probe tests.""" - queued_batch = self._load_request_queue.get_batch() - if queued_batch is None: - return - self._load_kv_specs( - {item.spec_id: item.spec for item in queued_batch}, - sample_tensor, - {item.req_id: item.future for item in queued_batch}, - ) + def _ensure_load_request_dispatcher(self, sample_tensor: torch.Tensor) -> None: + """Drain queued requests synchronously for worker probe tests.""" + while not self._load_request_queue.empty(): + item = self._load_request_queue.get() + if item is None: + return + self._load_kv_specs( + {item.spec_id: item.spec}, + sample_tensor, + {item.req_id: item.future}, + ) @property def transfer_ready(self): @@ -534,6 +522,10 @@ def result(self, timeout: float | None = None) -> None: if self._error is not None: raise self._error + def mark_done(self) -> None: + """Mark this fake future as complete.""" + self._done = True + class _AsyncLoadProbe(WorkerConnectorMixin): """Worker probe with pending async load futures.""" @@ -545,22 +537,11 @@ def __init__(self) -> None: self._save_futures = [] self._invalid_load_block_ids = set() self.released: list[object] = [] - self.load_executor_shutdown = False self.load_loop_stopped = False self.store_loop_stopped = False self.load_thread_joined = False self.store_thread_joined = False - class _Executor: - def __init__(self, probe: _AsyncLoadProbe) -> None: - self._probe = probe - - def shutdown(self, wait: bool) -> None: - del wait - self._probe.load_executor_shutdown = True - - self._load_executor = _Executor(self) - class _Loop: def __init__(self, callback) -> None: self._callback = callback @@ -1367,70 +1348,104 @@ def test_request_load_completion_handles_split_staging_ids() -> None: assert future.done() -def test_load_request_queue_groups_requests_across_puts() -> None: - """Worker-side load queue should group requests across scheduler windows.""" - queue = LoadRequestQueue(batch_size=2, coalesce_timeout_s=0.0) - first = SimpleNamespace(req_id="req-a") - second = SimpleNamespace(req_id="req-b") +def test_load_request_dispatcher_limits_active_requests_by_staging_depth() -> None: + """Request dispatcher should cap active loads by max inflight and buffers.""" + dispatcher = LoadRequestDispatcher(max_inflight=8, staging_depth=3) - queue.put(first) - queue.put(second) + assert dispatcher.effective_inflight == 3 - batch = queue.get_batch() - assert [item.req_id for item in batch or []] == ["req-a", "req-b"] +def test_load_staging_pool_depth_respects_available_cuda_memory(monkeypatch) -> None: + """Load staging depth should not preallocate past available CUDA headroom.""" + monkeypatch.setattr( + torch.cuda, + "mem_get_info", + lambda device=None: ((4 << 30), 80 << 30), + ) + depth = _load_staging_pool_depth( + buffer_bytes=1536 << 20, + pending_limit_bytes=12 << 30, + device=torch.device("cuda"), + ) -def test_load_request_queue_pump_batches_two_requests() -> None: - """Persistent load pump should pass two queued requests to one load call.""" + assert depth == 2 + + +def test_load_request_dispatcher_submits_without_waiting_for_batch() -> None: + """Queued requests should be submitted immediately while slots are free.""" class Probe(_AsyncLoadProbe): def __init__(self) -> None: super().__init__() - self._load_request_queue = LoadRequestQueue( - batch_size=2, - coalesce_timeout_s=0.0, + self.submitted: list[str] = [] + self.completed: list[str] = [] + + def _submit_request_load_for_dispatcher( + self, item, buffer_index, sample_tensor + ): + del sample_tensor + self.submitted.append(item.req_id) + future = _AsyncLoadFuture(done=False) + active = SimpleNamespace(future=future) + return _InflightRequestLoad( + item=item, + buffer_index=buffer_index, + batches=[], + next_batch=0, + remaining_batches=1, + active=active, + completed=[], ) - self.load_calls: list[list[str]] = [] - def _load_kv_specs( - self, - reqs_to_load, - sample_tensor, - request_futures=None, - ) -> None: + def _consume_dispatcher_load(self, state, sample_tensor): del sample_tensor - self.load_calls.append(list(reqs_to_load)) - for future in (request_futures or {}).values(): - future.set_result() - self._load_request_queue.stop() + self.completed.append(state.item.req_id) + state.item.future.set_result() + return 0, True probe = Probe() + dispatcher = LoadRequestDispatcher(max_inflight=8, staging_depth=2) sample = torch.empty(1) first_future = _RequestLoadFuture() second_future = _RequestLoadFuture() - probe._load_request_queue.put( # noqa: SLF001 + third_future = _RequestLoadFuture() + queued = [ SimpleNamespace( req_id="req-a", spec_id="req-a", spec=ReqLoadSpec("hit-a", 0, 1, [1], 0, 4), future=first_future, - ) - ) - probe._load_request_queue.put( # noqa: SLF001 + ), SimpleNamespace( req_id="req-b", spec_id="req-b", spec=ReqLoadSpec("hit-b", 1, 1, [2], 4, 4), future=second_future, - ) - ) + ), + SimpleNamespace( + req_id="req-c", + spec_id="req-c", + spec=ReqLoadSpec("hit-c", 2, 1, [3], 8, 4), + future=third_future, + ), + ] + + active = dispatcher.submit_ready(probe, queued, sample) + + assert probe.submitted == ["req-a", "req-b"] + assert [item.req_id for item in queued] == ["req-c"] + assert len(active) == 2 - probe._run_load_request_queue(sample) # noqa: SLF001 + active_batch = active[0].active + assert active_batch is not None + active_batch.future.mark_done() + dispatcher.consume_ready(probe, active, sample) + dispatcher.submit_ready(probe, queued, sample) - assert probe.load_calls == [["req-a", "req-b"]] + assert probe.completed == ["req-a"] + assert probe.submitted == ["req-a", "req-b", "req-c"] assert first_future.done() - assert second_future.done() def test_get_finished_releases_request_after_async_load_failure() -> None: @@ -1457,7 +1472,6 @@ def test_shutdown_collects_failed_async_load_without_raising() -> None: assert connector._pending_loads == {} # noqa: SLF001 assert connector._invalid_load_block_ids == {4, 5} # noqa: SLF001 - assert connector.load_executor_shutdown assert connector.load_loop_stopped assert connector.store_loop_stopped assert connector.load_thread_joined