From 4acfa519bb2ddc3a0b6633ccc28f74206156a5a1 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Sun, 12 Jul 2026 15:36:44 +0800 Subject: [PATCH 1/4] refactor(connector): deepen execution pipeline ownership - isolate worker load, store, memory, and runtime responsibilities - centralize scheduler request lifecycle and immutable store intents - narrow staging to tensor transforms and document ownership flows - preserve TP2 cold and warm cache behavior with regression coverage --- daser/connector/daser_connector.py | 95 +- daser/connector/load_pipeline.py | 324 ++++ daser/connector/request_lifecycle.py | 770 ++++++++ daser/connector/reuse.py | 170 +- daser/connector/scheduler.py | 946 +--------- daser/connector/scheduler_planning.py | 296 +++ daser/connector/staging.py | 548 ------ daser/connector/store_pipeline.py | 145 ++ daser/connector/worker.py | 2162 +--------------------- daser/connector/worker_memory.py | 235 +++ daser/connector/worker_runtime.py | 2156 +++++++++++++++++++++ docs/design/architecture.md | 30 +- docs/design/components.md | 10 +- docs/design/flows.md | 80 +- tests/connector/test_daser_connector.py | 443 +++-- tests/connector/test_gds_transfer.py | 6 +- tests/unit/test_block_aligned_bug.py | 17 +- tests/unit/test_chunk_reuse_scheduler.py | 4 +- tests/unit/test_skip_save_flag.py | 7 +- 19 files changed, 4502 insertions(+), 3942 deletions(-) create mode 100644 daser/connector/load_pipeline.py create mode 100644 daser/connector/request_lifecycle.py create mode 100644 daser/connector/scheduler_planning.py create mode 100644 daser/connector/store_pipeline.py create mode 100644 daser/connector/worker_memory.py create mode 100644 daser/connector/worker_runtime.py diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index 32e00dd..e43cbaa 100644 --- a/daser/connector/daser_connector.py +++ b/daser/connector/daser_connector.py @@ -1,12 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # Standard -import asyncio -import threading from typing import TYPE_CHECKING, Any # Third Party -import torch from vllm.distributed import get_tensor_model_parallel_rank from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, @@ -18,9 +15,12 @@ from vllm.config import VllmConfig # First Party -from daser.connector.helpers import PendingStore -from daser.connector.ipc_client import IPCClientAsync, IPCClientSync +from daser.connector.ipc_client import IPCClientSync +from daser.connector.load_pipeline import ( + build_load_read_plan as _build_load_read_plan, +) from daser.connector.metadata import DaserConnectorMeta, ReqLoadSpec, ReqStoreSpec +from daser.connector.request_lifecycle import RequestLifecycle from daser.connector.reuse import CacheReuseStrategy, build_cache_reuse_strategy from daser.connector.scheduler import ( SchedulerConnectorMixin, @@ -34,13 +34,11 @@ from daser.connector.staging import ( apply_rope_delta_to_key_block as _apply_rope_delta_to_key_block, ) -from daser.connector.staging import ( - build_load_read_plan as _build_load_read_plan, -) from daser.connector.staging import ( copy_staging_to_kv_cache as _copy_staging_to_kv_cache, ) -from daser.connector.worker import _LOAD_REQUEST_MAX_INFLIGHT, WorkerConnectorMixin +from daser.connector.worker import WorkerConnectorMixin +from daser.connector.worker_runtime import WorkerRuntime from daser.logging import init_logger logger = init_logger(__name__) @@ -109,8 +107,9 @@ def __init__( self._block_tokens: int = 16 self._model_id: str = "default" self._skip_l2: bool = bool(extra.get("skip_l2", False)) + cache_reuse_mode = str(extra.get("cache_reuse_mode", "chunk")) self._cache_reuse_strategy: CacheReuseStrategy - self._set_cache_reuse_strategy(str(extra.get("cache_reuse_mode", "chunk"))) + self._set_cache_reuse_strategy(cache_reuse_mode) self._runtime_config_ready = False self._rope_base: float = 10000.0 self._rope_rotary_dim: int = 0 @@ -123,55 +122,34 @@ def __init__( self._init_rope_config(vllm_config) if role == KVConnectorRole.SCHEDULER: - self._ipc_sync = IPCClientSync(self._socket_path) self._refresh_runtime_config() - self._pending_loads: dict[str, dict[str, Any]] = {} - self._pending_stores: dict[str, dict[str, Any]] = {} - self._pending_alloc: dict[str, PendingStore] = {} - self._pending_async_saves: set[str] = set() - self._req_tokens: dict[str, list[int]] = {} - else: - self._transfer_ready = False - self._transfer_mode = str(extra.get("transfer_mode", "iouring")) - self._ipc_load_async = IPCClientAsync(self._socket_path) - self._ipc_load_async_pool = [ - self._ipc_load_async, - *[ - IPCClientAsync(self._socket_path) - for _ in range(max(0, _LOAD_REQUEST_MAX_INFLIGHT - 1)) - ], - ] - self._ipc_store_async = IPCClientAsync(self._socket_path) - self._kv_caches: dict[str, torch.Tensor] = {} - self._layer_names: list[str] = [] - self._layer_idx_map: dict[str, int] = {} - self._meta: DaserConnectorMeta | None = None - self._save_futures: list = [] - self._pending_save_staging_bytes = 0 - self._store_staging_bytes = 0 - self._pending_store_staging_limit_bytes = 0 - self._store_staging_pool = None - self._pending_commits: set[str] = set() - self._pending_finished_saves: dict[str, Any] = {} - self._pending_loads: dict[str, Any] = {} - self._invalid_load_block_ids: set[int] = set() - self._load_request_queue = None - self._load_request_queue_lock = threading.Lock() - self._load_request_dispatcher_future = None - self._load_loop = asyncio.new_event_loop() - self._store_loop = asyncio.new_event_loop() - self._load_thread = threading.Thread( - target=self._run_load_loop, - daemon=True, - name="daser-load-io", + self._request_lifecycle = RequestLifecycle( + ipc_client=IPCClientSync(self._socket_path), + block_tokens=self._block_tokens, + slot_size=self._slot_size, + model_id=self._model_id, + cache_reuse_mode=self._cache_reuse_mode, + runtime_config_ready=self._runtime_config_ready, ) - self._store_thread = threading.Thread( - target=self._run_store_loop, - daemon=True, - name="daser-store-io", + else: + self._worker_runtime = WorkerRuntime( + socket_path=self._socket_path, + transfer_mode=str(extra.get("transfer_mode", "iouring")), + skip_l2=self._skip_l2, + tp_size=self._tp_size, + tp_rank=self._tp_rank, + server_tp_size=self._server_tp_size, + slot_size=self._slot_size, + store_path=self._store_path, + rank_stride_bytes=self._rank_stride_bytes, + rope_base=self._rope_base, + rope_rotary_dim=self._rope_rotary_dim, + rope_is_neox_style=self._rope_is_neox_style, + rope_delta_scale=self._rope_delta_scale, + load_key_scale=self._load_key_scale, + load_value_scale=self._load_value_scale, + kv_cache_config=kv_cache_config, ) - self._load_thread.start() - self._store_thread.start() logger.info("[CONNECTOR] role=%s socket=%s", role.name, self._socket_path) @@ -218,7 +196,9 @@ def _refresh_runtime_config(self) -> None: return finally: if owns_client: - client.close() + close = getattr(client, "close", None) + if close is not None: + close() self._store_path = str(config.get("store_path", self._store_path)) self._slot_size = int(config.get("slot_size", self._slot_size)) @@ -255,6 +235,7 @@ def _set_cache_reuse_strategy(self, cache_reuse_mode: str) -> None: Args: cache_reuse_mode: either ``"chunk"`` or ``"prefix"``. """ + self._cache_reuse_mode = cache_reuse_mode self._cache_reuse_strategy = build_cache_reuse_strategy( cache_reuse_mode, self._block_tokens, diff --git a/daser/connector/load_pipeline.py b/daser/connector/load_pipeline.py new file mode 100644 index 0000000..5c9ebb9 --- /dev/null +++ b/daser/connector/load_pipeline.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, replace +import threading +from typing import Any + +from daser.connector.ipc_client import IPCClientAsync +from daser.connector.metadata import ReqLoadSpec +from daser.connector.worker_memory import FixedCudaStagingPool + + +class LoadPipeline: + """Own load-side asyncio, IPC, staging, queue, and completion state. + + Args: + socket_path: DaseR server Unix socket path. + client_count: Number of independent async IPC connections. + + Async/thread-safety: + Construction and public methods run on the vLLM worker thread. Async + IPC runs exclusively on the private ``daser-load-io`` thread. + """ + + def __init__(self, socket_path: str, client_count: int) -> None: + self._clients = [IPCClientAsync(socket_path) for _ in range(client_count)] + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread( + target=self._run_loop, + daemon=True, + name="daser-load-io", + ) + self._queue: asyncio.Queue[Any] | None = None + self._queue_lock = threading.Lock() + self._dispatcher_future: Any | None = None + self._pending: dict[str, Any] = {} + self._invalid_block_ids: set[int] = set() + self._staging_pool: FixedCudaStagingPool | None = None + self._staging_registered = False + self._thread.start() + + @property + def pending(self) -> dict[str, Any]: + """Return worker-thread-owned request completion records.""" + return self._pending + + @property + def invalid_block_ids(self) -> set[int]: + """Return load-error block IDs accumulated by completion polling.""" + return self._invalid_block_ids + + @property + def staging_pool(self) -> FixedCudaStagingPool | None: + """Return the configured fixed staging pool, if registered.""" + return self._staging_pool + + @property + def staging_registered(self) -> bool: + """Return whether every fixed staging buffer has CUDA IPC registration.""" + return self._staging_registered + + @staging_registered.setter + def staging_registered(self, value: bool) -> None: + self._staging_registered = value + + def set_staging_pool(self, pool: FixedCudaStagingPool) -> None: + """Install the fixed load staging pool before request traffic.""" + self._staging_pool = pool + self._staging_registered = False + + def submit(self, coro: Any) -> Any: + """Submit a coroutine to the private load event loop.""" + return asyncio.run_coroutine_threadsafe(coro, self._loop) + + def client(self, buffer_index: int | None = None) -> IPCClientAsync: + """Return the IPC client assigned to a staging buffer index.""" + index = 0 if buffer_index is None else int(buffer_index) + return self._clients[index % len(self._clients)] + + def clients(self) -> tuple[IPCClientAsync, ...]: + """Return all load IPC clients for startup and shutdown orchestration.""" + return tuple(self._clients) + + def ensure_queue(self) -> asyncio.Queue[Any]: + """Create the load request queue on its owning event loop once.""" + with self._queue_lock: + if self._queue is None: + future = self.submit(self._create_queue()) + self._queue = future.result(timeout=5.0) + return self._queue + + def enqueue(self, item: Any) -> None: + """Append one request from the worker thread without blocking.""" + queue = self.ensure_queue() + self._loop.call_soon_threadsafe(queue.put_nowait, item) + + def ensure_dispatcher(self, coro: Any) -> None: + """Start the single queue dispatcher when none is active.""" + with self._queue_lock: + if ( + self._dispatcher_future is not None + and not self._dispatcher_future.done() + ): + coro.close() + return + self._dispatcher_future = self.submit(coro) + + def shutdown(self) -> None: + """Drain queue ownership, close IPC clients, and stop the load loop.""" + if self._queue is not None: + self._loop.call_soon_threadsafe(self._queue.put_nowait, None) + if self._dispatcher_future is not None: + self._dispatcher_future.result(timeout=120.0) + for client in self._clients: + self.submit(client.close()).result(timeout=5.0) + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=5.0) + + async def _create_queue(self) -> asyncio.Queue[Any]: + return asyncio.Queue() + + def _run_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + +@dataclass(frozen=True) +class LoadCopyRun: + """Describe one contiguous staging range with a shared KV transform.""" + + start: int + end: int + block_ids: list[int] + pos_offset: int + + +def build_load_read_plan( + reqs_to_load: dict[str, ReqLoadSpec], + slot_size: int, + include_req_ids: bool = False, +) -> tuple[int, list[dict[str, int]], list[Any]]: + """Build one combined server read and staging restore plan. + + Args: + reqs_to_load: Request IDs mapped to scheduler load specifications. + slot_size: Bytes stored for one rank-local KV slot. + include_req_ids: Include request IDs in restore ranges when true. + + Returns: + Total bytes, server read spans, and ranges mapping staging back to + request specifications. + + Async/thread-safety: + Pure CPU planning; safe to call from worker or load-loop threads. + """ + total_bytes = 0 + spans: list[dict[str, int]] = [] + per_req_ranges: list[Any] = [] + source_ranges: dict[tuple[str, int, int, int, int], tuple[int, int]] = {} + for req_id, spec in reqs_to_load.items(): + num_slots = len(spec.block_ids) + if num_slots == 0: + continue + nbytes = num_slots * slot_size + source_key = ( + spec.chunk_key, + spec.start_slot, + spec.num_slots, + spec.file_offset, + nbytes, + ) + existing = source_ranges.get(source_key) + if existing is None: + start = total_bytes + end = start + nbytes + spans.append( + { + "target_offset": start, + "nbytes": nbytes, + "file_offset": spec.file_offset, + } + ) + source_ranges[source_key] = (start, end) + total_bytes = end + else: + start, end = existing + 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 + + +def build_load_read_batches( + reqs_to_load: dict[str, ReqLoadSpec], + slot_size: int, + max_batch_bytes: int, + include_req_ids: bool = False, +) -> list[tuple[int, list[dict[str, int]], list[Any]]]: + """Split load work into staging-capacity-bounded read plans. + + Args: + reqs_to_load: Request IDs mapped to scheduler load specifications. + slot_size: Bytes stored for one rank-local KV slot. + max_batch_bytes: Maximum staging bytes in one transfer. + include_req_ids: Include request IDs in restore ranges when true. + + Returns: + Ordered read plans; requests larger than the cap are split on slot + boundaries. + + Async/thread-safety: + Pure CPU planning; safe to call from worker or load-loop threads. + """ + if slot_size <= 0: + raise ValueError("slot_size must be positive") + 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[Any]]] = [] + current: dict[str, ReqLoadSpec] = {} + current_slots = 0 + synthetic_id = 0 + + def flush() -> None: + nonlocal current, current_slots + if current: + batches.append( + build_load_read_plan( + current, + slot_size, + include_req_ids=include_req_ids, + ) + ) + current = {} + current_slots = 0 + + for req_id, spec in reqs_to_load.items(): + cursor = 0 + while cursor < len(spec.block_ids): + if current_slots >= max_slots: + flush() + available = max_slots - current_slots + take = min(available, len(spec.block_ids) - cursor) + if take <= 0: + flush() + continue + part = spec.block_ids[cursor : cursor + take] + batch_spec = replace( + spec, + start_slot=spec.start_slot + cursor, + num_slots=take, + block_ids=part, + file_offset=spec.file_offset + cursor * slot_size, + ) + key = ( + req_id + if cursor == 0 and take == len(spec.block_ids) + else f"{req_id}#{synthetic_id}" + ) + synthetic_id += 1 + current[key] = batch_spec + current_slots += take + cursor += take + flush() + return batches + + +def build_load_copy_runs( + per_req_ranges: list[tuple[int, int, ReqLoadSpec]], +) -> list[LoadCopyRun]: + """Merge adjacent restore ranges with the same position transform. + + Args: + per_req_ranges: Per-request staging ranges from a read plan. + + Returns: + Ordered contiguous copy runs. + + Async/thread-safety: + Pure CPU planning; safe to call from worker or load-loop threads. + """ + runs: list[LoadCopyRun] = [] + run_start = -1 + run_end = -1 + run_pos_offset = 0 + run_block_ids: list[int] = [] + + def flush() -> None: + nonlocal run_start, run_end, run_pos_offset, run_block_ids + if run_start >= 0 and run_block_ids: + runs.append( + LoadCopyRun( + start=run_start, + end=run_end, + block_ids=run_block_ids, + pos_offset=run_pos_offset, + ) + ) + run_start = -1 + run_end = -1 + run_pos_offset = 0 + run_block_ids = [] + + 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: + run_end = end + run_block_ids.extend(spec.block_ids) + continue + flush() + run_start = start + run_end = end + run_pos_offset = spec.pos_offset + run_block_ids = list(spec.block_ids) + flush() + return runs diff --git a/daser/connector/request_lifecycle.py b/daser/connector/request_lifecycle.py new file mode 100644 index 0000000..338d611 --- /dev/null +++ b/daser/connector/request_lifecycle.py @@ -0,0 +1,770 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import logging +import math +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from vllm.v1.core.kv_cache_utils import KVCacheBlocks + from vllm.v1.core.scheduler import SchedulerOutput + from vllm.v1.request import Request + +from daser.connector.helpers import PendingStore +from daser.connector.metadata import DaserConnectorMeta, ReqLoadSpec, ReqStoreSpec +from daser.connector.reuse import build_cache_reuse_strategy +from daser.connector.scheduler_planning import ( + _base_req_id, + _computed_tokens_after_step, + _contiguous_prefix_tokens, + _get_kv_transfer_flag, + _load_spec_from_chunk, + _matches_request_or_store_id, + _merge_adjacent_load_specs, + _store_slot_index, + _trim_chunk_to_external_window, +) +from daser.logging import init_logger + +logger = init_logger(__name__) + + +class RequestLifecycle: + """Own scheduler request state and synchronous IPC orchestration. + + Async/thread-safety: + These methods run on vLLM's scheduler thread and use the synchronous + IPC client owned by the connector instance. + """ + + def __init__( + self, + *, + ipc_client: Any, + block_tokens: int, + slot_size: int, + model_id: str, + cache_reuse_mode: str, + runtime_config_ready: bool, + ) -> None: + self._ipc_sync = ipc_client + self._block_tokens = block_tokens + self._slot_size = slot_size + self._model_id = model_id + self._cache_reuse_mode = cache_reuse_mode + self._runtime_config_ready = runtime_config_ready + self._cache_reuse_strategy = build_cache_reuse_strategy( + cache_reuse_mode, + block_tokens, + ) + self._pending_loads: dict[str, dict[str, Any]] = {} + self._pending_stores: dict[str, dict[str, Any]] = {} + self._pending_alloc: dict[str, PendingStore] = {} + self._pending_async_saves: set[str] = set() + self._req_tokens: dict[str, list[int]] = {} + + def get_num_new_matched_tokens( + self, + request: "Request", + num_computed_tokens: int, + ) -> "tuple[int | None, bool]": + """Query DaseR for cached KV matching request tokens. + + Args: + request: vLLM Request with prompt_token_ids. + num_computed_tokens: tokens already in vLLM's KV cache. + + Returns: + (num_external_tokens, is_async) - (0, False) on miss. + """ + tokens = list(request.prompt_token_ids) + self._req_tokens[request.request_id] = tokens + if not getattr(self, "_runtime_config_ready", True): + self._refresh_runtime_config() + + start = num_computed_tokens + available = len(tokens) - start + if available < self._block_tokens: + self._record_external_prefix_cache_miss(available) + return 0, False + + skip_load = bool(_get_kv_transfer_flag(request, "daser_skip_load")) + if skip_load: + logger.debug("[CONNECTOR] skip load req=%s", request.request_id[:8]) + self._record_external_prefix_cache_miss(available) + full_aligned = (len(tokens) // self._block_tokens) * self._block_tokens + skip_save = bool(_get_kv_transfer_flag(request, "daser_skip_save")) + pending_store = ( + None + if skip_save + else self._reuse_strategy().prepare_store(tokens, full_aligned) + ) + if pending_store is not None: + self._pending_alloc[request.request_id] = pending_store + return 0, False + + aligned = (available // self._block_tokens) * self._block_tokens + prefix = tokens[: start + aligned] + full_aligned = (len(tokens) // self._block_tokens) * self._block_tokens + skip_save = bool(_get_kv_transfer_flag(request, "daser_skip_save")) + + try: + chunks = self._lookup_with_external_prefix_metrics( + prefix, + self._model_id, + max(0, len(tokens) - num_computed_tokens), + num_computed_tokens, + ) + except Exception as exc: + logger.warning("[CONNECTOR] lookup failed: %s", exc) + self._runtime_config_ready = False + return 0, False + + if not chunks: + pending_store = ( + None + if skip_save + else self._reuse_strategy().prepare_store(tokens, full_aligned) + ) + if pending_store is not None: + self._pending_alloc[request.request_id] = pending_store + logger.debug("[CONNECTOR] cache miss req=%s", request.request_id[:8]) + return 0, False + + extra_tokens = _contiguous_prefix_tokens(chunks, num_computed_tokens) + if extra_tokens <= 0: + return 0, False + + pending_store = ( + None + if skip_save + else self._reuse_strategy().prepare_store( + tokens, + full_aligned, + chunks, + ) + ) + if pending_store is not None: + self._pending_alloc[request.request_id] = pending_store + + available = len(tokens) - num_computed_tokens + if extra_tokens >= available: + extra_tokens = available - 1 + if extra_tokens <= 0: + return 0, False + + if len(chunks) == 1: + self._pending_loads[request.request_id] = dict( + chunks[0], num_computed_tokens=num_computed_tokens + ) + else: + self._pending_loads[request.request_id] = { + str(i): dict(chunk, num_computed_tokens=num_computed_tokens) + for i, chunk in enumerate(chunks) + } + + logger.debug( + "[CONNECTOR] cache hit req=%s chunks=%d prefix_tokens=%d", + request.request_id[:8], + len(chunks), + extra_tokens, + ) + return extra_tokens, True + + def update_state_after_alloc( + self, + request: "Request", + blocks: "KVCacheBlocks", + num_external_tokens: int, + ) -> None: + """Record block IDs for requests that will load or store KV. + + Args: + request: vLLM Request. + blocks: vLLM KV cache block allocation for this request. + num_external_tokens: tokens from DaseR (0 if miss). + """ + req_id = request.request_id + block_ids: list[int] = [blk.block_id for blk in blocks.blocks[0]] + + if req_id in self._pending_loads: + chunks = self._pending_loads[req_id] + if "chunk_key" in chunks: + chunk = chunks + if not _trim_chunk_to_external_window( + chunk=chunk, + block_ids=block_ids, + external_start=int(chunk.get("num_computed_tokens", 0)), + num_external_tokens=num_external_tokens, + block_tokens=self._block_tokens, + slot_size=self._slot_size, + ): + del self._pending_loads[req_id] + self._record_pending_store_blocks(req_id, block_ids) + return + logger.debug( + "[CONNECTOR] load blocks req=%s blocks=%s", + req_id, + chunk["block_ids"], + ) + self._record_pending_store_blocks(req_id, block_ids) + return + for key, chunk in list(chunks.items()): + if not _trim_chunk_to_external_window( + chunk=chunk, + block_ids=block_ids, + external_start=int(chunk.get("num_computed_tokens", 0)), + num_external_tokens=num_external_tokens, + block_tokens=self._block_tokens, + slot_size=self._slot_size, + ): + logger.debug( + "[CONNECTOR] skip load req=%s key=%s target=%d slots=%d", + req_id[:8], + chunk.get("chunk_key", "")[:8], + int(chunk.get("target_token_start", 0)), + int(chunk["num_slots"]), + ) + del chunks[key] + continue + logger.debug( + "[CONNECTOR] load blocks req=%s key=%s blocks=%s", + req_id, + chunk.get("chunk_key", "")[:8], + chunk["block_ids"], + ) + self._record_pending_store_blocks(req_id, block_ids) + + def build_connector_meta( + self, scheduler_output: "SchedulerOutput" + ) -> DaserConnectorMeta: + """Package pending load/store specs into connector metadata. + + Args: + scheduler_output: vLLM SchedulerOutput for this step. + + Returns: + DaserConnectorMeta with reqs_to_load and reqs_to_store. + """ + meta = DaserConnectorMeta() + self._drop_preempted_pending_state(scheduler_output) + scheduled_ids: set[str] = set(scheduler_output.num_scheduled_tokens.keys()) + computed_after = _computed_tokens_after_step(scheduler_output) + self._record_cached_store_blocks(scheduler_output) + + for req_id, chunks in list(self._pending_loads.items()): + if "chunk_key" in chunks: + chunk = chunks + if "block_ids" in chunk: + meta.reqs_to_load[req_id] = _load_spec_from_chunk(chunk) + del self._pending_loads[req_id] + continue + ready = True + load_specs: list[ReqLoadSpec] = [] + for chunk in chunks.values(): + if "block_ids" not in chunk: + ready = False + continue + load_specs.append(_load_spec_from_chunk(chunk)) + merged_specs = _merge_adjacent_load_specs(load_specs, self._slot_size) + for idx, spec in enumerate(merged_specs): + load_id = req_id if len(merged_specs) == 1 else f"{req_id}:load:{idx}" + meta.reqs_to_load[load_id] = spec + if ready: + del self._pending_loads[req_id] + + for req_id, alloc in list(self._pending_stores.items()): + scheduled_tokens = scheduler_output.num_scheduled_tokens.get(req_id, 0) + if scheduled_tokens <= 0 and ":store:" in req_id: + scheduled_tokens = scheduler_output.num_scheduled_tokens.get( + _base_req_id(req_id), + 0, + ) + base_req_id = _base_req_id(req_id) + computed_tokens = computed_after.get(base_req_id, scheduled_tokens) + slot_index = _store_slot_index(req_id) + required_tokens = ( + (slot_index + 1) * self._block_tokens + if slot_index is not None + else int(alloc["token_count"]) + ) + should_store = ( + base_req_id in scheduled_ids + and scheduled_tokens > 0 + and computed_tokens >= required_tokens + and "block_ids" in alloc + ) + if should_store: + meta.reqs_to_store[req_id] = ReqStoreSpec( + chunk_key=alloc["chunk_key"], + start_slot=alloc["start_slot"], + num_slots=alloc["num_slots"], + block_ids=alloc["block_ids"], + file_offset=alloc["file_offset"], + token_count=alloc["token_count"], + ) + del self._pending_stores[req_id] + + if meta.reqs_to_store: + meta.reqs_to_store = self._filter_live_store_specs(meta.reqs_to_store) + pending_async_saves = self._pending_async_save_ids() + for req_id in meta.reqs_to_store: + pending_async_saves.add(_base_req_id(req_id)) + + if logger.isEnabledFor(logging.DEBUG): + for req_id, spec in meta.reqs_to_load.items(): + logger.debug( + "[CONNECTOR] meta LOAD req=%s start_slot=%d blocks=%d tokens=%d", + req_id[:8], + spec.start_slot, + len(spec.block_ids), + spec.token_count, + ) + for req_id, spec in meta.reqs_to_store.items(): + logger.debug( + "[CONNECTOR] meta STORE req=%s start_slot=%d blocks=%d tokens=%d", + req_id[:8], + spec.start_slot, + len(spec.block_ids), + spec.token_count, + ) + return meta + + def _drop_preempted_pending_state( + self, + scheduler_output: "SchedulerOutput", + ) -> None: + """Discard pending scheduler state whose KV blocks were preempted. + + Args: + scheduler_output: vLLM SchedulerOutput for this step. + + Async/thread-safety: + Runs on the scheduler thread before metadata is handed to workers. + """ + preempted_req_ids = getattr(scheduler_output, "preempted_req_ids", set()) + pending_async_saves = self._pending_async_save_ids() + for req_id in preempted_req_ids: + base_req_id = str(req_id) + pending_async_saves.discard(base_req_id) + for pending_req_id in list(self._pending_loads): + if _matches_request_or_store_id(pending_req_id, base_req_id): + self._pending_loads.pop(pending_req_id, None) + for pending_req_id in list(self._pending_stores): + if _matches_request_or_store_id(pending_req_id, base_req_id): + self._drop_pending_store(pending_req_id) + for pending_req_id in list(self._pending_alloc): + if _matches_request_or_store_id(pending_req_id, base_req_id): + self._pending_alloc.pop(pending_req_id, None) + + def _filter_live_store_specs( + self, + specs: dict[str, ReqStoreSpec], + ) -> dict[str, ReqStoreSpec]: + """Drop store specs whose server allocation was already evicted. + + Args: + specs: Store specs built for the current scheduler step. + + Returns: + Specs that still own their allocated server slot ranges. + """ + try: + live_keys = self._ipc_sync.live_allocations( + [ + { + "chunk_key": spec.chunk_key, + "start_slot": spec.start_slot, + "num_slots": spec.num_slots, + } + for spec in specs.values() + ] + ) + except Exception as exc: + logger.warning("[CONNECTOR] live_allocations failed: %s", exc) + return specs + return { + req_id: spec + for req_id, spec in specs.items() + if spec.chunk_key in live_keys + } + + def _record_cached_store_blocks(self, scheduler_output: "SchedulerOutput") -> None: + """Append blocks from later chunked-prefill steps to store trackers. + + Args: + scheduler_output: vLLM SchedulerOutput for this step. + """ + cached_reqs = getattr(scheduler_output, "scheduled_cached_reqs", None) + if cached_reqs is None: + return + + req_ids = getattr(cached_reqs, "req_ids", []) + new_block_ids = getattr(cached_reqs, "new_block_ids", []) + for req_id, block_groups in zip(req_ids, new_block_ids, strict=False): + pending_store = self._pending_alloc.get(req_id) + if pending_store is None or block_groups is None: + continue + if not block_groups: + continue + block_group = block_groups[0] + if block_group is None: + continue + block_ids = list(block_group) + if ( + req_id in getattr(cached_reqs, "resumed_req_ids", set()) + and block_ids[: len(pending_store.block_ids)] == pending_store.block_ids + ): + pending_store.block_ids = block_ids + else: + pending_store.block_ids.extend(block_ids) + needed_slots = math.ceil(pending_store.token_count / self._block_tokens) + if len(pending_store.block_ids) > needed_slots: + pending_store.block_ids = pending_store.block_ids[:needed_slots] + self._maybe_allocate_pending_store(req_id, pending_store) + + def _lookup_with_external_prefix_metrics( + self, + tokens: list[int], + model_id: str, + queries: int, + num_computed_tokens: int, + ) -> list[dict[str, Any]]: + """Run lookup while passing vLLM external-prefix query count when supported. + + Args: + tokens: token prefix sent to DaseR lookup. + model_id: model identifier. + queries: vLLM external prefix query token count. + num_computed_tokens: tokens already computed locally by vLLM. + + Returns: + Matching chunk dicts returned by the IPC client. + + Thread-safety: + Runs on the scheduler thread and uses the synchronous IPC client. + """ + try: + return self._ipc_sync.lookup( + tokens, + model_id, + external_prefix_queries=queries, + num_computed_tokens=num_computed_tokens, + ) + except TypeError as exc: + if "external_prefix_queries" not in str(exc): + raise + return self._ipc_sync.lookup(tokens, model_id) + + def _record_external_prefix_cache_miss(self, queries: int) -> None: + """Record a connector external-prefix miss when lookup is skipped. + + Args: + queries: vLLM external prefix query token count. + + Returns: + None. + + Thread-safety: + Runs on the scheduler thread and uses the synchronous IPC client + when it supports the diagnostic operation. + """ + recorder = getattr(self._ipc_sync, "record_external_prefix_cache", None) + if recorder is None: + return + recorder(max(0, int(queries)), 0) + + def _record_pending_store_blocks(self, req_id: str, block_ids: list[int]) -> None: + """Record request KV blocks for a pending scheduler store. + + Args: + req_id: vLLM request ID. + block_ids: KV block IDs allocated for the request. + """ + pending_store = self._pending_alloc.get(req_id) + if pending_store is None: + return + requested_tokens = pending_store.token_count + pending_store.block_ids = block_ids[ + : math.ceil(requested_tokens / self._block_tokens) + ] + self._maybe_allocate_pending_store(req_id, pending_store) + + def _refresh_runtime_config(self) -> None: + """Refresh scheduler geometry and reuse policy over owned sync IPC.""" + try: + config = self._ipc_sync.get_runtime_config() + except Exception as exc: # noqa: BLE001 + logger.info("[CONNECTOR] runtime config unavailable: %s", exc) + return + self._slot_size = int(config.get("slot_size", self._slot_size)) + block_tokens = int(config.get("block_tokens", self._block_tokens)) + self._model_id = str(config.get("model_id", self._model_id)) + cache_reuse_mode = str(config.get("cache_reuse_mode", self._cache_reuse_mode)) + if ( + cache_reuse_mode != self._cache_reuse_mode + or block_tokens != self._block_tokens + ): + self._cache_reuse_mode = cache_reuse_mode + self._block_tokens = block_tokens + self._cache_reuse_strategy = build_cache_reuse_strategy( + cache_reuse_mode, + self._block_tokens, + ) + else: + self._block_tokens = block_tokens + self._runtime_config_ready = bool(self._slot_size) + + def _init_reuse_strategy(self) -> None: + """Initialize the scheduler cache reuse strategy from current config.""" + self._cache_reuse_strategy = build_cache_reuse_strategy( + "chunk", + self._block_tokens, + ) + + def _reuse_strategy(self) -> Any: + """Return the configured cache reuse strategy. + + Returns: + Cache reuse strategy initialized from connector runtime config. + """ + strategy = getattr(self, "_cache_reuse_strategy", None) + if strategy is None: + self._init_reuse_strategy() + strategy = self._cache_reuse_strategy + return strategy + + def allocate_store_chunk( + self, + chunk_key: str, + token_count: int, + ) -> dict[str, Any]: + """Allocate server metadata for a pending scheduler store. + + Args: + chunk_key: cache key to allocate. + token_count: number of tokens covered by the allocation. + + Returns: + Mutable server allocation metadata. + """ + return self._ipc_sync.alloc_chunk(chunk_key, token_count, self._model_id) + + def allocate_store_chunks( + self, + chunks: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + """Allocate server metadata for multiple pending scheduler stores. + + Args: + chunks: chunk descriptors with chunk_key and token_count. + + Returns: + Mutable server allocation metadata for each chunk. + """ + alloc_chunks = getattr(self._ipc_sync, "alloc_chunks", None) + if alloc_chunks is None: + return [ + { + **self.allocate_store_chunk( + str(chunk["chunk_key"]), + int(chunk["token_count"]), + ), + "chunk_key": str(chunk["chunk_key"]), + } + for chunk in chunks + ] + return alloc_chunks(chunks, self._model_id) + + def set_pending_store(self, req_id: str, alloc: dict[str, Any]) -> None: + """Record an allocated store for later connector metadata packaging. + + Args: + req_id: vLLM request ID or synthetic store work ID. + alloc: mutable allocation metadata. + """ + self._pending_stores[req_id] = alloc + + def has_pending_store(self, req_id: str) -> bool: + """Return whether a pending store entry already exists. + + Args: + req_id: vLLM request ID or synthetic store work ID. + + Returns: + True when the store is already pending. + """ + return req_id in self._pending_stores + + def drop_pending_alloc(self, req_id: str) -> None: + """Remove pending allocation state for a request. + + Args: + req_id: vLLM request ID. + """ + self._pending_alloc.pop(req_id, None) + + def _drop_pending_store(self, req_id: str) -> None: + """Remove a pending store and release its server writer claim. + + Args: + req_id: vLLM request ID or synthetic store work ID. + """ + alloc = self._pending_stores.pop(req_id, None) + if alloc is None: + return + release = getattr(self._ipc_sync, "release_chunk_writer", None) + if release is None: + return + try: + release( + str(alloc["chunk_key"]), + int(alloc["start_slot"]), + int(alloc["num_slots"]), + ) + except Exception as exc: # noqa: BLE001 + logger.warning("[CONNECTOR] release_chunk_writer failed: %s", exc) + + def _discard_pending_request(self, req_id: str) -> None: + """Clear scheduler-side pending state for a request. + + Args: + req_id: vLLM request ID. + """ + self._pending_loads.pop(req_id, None) + if req_id in self._pending_stores: + self._drop_pending_store(req_id) + for pending_req_id in list(self._pending_stores): + if pending_req_id.startswith(f"{req_id}:store:"): + self._drop_pending_store(pending_req_id) + self._pending_alloc.pop(req_id, None) + + def _pending_async_save_ids(self) -> set[str]: + """Return request IDs whose worker-side saves are still pending. + + Returns: + Mutable set of base vLLM request IDs. + + Thread-safety: + Runs on the scheduler thread. The lazy initialization supports + tests and mixin probes that do not call ``DaserConnector.__init__``. + """ + pending = getattr(self, "_pending_async_saves", None) + if pending is None: + pending = set() + self._pending_async_saves = pending + return pending + + def _maybe_allocate_pending_store( + self, req_id: str, pending_store: PendingStore + ) -> None: + """Allocate a DaseR chunk once a pending store has full KV coverage. + + Args: + req_id: vLLM request ID being tracked. + pending_store: store tracker for the request. + """ + requested_tokens = pending_store.token_count + strategy = self._reuse_strategy() + if not strategy.ready_to_allocate(pending_store): + return + tokens = self._req_tokens.get(req_id, []) + if len(tokens) < requested_tokens: + return + plan = strategy.plan_store( + req_id, + pending_store, + tokens, + set(self._pending_stores), + ) + if plan.invalid: + self._pending_alloc.pop(req_id, None) + return + if plan.intents: + try: + if len(plan.intents) == 1 and plan.intents[0].req_id == req_id: + intent = plan.intents[0] + allocations = [ + self.allocate_store_chunk( + intent.chunk_key, + intent.token_count, + ) + ] + else: + allocations = self.allocate_store_chunks( + [ + { + "chunk_key": intent.chunk_key, + "token_count": intent.token_count, + } + for intent in plan.intents + ] + ) + except Exception as exc: # noqa: BLE001 + logger.warning("[CONNECTOR] store allocation failed: %s", exc) + return + if len(allocations) != len(plan.intents): + logger.warning( + "[CONNECTOR] allocation returned %d entries for %d intents", + len(allocations), + len(plan.intents), + ) + return + for intent, alloc in zip(plan.intents, allocations, strict=True): + if bool(alloc.get("skipped", False)): + continue + alloc["chunk_key"] = str(alloc.get("chunk_key", intent.chunk_key)) + alloc["token_count"] = intent.token_count + alloc["num_slots"] = len(intent.block_ids) + alloc["block_ids"] = intent.block_ids + self._pending_stores[intent.req_id] = alloc + pending_store.rolling_key = plan.next_key + pending_store.rolling_slot_index = plan.next_slot + if plan.complete: + pending_store.chunk_key = plan.next_key + self._pending_alloc.pop(req_id, None) + + def request_finished( + self, + request: "Request", + block_ids: list[int], + ) -> "tuple[bool, dict[str, Any] | None]": + """Clean up per-request state after inference completes. + + Args: + request: finished vLLM Request. + block_ids: block IDs being freed. + + Returns: + (True, None) when DaseR is still storing this request's KV blocks, + otherwise (False, None). + """ + del block_ids + if request.request_id in self._pending_async_save_ids(): + return True, None + self._req_tokens.pop(request.request_id, None) + self._discard_pending_request(request.request_id) + return False, None + + def update_connector_output(self, connector_output: Any) -> None: + """Update scheduler state from worker-side transfer completions. + + Args: + connector_output: vLLM KVConnectorOutput carrying finished request + IDs from workers. + + Async/thread-safety: + Runs on vLLM's scheduler thread after worker connector polling. + """ + pending_async_saves = self._pending_async_save_ids() + for req_id in getattr(connector_output, "finished_sending", None) or (): + pending_async_saves.discard(req_id) + self._req_tokens.pop(req_id, None) + self._discard_pending_request(req_id) + for req_id in getattr(connector_output, "finished_recving", None) or (): + if req_id not in self._pending_alloc: + self._req_tokens.pop(req_id, None) + for pending_req_id in list(self._pending_loads): + if _matches_request_or_store_id(pending_req_id, req_id): + self._pending_loads.pop(pending_req_id, None) diff --git a/daser/connector/reuse.py b/daser/connector/reuse.py index 2920199..1fd03be 100644 --- a/daser/connector/reuse.py +++ b/daser/connector/reuse.py @@ -3,6 +3,7 @@ # Standard from abc import ABC, abstractmethod +from dataclasses import dataclass import math from typing import Any @@ -19,6 +20,27 @@ logger = init_logger(__name__) +@dataclass(frozen=True) +class StoreIntent: + """Describe one server allocation needed for pending store work.""" + + req_id: str + chunk_key: str + token_count: int + block_ids: list[int] + + +@dataclass(frozen=True) +class StoreIntentPlan: + """Return store intents together with the next strategy cursor state.""" + + intents: tuple[StoreIntent, ...] + next_key: str + next_slot: int + complete: bool + invalid: bool = False + + class CacheReuseStrategy(ABC): """Compute store keys and allocate scheduler-side store work. @@ -65,21 +87,23 @@ def ready_to_allocate(self, pending_store: PendingStore) -> bool: return len(pending_store.block_ids) >= num_slots @abstractmethod - def allocate_store( + def plan_store( self, - owner: Any, req_id: str, pending_store: PendingStore, tokens: list[int], - ) -> None: - """Allocate server-side store metadata once block IDs are known. + pending_store_ids: set[str], + ) -> StoreIntentPlan: + """Build store allocation intents once block IDs are known. Args: - owner: scheduler connector object with ``_ipc_sync``, - ``_model_id``, ``_slot_size``, and ``_pending_*`` attributes. req_id: vLLM request ID. pending_store: pending store state for this request. tokens: full prompt token IDs. + pending_store_ids: synthetic or base IDs already allocated. + + Returns: + Immutable allocation intent plan for request lifecycle execution. """ @@ -122,56 +146,43 @@ def prepare_store( return None return PendingStore(chunk_key=chunk_key, token_count=aligned_tokens) - def allocate_store( + def plan_store( self, - owner: Any, req_id: str, pending_store: PendingStore, tokens: list[int], - ) -> None: - """Allocate one store covering the whole aligned prefix. + pending_store_ids: set[str], + ) -> StoreIntentPlan: + """Plan one store covering the whole aligned prefix. Args: - owner: scheduler connector object. req_id: vLLM request ID. pending_store: pending store state for this request. tokens: full prompt token IDs. + pending_store_ids: existing allocated work IDs. + + Returns: + One whole-prefix intent, or an invalid plan on key mismatch. """ + del pending_store_ids requested_tokens = pending_store.token_count num_slots = math.ceil(requested_tokens / self._block_tokens) chunk_key = pending_store.chunk_key if chunk_key != self.store_key(tokens, requested_tokens): logger.warning("[CONNECTOR] pending store key mismatch req=%s", req_id[:8]) - owner.drop_pending_alloc(req_id) - return - try: - alloc = owner.allocate_store_chunk( - chunk_key, - requested_tokens, - ) - except Exception as exc: # noqa: BLE001 - logger.warning("[CONNECTOR] alloc_chunk failed: %s", exc) - return - if bool(alloc.get("skipped", False)): - owner.drop_pending_alloc(req_id) - logger.debug( - "[CONNECTOR] skip duplicate store req=%s key=%s", - req_id[:8], - chunk_key[:8], - ) - return - alloc["chunk_key"] = chunk_key - alloc["token_count"] = requested_tokens - alloc["num_slots"] = num_slots - alloc["block_ids"] = pending_store.block_ids[:num_slots] - owner.set_pending_store(req_id, alloc) - owner.drop_pending_alloc(req_id) - logger.debug( - "[CONNECTOR] alloc store req=%s key=%s tokens=%d/%d", - req_id, - alloc["chunk_key"][:8], - requested_tokens, - requested_tokens, + return StoreIntentPlan((), chunk_key, num_slots, True, invalid=True) + return StoreIntentPlan( + intents=( + StoreIntent( + req_id=req_id, + chunk_key=chunk_key, + token_count=requested_tokens, + block_ids=pending_store.block_ids[:num_slots], + ), + ), + next_key=chunk_key, + next_slot=num_slots, + complete=True, ) @@ -227,20 +238,23 @@ def ready_to_allocate(self, pending_store: PendingStore) -> bool: """ return len(pending_store.block_ids) > pending_store.rolling_slot_index - def allocate_store( + def plan_store( self, - owner: Any, req_id: str, pending_store: PendingStore, tokens: list[int], - ) -> None: - """Allocate one store target for each missing rolling-prefix slot. + pending_store_ids: set[str], + ) -> StoreIntentPlan: + """Plan one store target for each missing rolling-prefix slot. Args: - owner: scheduler connector object. req_id: vLLM request ID. pending_store: pending store state for this request. tokens: full prompt token IDs. + pending_store_ids: existing allocated work IDs. + + Returns: + Missing slot intents and the next rolling-prefix cursor state. """ requested_tokens = pending_store.token_count num_slots = math.ceil(requested_tokens / self._block_tokens) @@ -260,60 +274,30 @@ def allocate_store( key = next_key if slot_i >= pending_store.start_slot_index: store_id = f"{req_id}:store:{slot_i}" - if not owner.has_pending_store(store_id): + if store_id not in pending_store_ids: run.append((slot_i, key)) slot_i += 1 - - if run: - try: - allocations = owner.allocate_store_chunks( - [ - {"chunk_key": chunk_key, "token_count": self._block_tokens} - for _slot_i, chunk_key in run - ] - ) - except Exception as exc: # noqa: BLE001 - logger.warning("[CONNECTOR] alloc_chunks failed: %s", exc) - return - if len(allocations) != len(run): - logger.warning( - "[CONNECTOR] alloc_chunks returned %d allocations for %d slots", - len(allocations), - len(run), - ) - return - for (store_slot_i, chunk_key), alloc in zip( - run, - allocations, - strict=True, - ): - if bool(alloc.get("skipped", False)): - continue - alloc["chunk_key"] = str(alloc.get("chunk_key", chunk_key)) - alloc["token_count"] = self._block_tokens - alloc["num_slots"] = 1 - alloc["block_ids"] = [pending_store.block_ids[store_slot_i]] - owner.set_pending_store(f"{req_id}:store:{store_slot_i}", alloc) - - pending_store.rolling_key = key - pending_store.rolling_slot_index = slot_i + intents = tuple( + StoreIntent( + req_id=f"{req_id}:store:{store_slot_i}", + chunk_key=chunk_key, + token_count=self._block_tokens, + block_ids=[pending_store.block_ids[store_slot_i]], + ) + for store_slot_i, chunk_key in run + ) if slot_i >= num_slots: if pending_store.chunk_key and pending_store.chunk_key != key: logger.warning( "[CONNECTOR] pending store key mismatch req=%s", req_id[:8] ) - owner.drop_pending_alloc(req_id) - return - pending_store.chunk_key = key - - if slot_i >= num_slots: - owner.drop_pending_alloc(req_id) - if run: - logger.debug( - "[CONNECTOR] alloc rolling-prefix stores req=%s slots=%d", - req_id[:8], - len(run), - ) + return StoreIntentPlan((), key, slot_i, True, invalid=True) + return StoreIntentPlan( + intents=intents, + next_key=key, + next_slot=slot_i, + complete=slot_i >= num_slots, + ) def build_cache_reuse_strategy( diff --git a/daser/connector/scheduler.py b/daser/connector/scheduler.py index 6064340..074d79b 100644 --- a/daser/connector/scheduler.py +++ b/daser/connector/scheduler.py @@ -1,419 +1,52 @@ # SPDX-License-Identifier: Apache-2.0 -# Standard -import logging -import math +from __future__ import annotations + from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - # Third Party from vllm.v1.core.kv_cache_utils import KVCacheBlocks from vllm.v1.core.scheduler import SchedulerOutput from vllm.v1.request import Request -# First Party -from daser.connector.helpers import PendingStore, base_req_id -from daser.connector.metadata import DaserConnectorMeta, ReqLoadSpec, ReqStoreSpec -from daser.connector.reuse import build_cache_reuse_strategy -from daser.logging import init_logger - -logger = init_logger(__name__) - - -def _base_req_id(req_id: str) -> str: - """Compatibility wrapper for tests importing the scheduler-private helper.""" - return base_req_id(req_id) - - -def _store_slot_index(req_id: str) -> int | None: - """Return the rolling-prefix store slot encoded in a synthetic request ID. - - Args: - req_id: vLLM request ID or ``:store:`` synthetic ID. - - Returns: - Slot index for synthetic store IDs, or None for regular IDs. - """ - if ":store:" not in req_id: - return None - try: - return int(req_id.rsplit(":store:", 1)[1]) - except ValueError: - return None - - -def _matches_request_or_store_id(req_id: str, base_req_id: str) -> bool: - """Return whether ``req_id`` belongs to a base request. - - Args: - req_id: vLLM request ID or synthetic connector work ID. - base_req_id: Base vLLM request ID to match. - - Returns: - True when ``req_id`` is the base request or one of its synthetic - store/load entries. - """ - return ( - req_id == base_req_id - or req_id.startswith(f"{base_req_id}:store:") - or req_id.startswith(f"{base_req_id}:load:") - ) - - -def _computed_tokens_after_step( - scheduler_output: "SchedulerOutput", -) -> dict[str, int]: - """Return per-request token counts that are valid after this step. - - Args: - scheduler_output: vLLM SchedulerOutput for this step. - - Returns: - Mapping from request ID to ``num_computed_tokens + scheduled_tokens``. - Falls back to the scheduled token count when older or test scheduler - outputs do not expose prior computed-token metadata. - """ - scheduled = dict(getattr(scheduler_output, "num_scheduled_tokens", {})) - computed_after = {req_id: int(tokens) for req_id, tokens in scheduled.items()} - - for req_data in getattr(scheduler_output, "scheduled_new_reqs", []) or []: - req_id = str(getattr(req_data, "req_id", "")) - if req_id in scheduled: - computed_after[req_id] = int( - getattr(req_data, "num_computed_tokens", 0) - ) + int(scheduled[req_id]) - - cached_reqs = getattr(scheduler_output, "scheduled_cached_reqs", None) - if cached_reqs is not None: - req_ids = getattr(cached_reqs, "req_ids", []) - prior_counts = getattr(cached_reqs, "num_computed_tokens", []) - for req_id, prior in zip(req_ids, prior_counts, strict=False): - req_id = str(req_id) - if req_id in scheduled: - computed_after[req_id] = int(prior) + int(scheduled[req_id]) - - return computed_after - - -def _get_kv_transfer_flag(request: "Request", key: str) -> Any: - """Return ``request.kv_transfer_params[key]`` if present, else ``None``. - - Args: - request: vLLM ``Request`` or compatible object. - key: connector-specific flag name to extract. - - Returns: - The value under ``key``, or ``None`` when absent. - """ - params = getattr(request, "kv_transfer_params", None) - if not isinstance(params, dict): - return None - return params.get(key) - - -def _block_ids_for_chunk( - block_ids: list[int], - target_token_start: int, - num_slots: int, - block_tokens: int, - max_tokens: int | None = None, -) -> list[int]: - """Return vLLM block IDs for a chunk's target prompt range. - - Args: - block_ids: all block IDs allocated to the request. - target_token_start: token offset where the chunk starts in the prompt. - num_slots: number of blocks/slots covered by the chunk. - block_tokens: tokens per vLLM block. - max_tokens: optional upper bound on accepted external tokens. - - Returns: - Slice of block_ids for the chunk, or an empty list when the range - is not block-aligned or exceeds the allocated blocks. - """ - if target_token_start % block_tokens != 0: - return [] - target_block_start = target_token_start // block_tokens - effective_slots = num_slots - if max_tokens is not None: - remaining_tokens = max_tokens - target_token_start - if remaining_tokens <= 0: - return [] - effective_slots = min(num_slots, math.ceil(remaining_tokens / block_tokens)) - target_block_end = target_block_start + effective_slots - if target_block_start < 0 or target_block_end > len(block_ids): - return [] - return block_ids[target_block_start:target_block_end] - - -def _trim_chunk_to_external_window( - chunk: dict[str, Any], - block_ids: list[int], - external_start: int, - num_external_tokens: int, - block_tokens: int, - slot_size: int, -) -> bool: - """Trim chunk metadata to the external token interval vLLM requested. - - Args: - chunk: Mutable chunk metadata returned by the server. - block_ids: Full vLLM block allocation for the request. - external_start: Token offset where external KV loading begins. - num_external_tokens: Number of tokens accepted from the connector. - block_tokens: Tokens per vLLM block. - slot_size: Bytes per DaseR slot. - - Returns: - True when the chunk still covers at least one whole KV block. - """ - if external_start % block_tokens != 0 or num_external_tokens <= 0: - return False - target_start = int(chunk.get("target_token_start", 0)) - target_end = target_start + int(chunk["token_count"]) - external_end = external_start + num_external_tokens - load_start = max(target_start, external_start) - load_end = min(target_end, external_end) - load_start = ((load_start + block_tokens - 1) // block_tokens) * block_tokens - load_end = ((load_end + block_tokens - 1) // block_tokens) * block_tokens - load_end = min(load_end, target_end) - if load_end <= load_start: - return False - - skip_slots = (load_start - target_start) // block_tokens - num_slots = (load_end - load_start) // block_tokens - if load_start < external_start: - return False - block_start = load_start // block_tokens - block_end = block_start + num_slots - if block_start < 0 or block_end > len(block_ids): - return False - - chunk["start_slot"] = int(chunk["start_slot"]) + skip_slots - chunk["file_offset"] = int(chunk["file_offset"]) + skip_slots * slot_size - chunk["num_slots"] = num_slots - chunk["token_count"] = num_slots * block_tokens - chunk["target_token_start"] = load_start - chunk["block_ids"] = block_ids[block_start:block_end] - return bool(chunk["block_ids"]) - - -def _contiguous_prefix_tokens( - chunks: list[dict[str, Any]], num_computed_tokens: int -) -> int: - """Return external tokens covered contiguously after computed tokens. - - Args: - chunks: server chunk payloads with target_token_start and token_count. - num_computed_tokens: tokens vLLM already has locally. - - Returns: - Number of additional contiguous prefix tokens covered by chunks. - """ - covered_until = num_computed_tokens - for chunk in sorted( - chunks, - key=lambda item: int(item.get("target_token_start", 0)), - ): - target_start = int(chunk.get("target_token_start", 0)) - token_count = int(chunk["token_count"]) - target_end = target_start + token_count - if target_end <= covered_until: - continue - if target_start > covered_until: - break - covered_until = target_end - return covered_until - num_computed_tokens - - -def _load_spec_from_chunk(chunk: dict[str, Any]) -> ReqLoadSpec: - """Build a worker load specification from scheduler chunk metadata. - - Args: - chunk: Chunk metadata returned by the server and annotated with vLLM - block IDs during allocation. - - Returns: - ReqLoadSpec consumed by the worker load path. - - Async/thread-safety: - Pure scheduler-thread helper; it does not mutate connector state. - """ - return ReqLoadSpec( - chunk_key=str(chunk["chunk_key"]), - start_slot=int(chunk["start_slot"]), - num_slots=int(chunk["num_slots"]), - block_ids=list(chunk["block_ids"]), - file_offset=int(chunk["file_offset"]), - token_count=int(chunk["token_count"]), - target_token_start=int(chunk.get("target_token_start", 0)), - pos_offset=int(chunk.get("pos_offset", 0)), - ) - - -def _merge_adjacent_load_specs( - specs: list[ReqLoadSpec], - slot_size: int, -) -> list[ReqLoadSpec]: - """Merge adjacent load specs that describe one continuous KV byte range. - - Args: - specs: Load specs for one request in prompt order. - slot_size: Bytes represented by one DaseR KV slot. - - Returns: - Coalesced load specs. Chunk keys from the first spec in a run are kept - only as diagnostics; the worker load path addresses data by byte range. - - Async/thread-safety: - Pure scheduler-thread helper; it does not mutate connector state. - """ - merged: list[ReqLoadSpec] = [] - for spec in specs: - if not spec.block_ids: - continue - if not merged: - merged.append(spec) - continue - prev = merged[-1] - prev_slots = len(prev.block_ids) - adjacent = ( - prev.pos_offset == spec.pos_offset - and prev.start_slot + prev_slots == spec.start_slot - and prev.file_offset + prev_slots * slot_size == spec.file_offset - and prev.target_token_start + prev.token_count == spec.target_token_start - ) - if not adjacent: - merged.append(spec) - continue - merged[-1] = ReqLoadSpec( - chunk_key=prev.chunk_key, - start_slot=prev.start_slot, - num_slots=prev_slots + len(spec.block_ids), - block_ids=[*prev.block_ids, *spec.block_ids], - file_offset=prev.file_offset, - token_count=prev.token_count + spec.token_count, - target_token_start=prev.target_token_start, - pos_offset=prev.pos_offset, - ) - return merged +from daser.connector.metadata import DaserConnectorMeta +from daser.connector.scheduler_planning import ( + _base_req_id, + _block_ids_for_chunk, + _computed_tokens_after_step, + _contiguous_prefix_tokens, + _get_kv_transfer_flag, + _matches_request_or_store_id, + _store_slot_index, + _trim_chunk_to_external_window, +) + +__all__ = [ + "SchedulerConnectorMixin", + "_base_req_id", + "_block_ids_for_chunk", + "_computed_tokens_after_step", + "_contiguous_prefix_tokens", + "_get_kv_transfer_flag", + "_matches_request_or_store_id", + "_store_slot_index", + "_trim_chunk_to_external_window", +] class SchedulerConnectorMixin: - """Scheduler-role vLLM connector behavior. - - Async/thread-safety: - These methods run on vLLM's scheduler thread and use the synchronous - IPC client owned by the connector instance. - """ + """Adapt vLLM scheduler hooks to the request lifecycle interface.""" def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int, - ) -> "tuple[int | None, bool]": - """Query DaseR for cached KV matching request tokens. - - Args: - request: vLLM Request with prompt_token_ids. - num_computed_tokens: tokens already in vLLM's KV cache. - - Returns: - (num_external_tokens, is_async) - (0, False) on miss. - """ - tokens = list(request.prompt_token_ids) - self._req_tokens[request.request_id] = tokens - if not getattr(self, "_runtime_config_ready", True): - self._refresh_runtime_config() - - start = num_computed_tokens - available = len(tokens) - start - if available < self._block_tokens: - self._record_external_prefix_cache_miss(available) - return 0, False - - skip_load = bool(_get_kv_transfer_flag(request, "daser_skip_load")) - if skip_load: - logger.debug("[CONNECTOR] skip load req=%s", request.request_id[:8]) - self._record_external_prefix_cache_miss(available) - full_aligned = (len(tokens) // self._block_tokens) * self._block_tokens - skip_save = bool(_get_kv_transfer_flag(request, "daser_skip_save")) - pending_store = ( - None - if skip_save - else self._reuse_strategy().prepare_store(tokens, full_aligned) - ) - if pending_store is not None: - self._pending_alloc[request.request_id] = pending_store - return 0, False - - aligned = (available // self._block_tokens) * self._block_tokens - prefix = tokens[: start + aligned] - full_aligned = (len(tokens) // self._block_tokens) * self._block_tokens - skip_save = bool(_get_kv_transfer_flag(request, "daser_skip_save")) - - try: - chunks = self._lookup_with_external_prefix_metrics( - prefix, - self._model_id, - max(0, len(tokens) - num_computed_tokens), - num_computed_tokens, - ) - except Exception as exc: - logger.warning("[CONNECTOR] lookup failed: %s", exc) - self._runtime_config_ready = False - return 0, False - - if not chunks: - pending_store = ( - None - if skip_save - else self._reuse_strategy().prepare_store(tokens, full_aligned) - ) - if pending_store is not None: - self._pending_alloc[request.request_id] = pending_store - logger.debug("[CONNECTOR] cache miss req=%s", request.request_id[:8]) - return 0, False - - extra_tokens = _contiguous_prefix_tokens(chunks, num_computed_tokens) - if extra_tokens <= 0: - return 0, False - - pending_store = ( - None - if skip_save - else self._reuse_strategy().prepare_store( - tokens, - full_aligned, - chunks, - ) - ) - if pending_store is not None: - self._pending_alloc[request.request_id] = pending_store - - available = len(tokens) - num_computed_tokens - if extra_tokens >= available: - extra_tokens = available - 1 - if extra_tokens <= 0: - return 0, False - - if len(chunks) == 1: - self._pending_loads[request.request_id] = dict( - chunks[0], num_computed_tokens=num_computed_tokens - ) - else: - self._pending_loads[request.request_id] = { - str(i): dict(chunk, num_computed_tokens=num_computed_tokens) - for i, chunk in enumerate(chunks) - } - - logger.debug( - "[CONNECTOR] cache hit req=%s chunks=%d prefix_tokens=%d", - request.request_id[:8], - len(chunks), - extra_tokens, + ) -> tuple[int | None, bool]: + """Return DaseR cache credit for one request.""" + return self._request_lifecycle.get_num_new_matched_tokens( + request, + num_computed_tokens, ) - return extra_tokens, True def update_state_after_alloc( self, @@ -421,517 +54,28 @@ def update_state_after_alloc( blocks: "KVCacheBlocks", num_external_tokens: int, ) -> None: - """Record block IDs for requests that will load or store KV. - - Args: - request: vLLM Request. - blocks: vLLM KV cache block allocation for this request. - num_external_tokens: tokens from DaseR (0 if miss). - """ - req_id = request.request_id - block_ids: list[int] = [blk.block_id for blk in blocks.blocks[0]] - - if req_id in self._pending_loads: - chunks = self._pending_loads[req_id] - if "chunk_key" in chunks: - chunk = chunks - if not _trim_chunk_to_external_window( - chunk=chunk, - block_ids=block_ids, - external_start=int(chunk.get("num_computed_tokens", 0)), - num_external_tokens=num_external_tokens, - block_tokens=self._block_tokens, - slot_size=self._slot_size, - ): - del self._pending_loads[req_id] - self._record_pending_store_blocks(req_id, block_ids) - return - logger.debug( - "[CONNECTOR] load blocks req=%s blocks=%s", - req_id, - chunk["block_ids"], - ) - self._record_pending_store_blocks(req_id, block_ids) - return - for key, chunk in list(chunks.items()): - if not _trim_chunk_to_external_window( - chunk=chunk, - block_ids=block_ids, - external_start=int(chunk.get("num_computed_tokens", 0)), - num_external_tokens=num_external_tokens, - block_tokens=self._block_tokens, - slot_size=self._slot_size, - ): - logger.debug( - "[CONNECTOR] skip load req=%s key=%s target=%d slots=%d", - req_id[:8], - chunk.get("chunk_key", "")[:8], - int(chunk.get("target_token_start", 0)), - int(chunk["num_slots"]), - ) - del chunks[key] - continue - logger.debug( - "[CONNECTOR] load blocks req=%s key=%s blocks=%s", - req_id, - chunk.get("chunk_key", "")[:8], - chunk["block_ids"], - ) - self._record_pending_store_blocks(req_id, block_ids) + """Bind allocated vLLM blocks to pending lifecycle work.""" + self._request_lifecycle.update_state_after_alloc( + request, + blocks, + num_external_tokens, + ) def build_connector_meta( - self, scheduler_output: "SchedulerOutput" - ) -> DaserConnectorMeta: - """Package pending load/store specs into connector metadata. - - Args: - scheduler_output: vLLM SchedulerOutput for this step. - - Returns: - DaserConnectorMeta with reqs_to_load and reqs_to_store. - """ - meta = DaserConnectorMeta() - self._drop_preempted_pending_state(scheduler_output) - scheduled_ids: set[str] = set(scheduler_output.num_scheduled_tokens.keys()) - computed_after = _computed_tokens_after_step(scheduler_output) - self._record_cached_store_blocks(scheduler_output) - - for req_id, chunks in list(self._pending_loads.items()): - if "chunk_key" in chunks: - chunk = chunks - if "block_ids" in chunk: - meta.reqs_to_load[req_id] = _load_spec_from_chunk(chunk) - del self._pending_loads[req_id] - continue - ready = True - load_specs: list[ReqLoadSpec] = [] - for chunk in chunks.values(): - if "block_ids" not in chunk: - ready = False - continue - load_specs.append(_load_spec_from_chunk(chunk)) - merged_specs = _merge_adjacent_load_specs(load_specs, self._slot_size) - for idx, spec in enumerate(merged_specs): - load_id = req_id if len(merged_specs) == 1 else f"{req_id}:load:{idx}" - meta.reqs_to_load[load_id] = spec - if ready: - del self._pending_loads[req_id] - - for req_id, alloc in list(self._pending_stores.items()): - scheduled_tokens = scheduler_output.num_scheduled_tokens.get(req_id, 0) - if scheduled_tokens <= 0 and ":store:" in req_id: - scheduled_tokens = scheduler_output.num_scheduled_tokens.get( - _base_req_id(req_id), - 0, - ) - base_req_id = _base_req_id(req_id) - computed_tokens = computed_after.get(base_req_id, scheduled_tokens) - slot_index = _store_slot_index(req_id) - required_tokens = ( - (slot_index + 1) * self._block_tokens - if slot_index is not None - else int(alloc["token_count"]) - ) - should_store = ( - base_req_id in scheduled_ids - and scheduled_tokens > 0 - and computed_tokens >= required_tokens - and "block_ids" in alloc - ) - if should_store: - meta.reqs_to_store[req_id] = ReqStoreSpec( - chunk_key=alloc["chunk_key"], - start_slot=alloc["start_slot"], - num_slots=alloc["num_slots"], - block_ids=alloc["block_ids"], - file_offset=alloc["file_offset"], - token_count=alloc["token_count"], - ) - del self._pending_stores[req_id] - - if meta.reqs_to_store: - meta.reqs_to_store = self._filter_live_store_specs(meta.reqs_to_store) - pending_async_saves = self._pending_async_save_ids() - for req_id in meta.reqs_to_store: - pending_async_saves.add(_base_req_id(req_id)) - - if logger.isEnabledFor(logging.DEBUG): - for req_id, spec in meta.reqs_to_load.items(): - logger.debug( - "[CONNECTOR] meta LOAD req=%s start_slot=%d blocks=%d tokens=%d", - req_id[:8], - spec.start_slot, - len(spec.block_ids), - spec.token_count, - ) - for req_id, spec in meta.reqs_to_store.items(): - logger.debug( - "[CONNECTOR] meta STORE req=%s start_slot=%d blocks=%d tokens=%d", - req_id[:8], - spec.start_slot, - len(spec.block_ids), - spec.token_count, - ) - return meta - - def _drop_preempted_pending_state( self, scheduler_output: "SchedulerOutput", - ) -> None: - """Discard pending scheduler state whose KV blocks were preempted. - - Args: - scheduler_output: vLLM SchedulerOutput for this step. - - Async/thread-safety: - Runs on the scheduler thread before metadata is handed to workers. - """ - preempted_req_ids = getattr(scheduler_output, "preempted_req_ids", set()) - pending_async_saves = self._pending_async_save_ids() - for req_id in preempted_req_ids: - base_req_id = str(req_id) - pending_async_saves.discard(base_req_id) - for pending_req_id in list(self._pending_loads): - if _matches_request_or_store_id(pending_req_id, base_req_id): - self._pending_loads.pop(pending_req_id, None) - for pending_req_id in list(self._pending_stores): - if _matches_request_or_store_id(pending_req_id, base_req_id): - self._drop_pending_store(pending_req_id) - for pending_req_id in list(self._pending_alloc): - if _matches_request_or_store_id(pending_req_id, base_req_id): - self._pending_alloc.pop(pending_req_id, None) - - def _filter_live_store_specs( - self, - specs: dict[str, ReqStoreSpec], - ) -> dict[str, ReqStoreSpec]: - """Drop store specs whose server allocation was already evicted. - - Args: - specs: Store specs built for the current scheduler step. - - Returns: - Specs that still own their allocated server slot ranges. - """ - try: - live_keys = self._ipc_sync.live_allocations( - [ - { - "chunk_key": spec.chunk_key, - "start_slot": spec.start_slot, - "num_slots": spec.num_slots, - } - for spec in specs.values() - ] - ) - except Exception as exc: - logger.warning("[CONNECTOR] live_allocations failed: %s", exc) - return specs - return { - req_id: spec - for req_id, spec in specs.items() - if spec.chunk_key in live_keys - } - - def _record_cached_store_blocks(self, scheduler_output: "SchedulerOutput") -> None: - """Append blocks from later chunked-prefill steps to store trackers. - - Args: - scheduler_output: vLLM SchedulerOutput for this step. - """ - cached_reqs = getattr(scheduler_output, "scheduled_cached_reqs", None) - if cached_reqs is None: - return - - req_ids = getattr(cached_reqs, "req_ids", []) - new_block_ids = getattr(cached_reqs, "new_block_ids", []) - for req_id, block_groups in zip(req_ids, new_block_ids, strict=False): - pending_store = self._pending_alloc.get(req_id) - if pending_store is None or block_groups is None: - continue - if not block_groups: - continue - block_group = block_groups[0] - if block_group is None: - continue - block_ids = list(block_group) - if ( - req_id in getattr(cached_reqs, "resumed_req_ids", set()) - and block_ids[: len(pending_store.block_ids)] == pending_store.block_ids - ): - pending_store.block_ids = block_ids - else: - pending_store.block_ids.extend(block_ids) - needed_slots = math.ceil(pending_store.token_count / self._block_tokens) - if len(pending_store.block_ids) > needed_slots: - pending_store.block_ids = pending_store.block_ids[:needed_slots] - self._maybe_allocate_pending_store(req_id, pending_store) - - def _lookup_with_external_prefix_metrics( - self, - tokens: list[int], - model_id: str, - queries: int, - num_computed_tokens: int, - ) -> list[dict[str, Any]]: - """Run lookup while passing vLLM external-prefix query count when supported. - - Args: - tokens: token prefix sent to DaseR lookup. - model_id: model identifier. - queries: vLLM external prefix query token count. - num_computed_tokens: tokens already computed locally by vLLM. - - Returns: - Matching chunk dicts returned by the IPC client. - - Thread-safety: - Runs on the scheduler thread and uses the synchronous IPC client. - """ - try: - return self._ipc_sync.lookup( - tokens, - model_id, - external_prefix_queries=queries, - num_computed_tokens=num_computed_tokens, - ) - except TypeError as exc: - if "external_prefix_queries" not in str(exc): - raise - return self._ipc_sync.lookup(tokens, model_id) - - def _record_external_prefix_cache_miss(self, queries: int) -> None: - """Record a connector external-prefix miss when lookup is skipped. - - Args: - queries: vLLM external prefix query token count. - - Returns: - None. - - Thread-safety: - Runs on the scheduler thread and uses the synchronous IPC client - when it supports the diagnostic operation. - """ - recorder = getattr(self._ipc_sync, "record_external_prefix_cache", None) - if recorder is None: - return - recorder(max(0, int(queries)), 0) - - def _record_pending_store_blocks(self, req_id: str, block_ids: list[int]) -> None: - """Record request KV blocks for a pending scheduler store. - - Args: - req_id: vLLM request ID. - block_ids: KV block IDs allocated for the request. - """ - pending_store = self._pending_alloc.get(req_id) - if pending_store is None: - return - requested_tokens = pending_store.token_count - pending_store.block_ids = block_ids[ - : math.ceil(requested_tokens / self._block_tokens) - ] - self._maybe_allocate_pending_store(req_id, pending_store) - - def _init_reuse_strategy(self) -> None: - """Initialize the scheduler cache reuse strategy from current config.""" - self._cache_reuse_strategy = build_cache_reuse_strategy( - "chunk", - self._block_tokens, - ) - - def _reuse_strategy(self) -> Any: - """Return the configured cache reuse strategy. - - Returns: - Cache reuse strategy initialized from connector runtime config. - """ - strategy = getattr(self, "_cache_reuse_strategy", None) - if strategy is None: - self._init_reuse_strategy() - strategy = self._cache_reuse_strategy - return strategy - - def allocate_store_chunk( - self, - chunk_key: str, - token_count: int, - ) -> dict[str, Any]: - """Allocate server metadata for a pending scheduler store. - - Args: - chunk_key: cache key to allocate. - token_count: number of tokens covered by the allocation. - - Returns: - Mutable server allocation metadata. - """ - return self._ipc_sync.alloc_chunk(chunk_key, token_count, self._model_id) - - def allocate_store_chunks( - self, - chunks: list[dict[str, Any]], - ) -> list[dict[str, Any]]: - """Allocate server metadata for multiple pending scheduler stores. - - Args: - chunks: chunk descriptors with chunk_key and token_count. - - Returns: - Mutable server allocation metadata for each chunk. - """ - alloc_chunks = getattr(self._ipc_sync, "alloc_chunks", None) - if alloc_chunks is None: - return [ - { - **self.allocate_store_chunk( - str(chunk["chunk_key"]), - int(chunk["token_count"]), - ), - "chunk_key": str(chunk["chunk_key"]), - } - for chunk in chunks - ] - return alloc_chunks(chunks, self._model_id) - - def set_pending_store(self, req_id: str, alloc: dict[str, Any]) -> None: - """Record an allocated store for later connector metadata packaging. - - Args: - req_id: vLLM request ID or synthetic store work ID. - alloc: mutable allocation metadata. - """ - self._pending_stores[req_id] = alloc - - def has_pending_store(self, req_id: str) -> bool: - """Return whether a pending store entry already exists. - - Args: - req_id: vLLM request ID or synthetic store work ID. - - Returns: - True when the store is already pending. - """ - return req_id in self._pending_stores - - def drop_pending_alloc(self, req_id: str) -> None: - """Remove pending allocation state for a request. - - Args: - req_id: vLLM request ID. - """ - self._pending_alloc.pop(req_id, None) - - def _drop_pending_store(self, req_id: str) -> None: - """Remove a pending store and release its server writer claim. - - Args: - req_id: vLLM request ID or synthetic store work ID. - """ - alloc = self._pending_stores.pop(req_id, None) - if alloc is None: - return - release = getattr(self._ipc_sync, "release_chunk_writer", None) - if release is None: - return - try: - release( - str(alloc["chunk_key"]), - int(alloc["start_slot"]), - int(alloc["num_slots"]), - ) - except Exception as exc: # noqa: BLE001 - logger.warning("[CONNECTOR] release_chunk_writer failed: %s", exc) - - def _discard_pending_request(self, req_id: str) -> None: - """Clear scheduler-side pending state for a request. - - Args: - req_id: vLLM request ID. - """ - self._pending_loads.pop(req_id, None) - if req_id in self._pending_stores: - self._drop_pending_store(req_id) - for pending_req_id in list(self._pending_stores): - if pending_req_id.startswith(f"{req_id}:store:"): - self._drop_pending_store(pending_req_id) - self._pending_alloc.pop(req_id, None) - - def _pending_async_save_ids(self) -> set[str]: - """Return request IDs whose worker-side saves are still pending. - - Returns: - Mutable set of base vLLM request IDs. - - Thread-safety: - Runs on the scheduler thread. The lazy initialization supports - tests and mixin probes that do not call ``DaserConnector.__init__``. - """ - pending = getattr(self, "_pending_async_saves", None) - if pending is None: - pending = set() - self._pending_async_saves = pending - return pending - - def _maybe_allocate_pending_store( - self, req_id: str, pending_store: PendingStore - ) -> None: - """Allocate a DaseR chunk once a pending store has full KV coverage. - - Args: - req_id: vLLM request ID being tracked. - pending_store: store tracker for the request. - """ - requested_tokens = pending_store.token_count - strategy = self._reuse_strategy() - if not strategy.ready_to_allocate(pending_store): - return - tokens = self._req_tokens.get(req_id, []) - if len(tokens) < requested_tokens: - return - strategy.allocate_store(self, req_id, pending_store, tokens) + ) -> DaserConnectorMeta: + """Build worker metadata from request lifecycle state.""" + return self._request_lifecycle.build_connector_meta(scheduler_output) def request_finished( self, request: "Request", block_ids: list[int], - ) -> "tuple[bool, dict[str, Any] | None]": - """Clean up per-request state after inference completes. - - Args: - request: finished vLLM Request. - block_ids: block IDs being freed. - - Returns: - (True, None) when DaseR is still storing this request's KV blocks, - otherwise (False, None). - """ - del block_ids - if request.request_id in self._pending_async_save_ids(): - return True, None - self._req_tokens.pop(request.request_id, None) - self._discard_pending_request(request.request_id) - return False, None + ) -> tuple[bool, dict[str, Any] | None]: + """Return whether worker save completion still holds request blocks.""" + return self._request_lifecycle.request_finished(request, block_ids) def update_connector_output(self, connector_output: Any) -> None: - """Update scheduler state from worker-side transfer completions. - - Args: - connector_output: vLLM KVConnectorOutput carrying finished request - IDs from workers. - - Async/thread-safety: - Runs on vLLM's scheduler thread after worker connector polling. - """ - pending_async_saves = self._pending_async_save_ids() - for req_id in getattr(connector_output, "finished_sending", None) or (): - pending_async_saves.discard(req_id) - self._req_tokens.pop(req_id, None) - self._discard_pending_request(req_id) - for req_id in getattr(connector_output, "finished_recving", None) or (): - if req_id not in self._pending_alloc: - self._req_tokens.pop(req_id, None) - for pending_req_id in list(self._pending_loads): - if _matches_request_or_store_id(pending_req_id, req_id): - self._pending_loads.pop(pending_req_id, None) + """Apply worker transfer completions to request lifecycle state.""" + self._request_lifecycle.update_connector_output(connector_output) diff --git a/daser/connector/scheduler_planning.py b/daser/connector/scheduler_planning.py new file mode 100644 index 0000000..986011f --- /dev/null +++ b/daser/connector/scheduler_planning.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 + +# Standard +import math +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + # Third Party + from vllm.v1.core.scheduler import SchedulerOutput + from vllm.v1.request import Request + +# First Party +from daser.connector.helpers import base_req_id +from daser.connector.metadata import ReqLoadSpec +from daser.logging import init_logger + +logger = init_logger(__name__) + + +def _base_req_id(req_id: str) -> str: + """Compatibility wrapper for tests importing the scheduler-private helper.""" + return base_req_id(req_id) + + +def _store_slot_index(req_id: str) -> int | None: + """Return the rolling-prefix store slot encoded in a synthetic request ID. + + Args: + req_id: vLLM request ID or ``:store:`` synthetic ID. + + Returns: + Slot index for synthetic store IDs, or None for regular IDs. + """ + if ":store:" not in req_id: + return None + try: + return int(req_id.rsplit(":store:", 1)[1]) + except ValueError: + return None + + +def _matches_request_or_store_id(req_id: str, base_req_id: str) -> bool: + """Return whether ``req_id`` belongs to a base request. + + Args: + req_id: vLLM request ID or synthetic connector work ID. + base_req_id: Base vLLM request ID to match. + + Returns: + True when ``req_id`` is the base request or one of its synthetic + store/load entries. + """ + return ( + req_id == base_req_id + or req_id.startswith(f"{base_req_id}:store:") + or req_id.startswith(f"{base_req_id}:load:") + ) + + +def _computed_tokens_after_step( + scheduler_output: "SchedulerOutput", +) -> dict[str, int]: + """Return per-request token counts that are valid after this step. + + Args: + scheduler_output: vLLM SchedulerOutput for this step. + + Returns: + Mapping from request ID to ``num_computed_tokens + scheduled_tokens``. + Falls back to the scheduled token count when older or test scheduler + outputs do not expose prior computed-token metadata. + """ + scheduled = dict(getattr(scheduler_output, "num_scheduled_tokens", {})) + computed_after = {req_id: int(tokens) for req_id, tokens in scheduled.items()} + + for req_data in getattr(scheduler_output, "scheduled_new_reqs", []) or []: + req_id = str(getattr(req_data, "req_id", "")) + if req_id in scheduled: + computed_after[req_id] = int( + getattr(req_data, "num_computed_tokens", 0) + ) + int(scheduled[req_id]) + + cached_reqs = getattr(scheduler_output, "scheduled_cached_reqs", None) + if cached_reqs is not None: + req_ids = getattr(cached_reqs, "req_ids", []) + prior_counts = getattr(cached_reqs, "num_computed_tokens", []) + for req_id, prior in zip(req_ids, prior_counts, strict=False): + req_id = str(req_id) + if req_id in scheduled: + computed_after[req_id] = int(prior) + int(scheduled[req_id]) + + return computed_after + + +def _get_kv_transfer_flag(request: "Request", key: str) -> Any: + """Return ``request.kv_transfer_params[key]`` if present, else ``None``. + + Args: + request: vLLM ``Request`` or compatible object. + key: connector-specific flag name to extract. + + Returns: + The value under ``key``, or ``None`` when absent. + """ + params = getattr(request, "kv_transfer_params", None) + if not isinstance(params, dict): + return None + return params.get(key) + + +def _block_ids_for_chunk( + block_ids: list[int], + target_token_start: int, + num_slots: int, + block_tokens: int, + max_tokens: int | None = None, +) -> list[int]: + """Return vLLM block IDs for a chunk's target prompt range. + + Args: + block_ids: all block IDs allocated to the request. + target_token_start: token offset where the chunk starts in the prompt. + num_slots: number of blocks/slots covered by the chunk. + block_tokens: tokens per vLLM block. + max_tokens: optional upper bound on accepted external tokens. + + Returns: + Slice of block_ids for the chunk, or an empty list when the range + is not block-aligned or exceeds the allocated blocks. + """ + if target_token_start % block_tokens != 0: + return [] + target_block_start = target_token_start // block_tokens + effective_slots = num_slots + if max_tokens is not None: + remaining_tokens = max_tokens - target_token_start + if remaining_tokens <= 0: + return [] + effective_slots = min(num_slots, math.ceil(remaining_tokens / block_tokens)) + target_block_end = target_block_start + effective_slots + if target_block_start < 0 or target_block_end > len(block_ids): + return [] + return block_ids[target_block_start:target_block_end] + + +def _trim_chunk_to_external_window( + chunk: dict[str, Any], + block_ids: list[int], + external_start: int, + num_external_tokens: int, + block_tokens: int, + slot_size: int, +) -> bool: + """Trim chunk metadata to the external token interval vLLM requested. + + Args: + chunk: Mutable chunk metadata returned by the server. + block_ids: Full vLLM block allocation for the request. + external_start: Token offset where external KV loading begins. + num_external_tokens: Number of tokens accepted from the connector. + block_tokens: Tokens per vLLM block. + slot_size: Bytes per DaseR slot. + + Returns: + True when the chunk still covers at least one whole KV block. + """ + if external_start % block_tokens != 0 or num_external_tokens <= 0: + return False + target_start = int(chunk.get("target_token_start", 0)) + target_end = target_start + int(chunk["token_count"]) + external_end = external_start + num_external_tokens + load_start = max(target_start, external_start) + load_end = min(target_end, external_end) + load_start = ((load_start + block_tokens - 1) // block_tokens) * block_tokens + load_end = ((load_end + block_tokens - 1) // block_tokens) * block_tokens + load_end = min(load_end, target_end) + if load_end <= load_start: + return False + + skip_slots = (load_start - target_start) // block_tokens + num_slots = (load_end - load_start) // block_tokens + if load_start < external_start: + return False + block_start = load_start // block_tokens + block_end = block_start + num_slots + if block_start < 0 or block_end > len(block_ids): + return False + + chunk["start_slot"] = int(chunk["start_slot"]) + skip_slots + chunk["file_offset"] = int(chunk["file_offset"]) + skip_slots * slot_size + chunk["num_slots"] = num_slots + chunk["token_count"] = num_slots * block_tokens + chunk["target_token_start"] = load_start + chunk["block_ids"] = block_ids[block_start:block_end] + return bool(chunk["block_ids"]) + + +def _contiguous_prefix_tokens( + chunks: list[dict[str, Any]], num_computed_tokens: int +) -> int: + """Return external tokens covered contiguously after computed tokens. + + Args: + chunks: server chunk payloads with target_token_start and token_count. + num_computed_tokens: tokens vLLM already has locally. + + Returns: + Number of additional contiguous prefix tokens covered by chunks. + """ + covered_until = num_computed_tokens + for chunk in sorted( + chunks, + key=lambda item: int(item.get("target_token_start", 0)), + ): + target_start = int(chunk.get("target_token_start", 0)) + token_count = int(chunk["token_count"]) + target_end = target_start + token_count + if target_end <= covered_until: + continue + if target_start > covered_until: + break + covered_until = target_end + return covered_until - num_computed_tokens + + +def _load_spec_from_chunk(chunk: dict[str, Any]) -> ReqLoadSpec: + """Build a worker load specification from scheduler chunk metadata. + + Args: + chunk: Chunk metadata returned by the server and annotated with vLLM + block IDs during allocation. + + Returns: + ReqLoadSpec consumed by the worker load path. + + Async/thread-safety: + Pure scheduler-thread helper; it does not mutate connector state. + """ + return ReqLoadSpec( + chunk_key=str(chunk["chunk_key"]), + start_slot=int(chunk["start_slot"]), + num_slots=int(chunk["num_slots"]), + block_ids=list(chunk["block_ids"]), + file_offset=int(chunk["file_offset"]), + token_count=int(chunk["token_count"]), + target_token_start=int(chunk.get("target_token_start", 0)), + pos_offset=int(chunk.get("pos_offset", 0)), + ) + + +def _merge_adjacent_load_specs( + specs: list[ReqLoadSpec], + slot_size: int, +) -> list[ReqLoadSpec]: + """Merge adjacent load specs that describe one continuous KV byte range. + + Args: + specs: Load specs for one request in prompt order. + slot_size: Bytes represented by one DaseR KV slot. + + Returns: + Coalesced load specs. Chunk keys from the first spec in a run are kept + only as diagnostics; the worker load path addresses data by byte range. + + Async/thread-safety: + Pure scheduler-thread helper; it does not mutate connector state. + """ + merged: list[ReqLoadSpec] = [] + for spec in specs: + if not spec.block_ids: + continue + if not merged: + merged.append(spec) + continue + prev = merged[-1] + prev_slots = len(prev.block_ids) + adjacent = ( + prev.pos_offset == spec.pos_offset + and prev.start_slot + prev_slots == spec.start_slot + and prev.file_offset + prev_slots * slot_size == spec.file_offset + and prev.target_token_start + prev.token_count == spec.target_token_start + ) + if not adjacent: + merged.append(spec) + continue + merged[-1] = ReqLoadSpec( + chunk_key=prev.chunk_key, + start_slot=prev.start_slot, + num_slots=prev_slots + len(spec.block_ids), + block_ids=[*prev.block_ids, *spec.block_ids], + file_offset=prev.file_offset, + token_count=prev.token_count + spec.token_count, + target_token_start=prev.target_token_start, + pos_offset=prev.pos_offset, + ) + return merged diff --git a/daser/connector/staging.py b/daser/connector/staging.py index 4e20cbd..d573b9d 100644 --- a/daser/connector/staging.py +++ b/daser/connector/staging.py @@ -2,16 +2,10 @@ from __future__ import annotations -# Standard -from collections.abc import Callable -from dataclasses import dataclass, replace -from typing import Any, Protocol - # Third Party import torch # First Party -from daser.connector.metadata import ReqLoadSpec, ReqStoreSpec, StoreWriteSpan from daser.logging import init_logger from daser.ops.rope_apply import ( apply_rope_delta_to_key_block as _apply_rope_delta_to_key_block, @@ -25,9 +19,6 @@ ) DEFAULT_ROPE_DELTA_SCALE = 1.0 -DEFAULT_STORE_STAGING_BYTES = 1536 << 20 -DEFAULT_PENDING_STORE_STAGING_BYTES = 3072 << 20 -MIN_STORE_STAGING_BYTES = 64 << 20 CROSS_LAYER_KV_CACHE_KEY = "__cross_layers__" FUSED_RESTORE_MIN_SLOTS = 32 @@ -38,261 +29,6 @@ ] = {} -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 a worker-side staging pool. - - Args: - pool: Owning pool that will receive the allocation on release. - tensor: Backing tensor, possibly larger than ``nbytes``. - nbytes: Logical byte count used by the current transfer. - - Async/thread-safety: - The lease is released on the vLLM worker thread after transfer - completion. It must not be reused while an async store future owns it. - """ - - pool: _CudaStagingLeaseOwner - tensor: torch.Tensor - nbytes: int - _released: bool = False - - @property - def view(self) -> torch.Tensor: - """Return the logical byte view for the active transfer. - - Returns: - A 1-D uint8 tensor slice with ``nbytes`` elements. - - Async/thread-safety: - The returned tensor remains valid until ``release`` is called. - """ - return self.tensor[: self.nbytes] - - def release(self) -> None: - """Return the lease to the owning pool. - - Async/thread-safety: - Call once no CUDA IPC transfer can access ``view``. - """ - if self._released: - return - self._released = True - self.pool.release(self) - - -class FixedCudaStagingPool: - """Fixed-size worker-side CUDA staging buffers. - - 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 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__( - 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._buffer_bytes = buffer_bytes - 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: - """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_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. - - 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, - 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``. - - 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_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) - - 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._buffers[index], - nbytes=nbytes, - ) - - def release(self, lease: CudaStagingLease) -> None: - """Return a fixed staging lease to the free list. - - Args: - lease: Lease previously returned by ``acquire``. - """ - 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") - - -StoreCudaStagingPool = FixedCudaStagingPool - - -@dataclass(frozen=True) -class StagedStoreBatch: - """Worker-owned CUDA staging batch ready for server-side transfer. - - Args: - buffer: Logical uint8 staging view exported through CUDA IPC. - ready_event: CUDA event recorded after KV -> staging copies. - spans: Server write spans targeting ``buffer``. - lease: Optional reusable staging lease backing ``buffer``. - - Async/thread-safety: - The batch crosses to the connector background loop; ``lease`` is - released only after the async server transfer future completes. - """ - - buffer: torch.Tensor - ready_event: torch.cuda.Event | None - spans: list[StoreWriteSpan] - lease: CudaStagingLease | None - - -@dataclass(frozen=True) -class LoadCopyRun: - """One contiguous staging range that can be copied with one layer loop.""" - - start: int - end: int - block_ids: list[int] - pos_offset: int - - -def _load_source_key(spec: ReqLoadSpec, nbytes: int) -> tuple[str, int, int, int, int]: - """Return the source identity for a load span.""" - return (spec.chunk_key, spec.start_slot, spec.num_slots, spec.file_offset, nbytes) - - -def _store_source_key(spec: ReqStoreSpec) -> tuple[str, int, int, int, int]: - """Return the destination identity for a full store spec.""" - return ( - spec.chunk_key, - spec.start_slot, - spec.num_slots, - spec.file_offset, - len(spec.block_ids), - ) - - def synchronize_cuda_tensor(tensor: torch.Tensor) -> None: """Synchronize pending CUDA work for a tensor before cross-process handoff. @@ -339,41 +75,6 @@ def contiguous_block_range(block_ids: list[int]) -> tuple[int, int] | None: return start, start + len(block_ids) -def derive_store_staging_limits(device: torch.device) -> tuple[int, int]: - """Return bounded GPU staging caps for a CUDA device. - - Args: - device: Device that will own worker-side staging tensors. - - Returns: - ``(single_batch_bytes, pending_bytes)``. The cap is based on both total - and currently free VRAM after vLLM has allocated KV cache. Defaults are - intentionally modest because staging is an IPC transport buffer, not a - persistent cache tier. - - Async/thread-safety: - Reads CUDA device properties only; safe during worker initialization. - """ - if device.type != "cuda": - return DEFAULT_STORE_STAGING_BYTES, DEFAULT_PENDING_STORE_STAGING_BYTES - props = torch.cuda.get_device_properties(device) - total = int(props.total_memory) - try: - free, _ = torch.cuda.mem_get_info(device) - free = int(free) - except (RuntimeError, TypeError, ValueError): - free = total - batch = min( - DEFAULT_STORE_STAGING_BYTES, - max(MIN_STORE_STAGING_BYTES, min(total // 50, free // 10)), - ) - pending = min( - DEFAULT_PENDING_STORE_STAGING_BYTES, - max(batch, min(total // 25, free // 5)), - ) - return batch, pending - - def apply_rope_delta_to_key_block( key_block: torch.Tensor, delta: int, @@ -864,252 +565,3 @@ def copy_cross_layer_kv_cache_to_staging( *src.shape[2:], ) dst.copy_(src) - - -def build_load_read_plan( - reqs_to_load: dict[str, ReqLoadSpec], - slot_size: int, - 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 - combined staging tensor and per-request ranges map slices back to - their original load specs. - """ - total_bytes = 0 - spans: list[dict[str, int]] = [] - per_req_ranges: list[Any] = [] - source_ranges: dict[tuple[str, int, int, int, int], tuple[int, int]] = {} - for req_id, spec in reqs_to_load.items(): - num_slots = len(spec.block_ids) - if num_slots == 0: - continue - nbytes = num_slots * slot_size - source_key = _load_source_key(spec, nbytes) - existing = source_ranges.get(source_key) - if existing is None: - start = total_bytes - end = start + nbytes - spans.append( - { - "target_offset": start, - "nbytes": nbytes, - "file_offset": spec.file_offset, - } - ) - source_ranges[source_key] = (start, end) - total_bytes = end - else: - start, end = existing - 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 - - -def build_load_read_batches( - reqs_to_load: dict[str, ReqLoadSpec], - slot_size: int, - max_batch_bytes: int, - 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 - split at block boundaries when one request exceeds the staging cap. - - Async/thread-safety: - Pure CPU helper. It does not mutate connector state. - """ - if slot_size <= 0: - raise ValueError("slot_size must be positive") - 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[Any]]] = [] - current: dict[str, ReqLoadSpec] = {} - current_slots = 0 - synthetic_id = 0 - - def flush() -> None: - nonlocal current, current_slots - if current: - batches.append( - build_load_read_plan( - current, - slot_size, - include_req_ids=include_req_ids, - ) - ) - current = {} - current_slots = 0 - - for req_id, spec in reqs_to_load.items(): - cursor = 0 - while cursor < len(spec.block_ids): - if current_slots >= max_slots: - flush() - available = max_slots - current_slots - take = min(available, len(spec.block_ids) - cursor) - if take <= 0: - flush() - continue - part = spec.block_ids[cursor : cursor + take] - batch_spec = replace( - spec, - start_slot=spec.start_slot + cursor, - num_slots=take, - block_ids=part, - file_offset=spec.file_offset + cursor * slot_size, - ) - key = ( - req_id - if cursor == 0 and take == len(spec.block_ids) - else (f"{req_id}#{synthetic_id}") - ) - synthetic_id += 1 - current[key] = batch_spec - current_slots += take - cursor += take - flush() - return batches - - -def build_load_copy_runs( - per_req_ranges: list[tuple[int, int, ReqLoadSpec]], -) -> list[LoadCopyRun]: - """Merge adjacent load ranges that share the same KV transform. - - Args: - per_req_ranges: Per-request staging ranges from ``build_load_read_plan``. - - Returns: - Ordered copy runs. Each run covers a contiguous staging slice and the - matching flattened block ID list. - - Async/thread-safety: - Pure CPU helper. It does not mutate connector state. - """ - runs: list[LoadCopyRun] = [] - run_start = -1 - run_end = -1 - run_pos_offset = 0 - run_block_ids: list[int] = [] - - def flush() -> None: - nonlocal run_start, run_end, run_pos_offset, run_block_ids - if run_start >= 0 and run_block_ids: - runs.append( - LoadCopyRun( - start=run_start, - end=run_end, - block_ids=run_block_ids, - pos_offset=run_pos_offset, - ) - ) - run_start = -1 - run_end = -1 - run_pos_offset = 0 - run_block_ids = [] - - 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: - run_end = end - run_block_ids.extend(spec.block_ids) - continue - flush() - run_start = start - run_end = end - run_pos_offset = spec.pos_offset - run_block_ids = list(spec.block_ids) - flush() - return runs - - -def build_staging_store_batches( - reqs_to_store: dict[str, ReqStoreSpec], - slot_size: int, - max_batch_bytes: int = DEFAULT_STORE_STAGING_BYTES, -) -> list[tuple[list[int], list[StoreWriteSpan]]]: - """Split store requests into bounded slot-major staging batches. - - Args: - reqs_to_store: Store specs keyed by request ID. - slot_size: DaseR bytes per KV slot. - max_batch_bytes: Maximum GPU staging bytes per batch. - - Returns: - List of ``(block_ids, spans)`` batches. Span source offsets are relative - to that batch's staging tensor. - - Async/thread-safety: - Pure CPU helper. It does not mutate connector state. - """ - if slot_size <= 0: - raise ValueError("slot_size must be positive") - max_slots = max(1, max_batch_bytes // slot_size) - batches: list[tuple[list[int], list[StoreWriteSpan]]] = [] - batch_blocks: list[int] = [] - batch_spans: list[StoreWriteSpan] = [] - written_specs: set[tuple[str, int, int, int, int]] = set() - - def flush_batch() -> None: - nonlocal batch_blocks, batch_spans - if batch_blocks: - batches.append((batch_blocks, batch_spans)) - batch_blocks = [] - batch_spans = [] - - for spec in reqs_to_store.values(): - source_key = _store_source_key(spec) - if source_key in written_specs: - continue - written_specs.add(source_key) - cursor = 0 - while cursor < len(spec.block_ids): - if len(batch_blocks) >= max_slots: - flush_batch() - available = max_slots - len(batch_blocks) - take = min(available, len(spec.block_ids) - cursor) - if take <= 0: - flush_batch() - continue - source_slot = len(batch_blocks) - part = spec.block_ids[cursor : cursor + take] - batch_blocks.extend(part) - batch_spans.append( - StoreWriteSpan( - source_offset=source_slot * slot_size, - nbytes=take * slot_size, - file_offset=spec.file_offset + cursor * slot_size, - chunk_key=spec.chunk_key, - start_slot=spec.start_slot, - num_slots=spec.num_slots, - ) - ) - cursor += take - flush_batch() - return batches diff --git a/daser/connector/store_pipeline.py b/daser/connector/store_pipeline.py new file mode 100644 index 0000000..42d2080 --- /dev/null +++ b/daser/connector/store_pipeline.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +import threading +from typing import Any + +import torch + +from daser.connector.ipc_client import IPCClientAsync +from daser.connector.metadata import ReqStoreSpec, StoreWriteSpan +from daser.connector.worker_memory import ( + DEFAULT_STORE_STAGING_BYTES, + CudaStagingLease, +) + + +class StorePipeline: + """Own store-side asyncio, IPC, staging, and future state. + + Args: + socket_path: DaseR server Unix socket path. + + Async/thread-safety: + Construction and public methods run on the vLLM worker thread. Async + IPC runs exclusively on the private ``daser-store-io`` thread. + """ + + def __init__(self, socket_path: str) -> None: + self._client = IPCClientAsync(socket_path) + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread( + target=self._run_loop, + daemon=True, + name="daser-store-io", + ) + self.save_futures: list[Any] = [] + self.pending_staging_bytes = 0 + self.staging_bytes = 0 + self.pending_staging_limit_bytes = 0 + self.staging_pool: Any | None = None + self.pending_finished_saves: dict[str, Any] = {} + self._thread.start() + + @property + def client(self) -> IPCClientAsync: + """Return the store IPC client for pipeline-owned coroutines.""" + return self._client + + def submit(self, coro: Any) -> Any: + """Submit a coroutine to the private store event loop.""" + return asyncio.run_coroutine_threadsafe(coro, self._loop) + + def shutdown(self) -> None: + """Close the IPC client and stop the private store event loop.""" + self.submit(self._client.close()).result(timeout=5.0) + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=5.0) + + def _run_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + +@dataclass(frozen=True) +class StagedStoreBatch: + """Hold a worker CUDA snapshot until its async store completes.""" + + buffer: torch.Tensor + ready_event: torch.cuda.Event | None + spans: list[StoreWriteSpan] + lease: CudaStagingLease | None + + +def build_staging_store_batches( + reqs_to_store: dict[str, ReqStoreSpec], + slot_size: int, + max_batch_bytes: int = DEFAULT_STORE_STAGING_BYTES, +) -> list[tuple[list[int], list[StoreWriteSpan]]]: + """Split store requests into bounded slot-major staging batches. + + Args: + reqs_to_store: Request IDs mapped to store specifications. + slot_size: Bytes stored for one rank-local KV slot. + max_batch_bytes: Maximum GPU staging bytes in one batch. + + Returns: + Ordered block ID and server write-span batches. + + Async/thread-safety: + Pure CPU planning; safe to call from worker or store-loop threads. + """ + if slot_size <= 0: + raise ValueError("slot_size must be positive") + max_slots = max(1, max_batch_bytes // slot_size) + batches: list[tuple[list[int], list[StoreWriteSpan]]] = [] + batch_blocks: list[int] = [] + batch_spans: list[StoreWriteSpan] = [] + written_specs: set[tuple[str, int, int, int, int]] = set() + + def flush_batch() -> None: + nonlocal batch_blocks, batch_spans + if batch_blocks: + batches.append((batch_blocks, batch_spans)) + batch_blocks = [] + batch_spans = [] + + for spec in reqs_to_store.values(): + source_key = ( + spec.chunk_key, + spec.start_slot, + spec.num_slots, + spec.file_offset, + len(spec.block_ids), + ) + if source_key in written_specs: + continue + written_specs.add(source_key) + cursor = 0 + while cursor < len(spec.block_ids): + if len(batch_blocks) >= max_slots: + flush_batch() + available = max_slots - len(batch_blocks) + take = min(available, len(spec.block_ids) - cursor) + if take <= 0: + flush_batch() + continue + source_slot = len(batch_blocks) + part = spec.block_ids[cursor : cursor + take] + batch_blocks.extend(part) + batch_spans.append( + StoreWriteSpan( + source_offset=source_slot * slot_size, + nbytes=take * slot_size, + file_offset=spec.file_offset + cursor * slot_size, + chunk_key=spec.chunk_key, + start_slot=spec.start_slot, + num_slots=spec.num_slots, + ) + ) + cursor += take + flush_batch() + return batches diff --git a/daser/connector/worker.py b/daser/connector/worker.py index 35b2375..5d9d046 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -2,1434 +2,49 @@ from __future__ import annotations -# Standard -import asyncio -from collections import deque -from dataclasses import dataclass, replace -import os -import queue -import threading -import time from typing import TYPE_CHECKING, Any -# Third Party -import cupy import torch -from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole if TYPE_CHECKING: - # Third Party from vllm.attention import AttentionMetadata from vllm.forward_context import ForwardContext -# First Party -from daser.connector.helpers import base_req_id -from daser.connector.metadata import ( - DaserConnectorMeta, - ReqStoreSpec, - StoreWriteSpan, -) -from daser.connector.staging import ( - CROSS_LAYER_KV_CACHE_KEY, - DEFAULT_PENDING_STORE_STAGING_BYTES, - DEFAULT_STORE_STAGING_BYTES, - FUSED_RESTORE_MIN_SLOTS, - CudaStagingLease, - FixedCudaStagingPool, - StagedStoreBatch, - StoreCudaStagingPool, -) -from daser.connector.staging import ( - build_load_copy_runs as _build_load_copy_runs, -) -from daser.connector.staging import ( - build_load_read_batches as _build_load_read_batches, -) -from daser.connector.staging import ( - build_staging_store_batches as _build_staging_store_batches, -) -from daser.connector.staging import ( - copy_cross_layer_kv_cache_to_staging as _copy_cross_layer_kv_cache_to_staging, -) -from daser.connector.staging import ( - copy_kv_cache_to_staging as _copy_kv_cache_to_staging, -) -from daser.connector.staging import ( - copy_staging_to_kv_cache as _copy_staging_to_kv_cache, -) -from daser.connector.staging import ( - derive_store_staging_limits as _derive_store_staging_limits, -) -from daser.connector.staging import ( - record_cuda_event as _record_cuda_event, -) -from daser.connector.staging import ( - synchronize_cuda_tensor as _synchronize_cuda_tensor, -) -from daser.logging import init_logger -from daser.ops.rope_apply import ( - apply_rope_delta_to_key_block as _apply_rope_delta_to_key_block, -) -from daser.ops.rope_apply import ( - apply_rope_delta_to_kv_key_block_table, - restore_cross_layer_kv_cache_table, -) -from daser.transfer.cuda_ipc import ( - cuda_array_device_id, - cuda_array_pointer, - export_cuda_ipc_handle, -) - -logger = init_logger(__name__) - -_ROPE_WARMUP_BLOCKS = 1 -_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]] - - -def _rank_lane_offset( - start_slot: int, - local_slot_size: int, - rank_stride_bytes: int, - tp_rank: int, -) -> int: - """Return the physical offset for one rank-local logical slot range. - - Args: - start_slot: First server-owned logical slot. - local_slot_size: Bytes stored per logical slot by one TP rank. - rank_stride_bytes: Byte distance between adjacent rank lanes. - tp_rank: Current vLLM tensor-parallel rank. - - Returns: - Physical store offset for ``start_slot`` in ``tp_rank``'s lane. - - Async/thread-safety: - Pure arithmetic used on the worker thread before IPC submission. - """ - return tp_rank * rank_stride_bytes + start_slot * local_slot_size - - -def _local_slot_bytes(connector: Any) -> int: - """Return per-rank slot bytes, falling back for TP=1 test probes.""" - local_slot_size = int(getattr(connector, "_local_slot_size", 0)) - return local_slot_size or int(connector._slot_size) # noqa: SLF001 - - -def _validate_tp_layout( - local_slot_size: int, - storage_slot_size: int, - tp_size: int, - server_tp_size: int, - tp_rank: int, - rank_stride_bytes: int = 0, -) -> None: - """Validate worker KV geometry against the server-owned TP layout. - - Args: - local_slot_size: Slot bytes measured from the worker KV tensor. - storage_slot_size: Aggregate slot bytes reported by the server. - tp_size: vLLM worker tensor-parallel size. - server_tp_size: Tensor-parallel size reported by the server. - tp_rank: Current vLLM tensor-parallel rank. - rank_stride_bytes: Byte distance between server-owned rank lanes. - - Raises: - ValueError: if rank counts or slot geometry do not match. - - Async/thread-safety: - Pure startup validation called before request traffic. - """ - if tp_size <= 0 or not 0 <= tp_rank < tp_size: - raise ValueError(f"invalid TP rank {tp_rank} for size {tp_size}") - if not storage_slot_size: - return - if server_tp_size != tp_size: - raise ValueError( - f"vLLM TP size {tp_size} does not match DaseR TP size {server_tp_size}" - ) - if local_slot_size * tp_size != storage_slot_size: - raise ValueError( - "worker KV slot geometry does not match DaseR storage layout: " - f"local={local_slot_size} tp={tp_size} storage={storage_slot_size}" - ) - if tp_size > 1 and rank_stride_bytes <= 0: - raise ValueError("DaseR runtime config is missing TP rank stride") - - -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) - - -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)) - - -@dataclass -class _DeferredFinishedSave: - """Store work held until vLLM reports a request as finished.""" - - commit_keys: set[str] - reqs_to_store: dict[str, ReqStoreSpec] - submitted: bool = False - future: Any | None = None - - -@dataclass -class _SaveFuture: - """One background save future and the staging it keeps alive. - - Attributes: - future: future returned by ``asyncio.run_coroutine_threadsafe``. - staging_bytes: GPU staging bytes held alive until completion. - lease: optional reusable staging lease released after completion. - """ - - future: Any - staging_bytes: int - lease: CudaStagingLease | None - - def release(self) -> None: - """Release the reusable staging lease, if any.""" - if self.lease is not None: - self.lease.release() - self.lease = None - - -@dataclass -class _PendingLoad: - """One background load future tracked until vLLM can resume the request. - - Attributes: - future: Future running the cross-layer load work. - block_ids: vLLM KV block IDs targeted by the load. - lease: optional staging lease held until the load completes. - """ - - future: Any - block_ids: list[int] - lease: CudaStagingLease | None - - def release(self) -> None: - """Release the reusable staging lease, if any.""" - if self.lease is not None: - self.lease.release() - 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[Any] - staging_lease: CudaStagingLease - future: Any - 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.""" - - 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. - - Args: - message: Error message raised when the future is collected. - - Async/thread-safety: - Immutable testable stand-in for a failed background future. It does not - spawn threads or perform IO. - """ - - def __init__(self, message: str) -> None: - self._message = message - - def done(self) -> bool: - """Return True because this failed future is already complete.""" - return True - - def result(self, timeout: float | None = None) -> None: - """Raise the seeded load-start failure. - - Args: - timeout: Ignored timeout for ``Future`` API compatibility. - """ - del timeout - 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() - - -@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 LoadRequestDispatcher: - """Bound request-level load concurrency by max in-flight and staging depth. - - Args: - max_inflight: Maximum request loads allowed in flight. - staging_depth: Number of fixed staging buffers available to request - loads. - - Async/thread-safety: - The dispatcher object is owned by the connector load asyncio loop. Test - helpers may call its pure synchronous scheduling helpers directly. - """ - - def __init__( - self, - max_inflight: int, - staging_depth: int, - ) -> None: - self._effective_inflight = max(1, min(max_inflight, staging_depth)) - self._free_buffers: deque[int] = deque(range(self._effective_inflight)) - - @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: - 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. - """ - 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 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. - - 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]: - """Return CUDA allocation base pointer and byte offset for ``device_ptr``. - - Args: - device_ptr: CUDA device pointer exported through IPC. - - Returns: - Tuple of ``(allocation_base_ptr, byte_offset)``. - """ - try: - from cuda.bindings import driver as cuda_driver - - result, base_ptr, _allocation_size = cuda_driver.cuMemGetAddressRange( - device_ptr - ) - if result == cuda_driver.CUresult.CUDA_SUCCESS: - base = int(base_ptr) - return base, int(device_ptr) - base - except Exception as exc: # noqa: BLE001 - logger.debug("[CONNECTOR] cuMemGetAddressRange failed: %s", exc) - return int(device_ptr), 0 - - -def _warm_rope_apply_backends( - device: torch.device, - dtype: torch.dtype, - block_tokens: int, - heads: int, - head_dim: int, - rotary_dim: int, - rope_base: float, - is_neox_style: bool, -) -> None: - """Warm dynamic-shape RoPE apply operators. - - Args: - device: device that owns the worker KV cache. - dtype: KV cache dtype. - block_tokens: tokens per cache block. - heads: number of KV heads. - head_dim: per-head dimension. - rotary_dim: number of dimensions covered by RoPE. - rope_base: RoPE theta/base. - is_neox_style: True for split-half rotation, False for interleaved. - - Async/thread-safety: - Runs synchronously during worker KV cache registration, before request - traffic starts. It launches CUDA work on the current stream. TileLang - failures are surfaced to avoid silently entering a slow restore path. - """ - if device.type != "cuda" or rotary_dim <= 0 or head_dim < rotary_dim: - return - sample = torch.empty( - (_ROPE_WARMUP_BLOCKS, block_tokens, heads, head_dim), - dtype=dtype, - device=device, - ) - _apply_rope_delta_to_key_block( - sample, - delta=1, - rope_base=rope_base, - rotary_dim=rotary_dim, - is_neox_style=is_neox_style, - ) - torch.cuda.synchronize(device) - - -def _warm_cross_layer_restore_backends( - device: torch.device, - dtype: torch.dtype, - layers: int, - block_tokens: int, - heads: int, - head_dim: int, - rotary_dim: int, - rope_base: float, - is_neox_style: bool, -) -> None: - """Warm cross-layer staging restore TileLang kernels. - - Args: - device: device that owns the worker KV cache. - dtype: KV cache dtype. - layers: number of model KV layers. - block_tokens: tokens per cache block. - heads: number of KV heads. - head_dim: per-head dimension. - rotary_dim: number of dimensions covered by RoPE. - rope_base: RoPE theta/base. - is_neox_style: True for split-half rotation, False for interleaved. - - Async/thread-safety: - Runs synchronously during worker KV cache registration, before request - traffic starts. TileLang import/compile failures are surfaced to avoid - silently entering a slow restore path. - """ - if device.type != "cuda" or rotary_dim <= 0 or head_dim < rotary_dim: - return - inv_freq = 1.0 / ( - rope_base - ** ( - torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device) - / rotary_dim - ) - ) - freqs = inv_freq - cos_table = freqs.cos().contiguous() - sin_table = freqs.sin().contiguous() - for blocks, use_fused_restore in ( - (_ROPE_WARMUP_BLOCKS, False), - (FUSED_RESTORE_MIN_SLOTS, True), - ): - sample = torch.empty( - blocks, - layers, - 2, - block_tokens, - heads, - head_dim, - dtype=dtype, - device=device, - ) - if use_fused_restore: - dst = torch.empty_like(sample) - restore_cross_layer_kv_cache_table( - sample, - dst, - cos_table=cos_table, - sin_table=sin_table, - rotary_dim=rotary_dim, - is_neox_style=is_neox_style, - ) - else: - apply_rope_delta_to_kv_key_block_table( - sample, - cos_table=cos_table, - sin_table=sin_table, - rotary_dim=rotary_dim, - is_neox_style=is_neox_style, - ) - torch.cuda.synchronize(device) +from daser.connector.metadata import DaserConnectorMeta class WorkerConnectorMixin: - """Worker-role vLLM connector behavior. - - Async/thread-safety: - Public methods are called on vLLM worker threads. Blocking NVMe work is - submitted to the connector's background asyncio loop. - """ + """Adapt vLLM worker hooks to the worker runtime interface.""" def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: - """Register the per-layer KV cache tensors. - - Args: - kv_caches: dict mapping layer_name -> KV tensor. - """ - self._kv_caches = kv_caches - self._layer_names = list(kv_caches.keys()) - self._layer_idx_map = {name: idx for idx, name in enumerate(self._layer_names)} - sample = next(iter(kv_caches.values()), None) - if kv_caches: - ( - self._store_staging_bytes, - self._pending_store_staging_limit_bytes, - ) = _derive_store_staging_limits(sample.device) - logger.info( - "[CONNECTOR] register_kv_caches: %d layers, first shape=%s dtype=%s", - len(kv_caches), - sample.shape, - sample.dtype, - ) - logger.info( - "[CONNECTOR] transient store staging caps: batch=%d pending=%d", - self._store_staging_bytes, - self._pending_store_staging_limit_bytes, - ) - - if self._layer_names and sample is not None: - num_blocks = sample.shape[1] if sample.dim() >= 2 else 1 - layer_size = sample.nbytes // num_blocks - local_slot_size = layer_size * len(self._layer_names) - tp_size = int(getattr(self, "_tp_size", 1)) - _validate_tp_layout( - local_slot_size, - self._slot_size, - tp_size, - int(getattr(self, "_server_tp_size", tp_size)), - int(getattr(self, "_tp_rank", 0)), - int(getattr(self, "_rank_stride_bytes", 0)), - ) - self._local_slot_size = local_slot_size - if self._slot_size == 0: - self._slot_size = local_slot_size * tp_size - logger.info( - "[CONNECTOR] registered local_slot_size=%d from %d layers", - self._local_slot_size, - len(self._layer_names), - ) - - if sample is not None: - self._store_staging_bytes = max( - self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, - self._local_slot_size, - ) - self._store_staging_pool = StoreCudaStagingPool( - device=sample.device, - 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, - buffer_bytes=self._store_staging_bytes, - 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_request_max_inflight=%d load_staging_depth=%d", - self._store_staging_bytes, - self._store_staging_bytes, - self._pending_store_staging_limit_bytes, - _LOAD_REQUEST_MAX_INFLIGHT, - self._load_staging_pool.depth, - ) - if sample.dim() >= 5: - _warm_rope_apply_backends( - device=sample.device, - dtype=sample.dtype, - block_tokens=int(sample.shape[-3]), - heads=int(sample.shape[-2]), - head_dim=int(sample.shape[-1]), - rotary_dim=int(getattr(self, "_rope_rotary_dim", 0)), - rope_base=float(getattr(self, "_rope_base", 10000.0)), - is_neox_style=bool(getattr(self, "_rope_is_neox_style", True)), - ) - - self._init_server_transfer() + """Register per-layer KV tensors with the worker runtime.""" + self._worker_runtime.register_kv_caches(kv_caches) def register_cross_layers_kv_cache( self, kv_cache: torch.Tensor, attn_backend: type[Any], ) -> None: - """Register vLLM's cross-layer KV cache tensor. - - Args: - kv_cache: vLLM tensor whose logical layout starts with - ``[blocks, layers, 2, block_tokens, heads, head_dim]`` for the - NHD layout DaseR requests. - attn_backend: Attention backend that created ``kv_cache``. - - Async/thread-safety: - Called once during worker initialization before request traffic. - """ - kv_cache_config = getattr(self, "_kv_cache_config", None) - layer_names: list[str] = [] - if kv_cache_config is not None: - for group in getattr(kv_cache_config, "kv_cache_groups", []): - layer_names.extend(list(getattr(group, "layer_names", []))) - if not layer_names: - layer_count = int(kv_cache.shape[1]) if kv_cache.dim() >= 2 else 0 - layer_names = [f"layer.{idx}" for idx in range(layer_count)] - self._kv_caches = {CROSS_LAYER_KV_CACHE_KEY: kv_cache} - self._layer_names = layer_names - self._layer_idx_map = {name: idx for idx, name in enumerate(self._layer_names)} - if kv_cache.dim() < 6: - logger.warning( - "[CONNECTOR] cross-layer KV cache has unsupported shape=%s", - tuple(kv_cache.shape), - ) - return - ( - self._store_staging_bytes, - self._pending_store_staging_limit_bytes, - ) = _derive_store_staging_limits(kv_cache.device) - layer_size = kv_cache[0, 0].nbytes - local_slot_size = layer_size * len(self._layer_names) - tp_size = int(getattr(self, "_tp_size", 1)) - _validate_tp_layout( - local_slot_size, - self._slot_size, - tp_size, - int(getattr(self, "_server_tp_size", tp_size)), - int(getattr(self, "_tp_rank", 0)), - int(getattr(self, "_rank_stride_bytes", 0)), - ) - self._local_slot_size = local_slot_size - if self._slot_size == 0: - self._slot_size = local_slot_size * tp_size - logger.info( - "[CONNECTOR] registered cross-layer local_slot_size=%d from %d layers", - self._local_slot_size, - len(self._layer_names), - ) - self._store_staging_bytes = max( - self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, - self._local_slot_size, - ) - self._store_staging_pool = StoreCudaStagingPool( - device=kv_cache.device, - 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, - buffer_bytes=self._store_staging_bytes, - 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_request_max_inflight=%d load_staging_depth=%d", - len(self._layer_names), - tuple(kv_cache.shape), - kv_cache.dtype, - _LOAD_REQUEST_MAX_INFLIGHT, - self._load_staging_pool.depth, - ) - _warm_rope_apply_backends( - device=kv_cache.device, - dtype=kv_cache.dtype, - block_tokens=int(kv_cache.shape[-3]), - heads=int(kv_cache.shape[-2]), - head_dim=int(kv_cache.shape[-1]), - rotary_dim=int(getattr(self, "_rope_rotary_dim", 0)), - rope_base=float(getattr(self, "_rope_base", 10000.0)), - is_neox_style=bool(getattr(self, "_rope_is_neox_style", True)), - ) - _warm_cross_layer_restore_backends( - device=kv_cache.device, - dtype=kv_cache.dtype, - layers=int(kv_cache.shape[1]), - block_tokens=int(kv_cache.shape[-3]), - heads=int(kv_cache.shape[-2]), - head_dim=int(kv_cache.shape[-1]), - rotary_dim=int(getattr(self, "_rope_rotary_dim", 0)), - rope_base=float(getattr(self, "_rope_base", 10000.0)), - is_neox_style=bool(getattr(self, "_rope_is_neox_style", True)), - ) - self._init_server_transfer() + """Register vLLM's cross-layer KV tensor with the worker runtime.""" + self._worker_runtime.register_cross_layers_kv_cache(kv_cache, attn_backend) def bind_connector_metadata(self, connector_metadata: DaserConnectorMeta) -> None: - """Receive scheduler metadata before each forward pass. - - Args: - connector_metadata: DaserConnectorMeta from build_connector_meta. - """ + """Bind one scheduler metadata step to the worker runtime.""" super().bind_connector_metadata(connector_metadata) - self._meta = connector_metadata - self._reap_save_futures(block=False) - self._pending_commits = set() - for spec in connector_metadata.reqs_to_store.values(): - if spec.block_ids: - self._pending_commits.add(spec.chunk_key) + self._worker_runtime.bind_connector_metadata(connector_metadata) def clear_connector_metadata(self) -> None: - """Clear metadata after forward pass completes.""" + """Clear the current metadata step from the worker runtime.""" super().clear_connector_metadata() - self._meta = None + self._worker_runtime.clear_connector_metadata() def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: - """Submit async KV cache loads for cache-hit requests. - - Args: - forward_context: vLLM ForwardContext for this forward pass. - """ - del forward_context, kwargs - if self._meta is None or not self._meta.reqs_to_load: - return - logger.debug( - "[CONNECTOR] start_load_kv: %d reqs to load", - len(self._meta.reqs_to_load), - ) - reqs_to_load = dict(self._meta.reqs_to_load) - if not self._ensure_transfer_ready(): - self._mark_load_start_failed( - reqs_to_load, - "server transfer config is not ready", - ) - return - - num_layers = len(self._layer_names) - if num_layers == 0: - self._mark_load_start_failed(reqs_to_load, "no registered KV cache layers") - return - - sample_tensor = next(iter(self._kv_caches.values()), None) - if sample_tensor is None: - self._mark_load_start_failed(reqs_to_load, "no registered KV cache tensor") - return - - pending_loads = getattr(self, "_pending_loads", None) - if pending_loads is None: - pending_loads = {} - self._pending_loads = pending_loads - 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() - pending_loads[req_id] = _PendingLoad( - future=request_future, - block_ids=list(spec.block_ids), - lease=None, - ) - self._enqueue_load_request( - load_queue, - _QueuedLoadRequest( - req_id=req_id, - spec_id=spec_id, - spec=spec, - future=request_future, - ), - ) - self._ensure_load_request_dispatcher(sample_tensor) - - def _mark_load_start_failed( - self, - reqs_to_load: dict[str, Any], - reason: str, - ) -> None: - """Record failed load submission so vLLM can release waiting requests. - - Args: - reqs_to_load: Load metadata that could not be submitted. - reason: Human-readable failure reason used in diagnostics. - - Async/thread-safety: - Called on the vLLM worker thread before any background load is - started. Completion is later reported through ``get_finished``. - """ - if not reqs_to_load: - return - block_ids = [ - block_id for spec in reqs_to_load.values() for block_id in spec.block_ids - ] - failed_future = _ImmediateLoadError(reason) - 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_id(req_id) for req_id in reqs_to_load}: - pending_loads[req_id] = _PendingLoad( - future=failed_future, - block_ids=list(block_ids), - lease=None, - ) - - 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: - 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: - 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 - - 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 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. - """ - 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: - dispatcher_future = getattr(self, "_load_request_dispatcher_future", None) - if dispatcher_future is not None and not dispatcher_future.done(): - return - self._load_request_dispatcher_future = asyncio.run_coroutine_threadsafe( - self._run_load_request_dispatcher(sample_tensor), - self._load_loop, - ) - - 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 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. - """ - if sample_tensor.device.type == "cuda": - torch.cuda.set_device(sample_tensor.device) - 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, - ) - 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: - try: - item = load_queue.get_nowait() - except asyncio.QueueEmpty: - return - 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 _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, - _local_slot_bytes(self), - ) - pool = FixedCudaStagingPool( - device=sample_tensor.device, - buffer_bytes=buffer_bytes, - 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 - - 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 - - submitted_at = time.perf_counter() - if getattr(self, "_load_staging_registered", False): - transfer_coro = self._transfer_load_registered_cuda( - buffer_index=buffer_index, - producer_pid=os.getpid(), - 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( - buffer_index=buffer_index, - 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, - ) - future = self._submit_load_coroutine(transfer_coro) - 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 _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) - spec = replace( - item.spec, - file_offset=_rank_lane_offset( - item.spec.start_slot, - _local_slot_bytes(self), - int(getattr(self, "_rank_stride_bytes", 0)), - int(getattr(self, "_tp_rank", 0)), - ), - ) - load_batches = _build_load_read_batches( - {item.spec_id: spec}, - _local_slot_bytes(self), - 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, - 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=_local_slot_bytes(self), - 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() + """Submit scheduler-selected cache loads through the load pipeline.""" + self._worker_runtime.start_load_kv(forward_context, **kwargs) def wait_for_layer_load(self, layer_name: str) -> None: - """No-op because async loads complete before vLLM resumes requests. - - Args: - layer_name: ignored. - """ - return + """Observe the runtime's request-level load completion contract.""" + self._worker_runtime.wait_for_layer_load(layer_name) def save_kv_layer( self, @@ -1438,742 +53,31 @@ def save_kv_layer( attn_metadata: "AttentionMetadata", **kwargs: Any, ) -> None: - """Submit this layer's KV blocks for server-owned transfer. - - Args: - layer_name: name of the current attention layer. - kv_layer: full KV cache tensor for this layer. - attn_metadata: attention metadata (not directly used). - """ - if self._meta is None or not self._meta.reqs_to_store: - return - if not self._ensure_transfer_ready(): - return - - if layer_name not in self._layer_idx_map: - logger.warning( - "[CONNECTOR] save_kv_layer: unknown layer %s, skipping", layer_name - ) + """Forward one vLLM save hook to the worker runtime.""" + self._worker_runtime.save_kv_layer( + layer_name, + kv_layer, + attn_metadata, + **kwargs, + ) def wait_for_save(self) -> None: - """Queue stores until vLLM reports request completion.""" - if self._meta is None: - return - - commit_keys = list(self._pending_commits) - reqs_to_store = dict(self._meta.reqs_to_store) - if commit_keys and reqs_to_store: - pending_finished = getattr(self, "_pending_finished_saves", None) - if pending_finished is None: - pending_finished = {} - self._pending_finished_saves = pending_finished - for req_id, spec in reqs_to_store.items(): - base_id = base_req_id(req_id) - save = pending_finished.get(base_id) - if save is None: - save = _DeferredFinishedSave(commit_keys=set(), reqs_to_store={}) - pending_finished[base_id] = save - save.reqs_to_store[req_id] = spec - if spec.chunk_key in commit_keys: - save.commit_keys.add(spec.chunk_key) - self._pending_commits.clear() + """Defer step stores through the worker runtime.""" + self._worker_runtime.wait_for_save() def get_finished( - self, finished_req_ids: set[str] + self, + finished_req_ids: set[str], ) -> tuple[set[str] | None, set[str] | None]: - """Collect completed background transfers after a worker step. - - Args: - finished_req_ids: Request IDs that vLLM finished in this step. - - Returns: - Finished-saving request IDs and finished async-loading request IDs. - """ - self._reap_save_futures(block=False) - finished_recving = self._collect_finished_loads() - pending_finished = getattr(self, "_pending_finished_saves", {}) - if not pending_finished: - return None, finished_recving or None - - finished_sending: set[str] = set() - candidates = set(finished_req_ids) - candidates.update( - req_id for req_id, save in pending_finished.items() if save.submitted - ) - for req_id in list(candidates): - save = pending_finished.get(req_id) - if save is None: - continue - if not save.submitted: - save.future = self._submit_finished_save(save) - save.submitted = True - future = save.future - if future is None: - finished_sending.add(req_id) - del pending_finished[req_id] - elif future.done(): - future.result(timeout=120.0) - finished_sending.add(req_id) - del pending_finished[req_id] - return finished_sending or None, finished_recving or None + """Return completed store and load request IDs from the runtime.""" + return self._worker_runtime.get_finished(finished_req_ids) def get_block_ids_with_load_errors(self) -> set[int]: - """Return and clear block IDs whose async load failed. - - Returns: - vLLM block IDs that should be treated as invalid. - """ - invalid_blocks = getattr(self, "_invalid_load_block_ids", None) - if invalid_blocks is None: - self._invalid_load_block_ids = set() - return set() - invalid = set(invalid_blocks) - invalid_blocks.clear() - return invalid - - def _collect_finished_loads(self) -> set[str]: - """Poll async load futures without blocking. - - Returns: - Base request IDs whose async load future completed in this poll. - """ - pending_loads = getattr(self, "_pending_loads", {}) - if not pending_loads: - return set() - - finished_recving: set[str] = set() - collected_futures: set[int] = set() - for req_id, load in list(pending_loads.items()): - if not load.future.done(): - continue - future_id = id(load.future) - try: - if future_id not in collected_futures: - load.future.result() - collected_futures.add(future_id) - except Exception as exc: # noqa: BLE001 - logger.warning( - "[CONNECTOR] async load failed req=%s blocks=%s: %s", - req_id, - load.block_ids, - exc, - ) - self._invalid_load_block_ids.update(load.block_ids) - collected_futures.add(future_id) - finally: - load.release() - del pending_loads[req_id] - finished_recving.add(req_id) - return finished_recving + """Return and clear runtime load-error block IDs.""" + return self._worker_runtime.get_block_ids_with_load_errors() def shutdown(self) -> None: - """Stop the background IO loop.""" - if self._role != KVConnectorRole.WORKER: - return - pending_loads = getattr(self, "_pending_loads", {}) - for load in {id(load.future): load for load in pending_loads.values()}.values(): - if not load.future.done(): - try: - load.future.result(timeout=120.0) - except Exception: # noqa: BLE001 - pass - self._collect_finished_loads() - 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: - 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: - dispatcher_future.result(timeout=120.0) - except Exception: # noqa: BLE001 - pass - 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) - 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) - self._store_loop.call_soon_threadsafe(self._store_loop.stop) - self._load_thread.join(timeout=5) - self._store_thread.join(timeout=5) - - def _run_load_loop(self) -> None: - """Run the foreground load asyncio IO loop.""" - asyncio.set_event_loop(self._load_loop) - self._load_loop.run_forever() - - def _run_store_loop(self) -> None: - """Run the background store asyncio IO loop.""" - asyncio.set_event_loop(self._store_loop) - self._store_loop.run_forever() - - def _submit_load_coroutine(self, coro: Any) -> Any: - """Submit foreground load work to the dedicated load event loop. - - Args: - coro: Coroutine object to schedule. - - Returns: - Future returned by ``asyncio.run_coroutine_threadsafe``. - - Async/thread-safety: - Called from vLLM worker threads. A load-only loop prevents cache-hit - reads from queueing behind background store coroutines. - """ - loop = self._load_loop - return asyncio.run_coroutine_threadsafe(coro, loop) - - def _submit_store_coroutine(self, coro: Any) -> Any: - """Submit background store and commit work to the store event loop. - - Args: - coro: Coroutine object to schedule. - - Returns: - Future returned by ``asyncio.run_coroutine_threadsafe``. - - Async/thread-safety: - Called from vLLM worker threads. Store work is serialized on the - store loop and does not occupy the foreground load loop. - """ - loop = self._store_loop - return asyncio.run_coroutine_threadsafe(coro, loop) - - def _ensure_transfer_ready(self) -> bool: - """Refresh server transfer config and mark worker data plane ready.""" - if getattr(self, "_transfer_ready", False): - return True - - self._refresh_runtime_config() - if not self._slot_size or ( - not getattr(self, "_skip_l2", False) and not self._store_path - ): - logger.warning( - "[CONNECTOR] server transfer config is not ready; start DaseR server " - "before sending requests", - ) - return False - - _validate_tp_layout( - _local_slot_bytes(self), - self._slot_size, - int(getattr(self, "_tp_size", 1)), - int(getattr(self, "_server_tp_size", 1)), - int(getattr(self, "_tp_rank", 0)), - int(getattr(self, "_rank_stride_bytes", 0)), - ) - - self._transfer_ready = True - logger.info("[CONNECTOR] server transfer mode=%s", self._transfer_mode) - return True - - def _init_server_transfer(self) -> None: - """Initialize the server-owned transfer layer on both IO loops. - - Async/thread-safety: - Called from a vLLM worker thread after KV-cache registration. Does - nothing until the server transfer config and both async IO loops - are ready. - """ - if not ( - self._ensure_transfer_ready() - and getattr(self, "_ipc_load_async", None) is not None - and getattr(self, "_ipc_store_async", None) is not None - and getattr(self, "_load_loop", None) is not None - and getattr(self, "_store_loop", None) is not None - ): - return - 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. - - 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 not self._load_ipc_clients(): - return - try: - for buffer_index in range(pool.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) - load_client = self._load_ipc_client_for_buffer(buffer_index) - self._submit_load_coroutine( - load_client.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", - pool.depth, - ) - - def _reap_save_futures(self, block: bool) -> None: - """Collect completed background save tasks. - - Args: - block: If True, wait for every pending save. If False, collect only - tasks that are already complete. - """ - remaining: list[_SaveFuture] = [] - pending_bytes = self._pending_save_staging_bytes - for record in self._save_futures: - if block or record.future.done(): - try: - record.future.result(timeout=120.0) - finally: - pending_bytes = max(0, pending_bytes - record.staging_bytes) - record.release() - else: - remaining.append(record) - self._save_futures = remaining - self._pending_save_staging_bytes = pending_bytes - - def _track_save_future( - self, - future: Any, - staging_bytes: int, - staging_lease: CudaStagingLease | None, - ) -> None: - """Track one background save future and its live staging bytes. - - Args: - future: Future returned by ``asyncio.run_coroutine_threadsafe``. - staging_bytes: GPU staging bytes kept alive by the future. - staging_lease: Optional reusable staging lease to release after - ``future`` completes. - - Async/thread-safety: - Called on the worker thread. Completion is collected by - ``_reap_save_futures``. - """ - self._pending_save_staging_bytes += staging_bytes - self._save_futures.append( - _SaveFuture(future=future, staging_bytes=staging_bytes, lease=staging_lease) - ) - - def _wait_for_save_staging_capacity(self, nbytes: int) -> None: - """Apply backpressure before allocating another store staging buffer. - - Args: - nbytes: Size of the next staging tensor. - - Async/thread-safety: - Called by vLLM's worker thread. It may wait for already-submitted - background stores when live staging would exceed the configured - cap. - """ - limit = max( - self._pending_store_staging_limit_bytes - or DEFAULT_PENDING_STORE_STAGING_BYTES, - nbytes, - ) - while self._pending_save_staging_bytes + nbytes > limit and self._save_futures: - 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 _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 oldest store futures until a lease is released. - """ - pool = getattr(self, "_store_staging_pool", None) - self._wait_for_save_staging_capacity(nbytes) - while pool is not None and pool.available == 0 and self._save_futures: - 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, - device: torch.device, - ) -> CudaStagingLease: - """Acquire a reusable staging buffer for a CUDA IPC transfer. - - Args: - nbytes: Logical byte count needed for the transfer. - device: Device used when the pool has not been initialized yet. - - Returns: - A staging lease whose ``view`` is safe to export through CUDA IPC. - - Async/thread-safety: - Called from the worker thread. Store-path callers must retain the - lease until the background server transfer completes. - """ - 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 = StoreCudaStagingPool( - device=device, - 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, - wait_for_release=lambda: self._wait_for_store_staging_release(nbytes), - ) - - def _stage_store_batch( - self, - block_ids: list[int], - spans: list[StoreWriteSpan], - ) -> StagedStoreBatch | None: - """Snapshot one bounded batch of KV blocks into CUDA staging. - - Args: - block_ids: vLLM KV block IDs to snapshot. - spans: Server store spans targeting this staging batch. - - Returns: - A staged batch ready for CUDA IPC transfer, or ``None`` when the - connector has no layer state. - - Async/thread-safety: - Runs on the vLLM worker thread so KV cache reads are launched before - vLLM can recycle the source blocks. The returned tensor is kept - alive by the background transfer future. - """ - num_layers = len(self._layer_names) - if num_layers == 0: - return None - sample_tensor = next(iter(self._kv_caches.values()), None) - if sample_tensor is None: - return None - if not block_ids or not spans: - return None - local_slot_size = _local_slot_bytes(self) - nbytes = len(block_ids) * local_slot_size - self._wait_for_save_staging_capacity(nbytes) - staging_lease = self._acquire_staging(nbytes, sample_tensor.device) - staging = staging_lease.view - block_index = torch.tensor( - block_ids, - dtype=torch.long, - device=sample_tensor.device, - ) - cross_layer_kv_cache = self._kv_caches.get(CROSS_LAYER_KV_CACHE_KEY) - if cross_layer_kv_cache is not None: - _copy_cross_layer_kv_cache_to_staging( - staging=staging, - kv_cache=cross_layer_kv_cache, - block_ids=block_ids, - num_layers=num_layers, - slot_size=local_slot_size, - block_index=block_index, - ) - else: - for layer_name in self._layer_names: - _copy_kv_cache_to_staging( - staging=staging, - kv_layer=self._kv_caches[layer_name], - layer_idx=self._layer_idx_map[layer_name], - block_ids=block_ids, - num_layers=num_layers, - slot_size=local_slot_size, - block_index=block_index, - ) - return StagedStoreBatch( - buffer=staging, - ready_event=_record_cuda_event(staging), - spans=spans, - lease=staging_lease, - ) - - def _submit_finished_save(self, save: _DeferredFinishedSave) -> Any | None: - """Submit one request's deferred KV store after request completion. - - Args: - save: Deferred store plan built during ``wait_for_save``. - - Returns: - Future that completes after all store batches are committed, or - ``None`` when no store batch could be staged. - - Async/thread-safety: - Called by vLLM's worker thread from ``get_finished`` while vLLM is - still holding the finished request's KV blocks. - """ - batch_futures = [] - reqs_to_store = { - req_id: replace( - spec, - file_offset=_rank_lane_offset( - spec.start_slot, - _local_slot_bytes(self), - int(getattr(self, "_rank_stride_bytes", 0)), - int(getattr(self, "_tp_rank", 0)), - ), - ) - for req_id, spec in save.reqs_to_store.items() - } - batches = _build_staging_store_batches( - reqs_to_store, - _local_slot_bytes(self), - max_batch_bytes=(self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES), - ) - for block_ids, spans in batches: - staged = self._stage_store_batch(block_ids, spans) - if staged is None: - continue - future = self._submit_store_coroutine( - self._write_cuda_buffer( - buffer=staged.buffer, - ready_event=staged.ready_event, - spans=staged.spans, - ) - ) - self._track_save_future(future, staged.buffer.nbytes, staged.lease) - batch_futures.append(future) - if not batch_futures: - return None - commit_future = self._submit_store_coroutine( - self._commit_after_store_futures(batch_futures, sorted(save.commit_keys)), - ) - self._track_save_future(commit_future, 0, None) - return commit_future - - async def _commit_after_store_futures( - self, - batch_futures: list[Any], - commit_keys: list[str], - ) -> None: - """Commit chunks after all staged transfer batches finish. - - Args: - batch_futures: Futures for each staged store batch. - commit_keys: Chunk keys to publish after stores complete. - - Async/thread-safety: - Runs on the connector background event loop and does not read vLLM - KV cache tensors. - """ - stored_keys: list[str] = [] - for future in batch_futures: - stored_keys.extend(await asyncio.wrap_future(future)) - await self._commit_stored_keys(stored_keys, commit_keys) - - async def _write_cuda_buffer( - self, - buffer: torch.Tensor, - ready_event: torch.cuda.Event | None, - spans: list[StoreWriteSpan], - ) -> list[str]: - """Write selected spans from one contiguous CUDA buffer. - - Args: - buffer: CUDA tensor exported over CUDA IPC. - ready_event: Producer-stream event for ``buffer``. - spans: Source/destination write spans. - - Returns: - Chunk keys accepted by the server for this buffer. - """ - if buffer.device.type == "cuda": - torch.cuda.set_device(buffer.device) - if ready_event is not None: - ready_event.synchronize() - else: - _synchronize_cuda_tensor(buffer) - cp_buffer = cupy.asarray(buffer) - cuda_ipc_handle = export_cuda_ipc_handle(cp_buffer) - device_id = cuda_array_device_id(cp_buffer) - device_ptr = cuda_array_pointer(cp_buffer) - ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) - stored_keys = await self._transfer_store_cuda( - cuda_ipc_handle=cuda_ipc_handle, - nbytes=buffer.nbytes, - device_id=device_id, - device_ptr=device_ptr, - allocation_base_ptr=ipc_base_ptr, - allocation_offset=ipc_offset, - producer_pid=os.getpid(), - spans=[ - { - "source_offset": span.source_offset, - "nbytes": span.nbytes, - "file_offset": span.file_offset, - "chunk_key": span.chunk_key, - "start_slot": span.start_slot, - "num_slots": span.num_slots, - } - for span in spans - ], - ) - return stored_keys - - async def _commit_stored_keys( - self, - stored_keys: list[str], - commit_keys: list[str], - ) -> None: - """Commit requested chunks whose store spans were accepted.""" - requested = set(commit_keys) - candidate_keys = [key for key in stored_keys if key in requested] - keys_to_commit = list(dict.fromkeys(candidate_keys)) - await self._ipc_store_async.commit_chunks( - keys_to_commit, - tp_rank=int(getattr(self, "_tp_rank", 0)), - tp_size=int(getattr(self, "_tp_size", 1)), - ) - - async def _transfer_load_cuda(self, **kwargs: Any) -> dict[str, Any]: - """Load through the dedicated worker load IPC client. - - Args: - **kwargs: forwarded CUDA transfer payload fields. - - Returns: - Server load response with timing counters. - - Async/thread-safety: - Runs on the worker load event loop. A dedicated client keeps - cache-hit loads from queueing behind store RPCs. - """ - 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. - - 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. - """ - 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. - - Args: - **kwargs: forwarded CUDA transfer payload fields. - - Returns: - Chunk keys accepted by the server. - - Async/thread-safety: - Runs on the worker store event loop and serializes only with other - store/commit traffic. - """ - return await self._ipc_store_async.transfer_store_cuda(**kwargs) + """Drain and stop the worker runtime when initialized.""" + runtime = getattr(self, "_worker_runtime", None) + if runtime is not None: + runtime.shutdown() diff --git a/daser/connector/worker_memory.py b/daser/connector/worker_memory.py new file mode 100644 index 0000000..2c0e008 --- /dev/null +++ b/daser/connector/worker_memory.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Bounded worker-side CUDA staging memory.""" + +from __future__ import annotations + +# Standard +from collections.abc import Callable +from dataclasses import dataclass +from typing import Protocol + +# Third Party +import torch + +DEFAULT_STORE_STAGING_BYTES = 1536 << 20 +DEFAULT_PENDING_STORE_STAGING_BYTES = 3072 << 20 +MIN_STORE_STAGING_BYTES = 64 << 20 + + +class _CudaStagingLeaseOwner(Protocol): + """Pool interface required by ``CudaStagingLease``.""" + + 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 a fixed CUDA pool. + + Args: + pool: Owning pool that receives the allocation on release. + tensor: Backing tensor, possibly larger than ``nbytes``. + nbytes: Logical byte count used by the current transfer. + + Async/thread-safety: + The worker thread releases a lease only after every CUDA IPC user has + finished with it. + """ + + pool: _CudaStagingLeaseOwner + tensor: torch.Tensor + nbytes: int + _released: bool = False + + @property + def view(self) -> torch.Tensor: + """Return the active one-dimensional uint8 tensor view.""" + return self.tensor[: self.nbytes] + + def release(self) -> None: + """Return this lease to its owning pool once.""" + if self._released: + return + self._released = True + self.pool.release(self) + + +class FixedCudaStagingPool: + """Preallocate and lease fixed-size worker CUDA staging buffers. + + 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: + One vLLM worker thread owns the pool. It never allocates after + construction; callers release leases after async transfer completes. + """ + + 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._buffer_bytes = buffer_bytes + 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: + """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_indices) + + @property + def depth(self) -> int: + """Return the number of fixed buffers in the pool.""" + return len(self._buffers) + + def buffer(self, index: int) -> torch.Tensor: + """Return a backing tensor for one-time CUDA IPC registration. + + Args: + index: Fixed buffer index. + + Returns: + The full preallocated backing tensor. + + Async/thread-safety: + The caller must not mutate the tensor outside the 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, + 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 every buffer + is leased. + + Returns: + Lease whose view is limited to ``nbytes``. + + Raises: + ValueError: If ``nbytes`` is invalid. + RuntimeError: If no buffer is available after the callback. + """ + 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_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) + + def acquire_index(self, index: int, nbytes: int) -> CudaStagingLease: + """Lease a specific preallocated staging buffer. + + Args: + index: Fixed buffer index. + nbytes: Logical transfer byte count. + + Returns: + Lease whose view is limited to ``nbytes``. + + Raises: + ValueError: If ``index`` or ``nbytes`` is invalid. + RuntimeError: If the requested buffer is already leased. + """ + 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._buffers[index], + nbytes=nbytes, + ) + + def release(self, lease: CudaStagingLease) -> None: + """Return a lease to this pool. + + Args: + lease: Lease previously returned by this pool. + + Raises: + ValueError: If the lease belongs to another pool. + """ + 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") + + +def derive_staging_limits(device: torch.device) -> tuple[int, int]: + """Return bounded staging batch and pending byte limits for a device. + + Args: + device: Device that owns worker-side staging tensors. + + Returns: + Tuple of single-buffer bytes and total pending bytes. + + Async/thread-safety: + Reads CUDA device properties during worker initialization before + request traffic starts. + """ + if device.type != "cuda": + return DEFAULT_STORE_STAGING_BYTES, DEFAULT_PENDING_STORE_STAGING_BYTES + props = torch.cuda.get_device_properties(device) + total = int(props.total_memory) + try: + free, _ = torch.cuda.mem_get_info(device) + free = int(free) + except (RuntimeError, TypeError, ValueError): + free = total + batch = min( + DEFAULT_STORE_STAGING_BYTES, + max(MIN_STORE_STAGING_BYTES, min(total // 50, free // 10)), + ) + pending = min( + DEFAULT_PENDING_STORE_STAGING_BYTES, + max(batch, min(total // 25, free // 5)), + ) + return batch, pending diff --git a/daser/connector/worker_runtime.py b/daser/connector/worker_runtime.py new file mode 100644 index 0000000..14c8e6f --- /dev/null +++ b/daser/connector/worker_runtime.py @@ -0,0 +1,2156 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +# Standard +import asyncio +from collections import deque +from dataclasses import dataclass, replace +import os +import threading +import time +from typing import TYPE_CHECKING, Any + +# Third Party +import cupy +import torch +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole + +if TYPE_CHECKING: + # Third Party + from vllm.attention import AttentionMetadata + from vllm.forward_context import ForwardContext + +# First Party +from daser.connector.helpers import base_req_id +from daser.connector.ipc_client import IPCClientSync +from daser.connector.load_pipeline import LoadPipeline +from daser.connector.load_pipeline import ( + build_load_copy_runs as _build_load_copy_runs, +) +from daser.connector.load_pipeline import ( + build_load_read_batches as _build_load_read_batches, +) +from daser.connector.metadata import ( + DaserConnectorMeta, + ReqStoreSpec, + StoreWriteSpan, +) +from daser.connector.staging import CROSS_LAYER_KV_CACHE_KEY, FUSED_RESTORE_MIN_SLOTS +from daser.connector.staging import ( + copy_cross_layer_kv_cache_to_staging as _copy_cross_layer_kv_cache_to_staging, +) +from daser.connector.staging import ( + copy_kv_cache_to_staging as _copy_kv_cache_to_staging, +) +from daser.connector.staging import ( + copy_staging_to_kv_cache as _copy_staging_to_kv_cache, +) +from daser.connector.staging import ( + record_cuda_event as _record_cuda_event, +) +from daser.connector.staging import ( + synchronize_cuda_tensor as _synchronize_cuda_tensor, +) +from daser.connector.store_pipeline import StagedStoreBatch, StorePipeline +from daser.connector.store_pipeline import ( + build_staging_store_batches as _build_staging_store_batches, +) +from daser.connector.worker_memory import ( + DEFAULT_PENDING_STORE_STAGING_BYTES, + DEFAULT_STORE_STAGING_BYTES, + CudaStagingLease, + FixedCudaStagingPool, +) +from daser.connector.worker_memory import ( + derive_staging_limits as _derive_staging_limits, +) +from daser.logging import init_logger +from daser.ops.rope_apply import ( + apply_rope_delta_to_key_block as _apply_rope_delta_to_key_block, +) +from daser.ops.rope_apply import ( + apply_rope_delta_to_kv_key_block_table, + restore_cross_layer_kv_cache_table, +) +from daser.transfer.cuda_ipc import ( + cuda_array_device_id, + cuda_array_pointer, + export_cuda_ipc_handle, +) + +logger = init_logger(__name__) + +_ROPE_WARMUP_BLOCKS = 1 +_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]] + + +def _rank_lane_offset( + start_slot: int, + local_slot_size: int, + rank_stride_bytes: int, + tp_rank: int, +) -> int: + """Return the physical offset for one rank-local logical slot range. + + Args: + start_slot: First server-owned logical slot. + local_slot_size: Bytes stored per logical slot by one TP rank. + rank_stride_bytes: Byte distance between adjacent rank lanes. + tp_rank: Current vLLM tensor-parallel rank. + + Returns: + Physical store offset for ``start_slot`` in ``tp_rank``'s lane. + + Async/thread-safety: + Pure arithmetic used on the worker thread before IPC submission. + """ + return tp_rank * rank_stride_bytes + start_slot * local_slot_size + + +def _local_slot_bytes(connector: Any) -> int: + """Return per-rank slot bytes, falling back for TP=1 test probes.""" + local_slot_size = int(getattr(connector, "_local_slot_size", 0)) + return local_slot_size or int(connector._slot_size) # noqa: SLF001 + + +def _validate_tp_layout( + local_slot_size: int, + storage_slot_size: int, + tp_size: int, + server_tp_size: int, + tp_rank: int, + rank_stride_bytes: int = 0, +) -> None: + """Validate worker KV geometry against the server-owned TP layout. + + Args: + local_slot_size: Slot bytes measured from the worker KV tensor. + storage_slot_size: Aggregate slot bytes reported by the server. + tp_size: vLLM worker tensor-parallel size. + server_tp_size: Tensor-parallel size reported by the server. + tp_rank: Current vLLM tensor-parallel rank. + rank_stride_bytes: Byte distance between server-owned rank lanes. + + Raises: + ValueError: if rank counts or slot geometry do not match. + + Async/thread-safety: + Pure startup validation called before request traffic. + """ + if tp_size <= 0 or not 0 <= tp_rank < tp_size: + raise ValueError(f"invalid TP rank {tp_rank} for size {tp_size}") + if not storage_slot_size: + return + if server_tp_size != tp_size: + raise ValueError( + f"vLLM TP size {tp_size} does not match DaseR TP size {server_tp_size}" + ) + if local_slot_size * tp_size != storage_slot_size: + raise ValueError( + "worker KV slot geometry does not match DaseR storage layout: " + f"local={local_slot_size} tp={tp_size} storage={storage_slot_size}" + ) + if tp_size > 1 and rank_stride_bytes <= 0: + raise ValueError("DaseR runtime config is missing TP rank stride") + + +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) + + +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)) + + +@dataclass +class _DeferredFinishedSave: + """Store work held until vLLM reports a request as finished.""" + + commit_keys: set[str] + reqs_to_store: dict[str, ReqStoreSpec] + submitted: bool = False + future: Any | None = None + + +@dataclass +class _SaveFuture: + """One background save future and the staging it keeps alive. + + Attributes: + future: future returned by ``asyncio.run_coroutine_threadsafe``. + staging_bytes: GPU staging bytes held alive until completion. + lease: optional reusable staging lease released after completion. + """ + + future: Any + staging_bytes: int + lease: CudaStagingLease | None + + def release(self) -> None: + """Release the reusable staging lease, if any.""" + if self.lease is not None: + self.lease.release() + self.lease = None + + +@dataclass +class _PendingLoad: + """One background load future tracked until vLLM can resume the request. + + Attributes: + future: Future running the cross-layer load work. + block_ids: vLLM KV block IDs targeted by the load. + lease: optional staging lease held until the load completes. + """ + + future: Any + block_ids: list[int] + lease: CudaStagingLease | None + + def release(self) -> None: + """Release the reusable staging lease, if any.""" + if self.lease is not None: + self.lease.release() + 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[Any] + staging_lease: CudaStagingLease + future: Any + 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.""" + + 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. + + Args: + message: Error message raised when the future is collected. + + Async/thread-safety: + Immutable testable stand-in for a failed background future. It does not + spawn threads or perform IO. + """ + + def __init__(self, message: str) -> None: + self._message = message + + def done(self) -> bool: + """Return True because this failed future is already complete.""" + return True + + def result(self, timeout: float | None = None) -> None: + """Raise the seeded load-start failure. + + Args: + timeout: Ignored timeout for ``Future`` API compatibility. + """ + del timeout + 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() + + +@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 LoadRequestDispatcher: + """Bound request-level load concurrency by max in-flight and staging depth. + + Args: + max_inflight: Maximum request loads allowed in flight. + staging_depth: Number of fixed staging buffers available to request + loads. + + Async/thread-safety: + The dispatcher object is owned by the connector load asyncio loop. Test + helpers may call its pure synchronous scheduling helpers directly. + """ + + def __init__( + self, + max_inflight: int, + staging_depth: int, + ) -> None: + self._effective_inflight = max(1, min(max_inflight, staging_depth)) + self._free_buffers: deque[int] = deque(range(self._effective_inflight)) + + @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: + 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. + """ + 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 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. + + 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]: + """Return CUDA allocation base pointer and byte offset for ``device_ptr``. + + Args: + device_ptr: CUDA device pointer exported through IPC. + + Returns: + Tuple of ``(allocation_base_ptr, byte_offset)``. + """ + try: + from cuda.bindings import driver as cuda_driver + + result, base_ptr, _allocation_size = cuda_driver.cuMemGetAddressRange( + device_ptr + ) + if result == cuda_driver.CUresult.CUDA_SUCCESS: + base = int(base_ptr) + return base, int(device_ptr) - base + except Exception as exc: # noqa: BLE001 + logger.debug("[CONNECTOR] cuMemGetAddressRange failed: %s", exc) + return int(device_ptr), 0 + + +def _warm_rope_apply_backends( + device: torch.device, + dtype: torch.dtype, + block_tokens: int, + heads: int, + head_dim: int, + rotary_dim: int, + rope_base: float, + is_neox_style: bool, +) -> None: + """Warm dynamic-shape RoPE apply operators. + + Args: + device: device that owns the worker KV cache. + dtype: KV cache dtype. + block_tokens: tokens per cache block. + heads: number of KV heads. + head_dim: per-head dimension. + rotary_dim: number of dimensions covered by RoPE. + rope_base: RoPE theta/base. + is_neox_style: True for split-half rotation, False for interleaved. + + Async/thread-safety: + Runs synchronously during worker KV cache registration, before request + traffic starts. It launches CUDA work on the current stream. TileLang + failures are surfaced to avoid silently entering a slow restore path. + """ + if device.type != "cuda" or rotary_dim <= 0 or head_dim < rotary_dim: + return + sample = torch.empty( + (_ROPE_WARMUP_BLOCKS, block_tokens, heads, head_dim), + dtype=dtype, + device=device, + ) + _apply_rope_delta_to_key_block( + sample, + delta=1, + rope_base=rope_base, + rotary_dim=rotary_dim, + is_neox_style=is_neox_style, + ) + torch.cuda.synchronize(device) + + +def _warm_cross_layer_restore_backends( + device: torch.device, + dtype: torch.dtype, + layers: int, + block_tokens: int, + heads: int, + head_dim: int, + rotary_dim: int, + rope_base: float, + is_neox_style: bool, +) -> None: + """Warm cross-layer staging restore TileLang kernels. + + Args: + device: device that owns the worker KV cache. + dtype: KV cache dtype. + layers: number of model KV layers. + block_tokens: tokens per cache block. + heads: number of KV heads. + head_dim: per-head dimension. + rotary_dim: number of dimensions covered by RoPE. + rope_base: RoPE theta/base. + is_neox_style: True for split-half rotation, False for interleaved. + + Async/thread-safety: + Runs synchronously during worker KV cache registration, before request + traffic starts. TileLang import/compile failures are surfaced to avoid + silently entering a slow restore path. + """ + if device.type != "cuda" or rotary_dim <= 0 or head_dim < rotary_dim: + return + inv_freq = 1.0 / ( + rope_base + ** ( + torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device) + / rotary_dim + ) + ) + freqs = inv_freq + cos_table = freqs.cos().contiguous() + sin_table = freqs.sin().contiguous() + for blocks, use_fused_restore in ( + (_ROPE_WARMUP_BLOCKS, False), + (FUSED_RESTORE_MIN_SLOTS, True), + ): + sample = torch.empty( + blocks, + layers, + 2, + block_tokens, + heads, + head_dim, + dtype=dtype, + device=device, + ) + if use_fused_restore: + dst = torch.empty_like(sample) + restore_cross_layer_kv_cache_table( + sample, + dst, + cos_table=cos_table, + sin_table=sin_table, + rotary_dim=rotary_dim, + is_neox_style=is_neox_style, + ) + else: + apply_rope_delta_to_kv_key_block_table( + sample, + cos_table=cos_table, + sin_table=sin_table, + rotary_dim=rotary_dim, + is_neox_style=is_neox_style, + ) + torch.cuda.synchronize(device) + + +class WorkerRuntime: + """Own worker KV layout, step metadata, pipelines, and completion state. + + Async/thread-safety: + Public methods are called on vLLM worker threads. Blocking NVMe work is + submitted to the runtime's load or store pipeline loop. + """ + + def __init__( + self, + *, + socket_path: str, + transfer_mode: str, + skip_l2: bool, + tp_size: int, + tp_rank: int, + server_tp_size: int, + slot_size: int, + store_path: str, + rank_stride_bytes: int, + rope_base: float, + rope_rotary_dim: int, + rope_is_neox_style: bool, + rope_delta_scale: float, + load_key_scale: float, + load_value_scale: float, + kv_cache_config: Any, + ) -> None: + self._socket_path = socket_path + self._transfer_mode = transfer_mode + self._skip_l2 = skip_l2 + self._tp_size = tp_size + self._tp_rank = tp_rank + self._server_tp_size = server_tp_size + self._slot_size = slot_size + self._local_slot_size = 0 + self._store_path = store_path + self._rank_stride_bytes = rank_stride_bytes + self._rope_base = rope_base + self._rope_rotary_dim = rope_rotary_dim + self._rope_is_neox_style = rope_is_neox_style + self._rope_delta_scale = rope_delta_scale + self._load_key_scale = load_key_scale + self._load_value_scale = load_value_scale + self._kv_cache_config = kv_cache_config + self._role = KVConnectorRole.WORKER + self._transfer_ready = False + self._load_pipeline = LoadPipeline(socket_path, _LOAD_REQUEST_MAX_INFLIGHT) + self._store_pipeline = StorePipeline(socket_path) + self._kv_caches: dict[str, torch.Tensor] = {} + self._layer_names: list[str] = [] + self._layer_idx_map: dict[str, int] = {} + self._meta: DaserConnectorMeta | None = None + self._pending_commits: set[str] = set() + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + """Register the per-layer KV cache tensors. + + Args: + kv_caches: dict mapping layer_name -> KV tensor. + """ + self._kv_caches = kv_caches + self._layer_names = list(kv_caches.keys()) + self._layer_idx_map = {name: idx for idx, name in enumerate(self._layer_names)} + sample = next(iter(kv_caches.values()), None) + if kv_caches: + ( + self._store_pipeline.staging_bytes, + self._store_pipeline.pending_staging_limit_bytes, + ) = _derive_staging_limits(sample.device) + logger.info( + "[CONNECTOR] register_kv_caches: %d layers, first shape=%s dtype=%s", + len(kv_caches), + sample.shape, + sample.dtype, + ) + logger.info( + "[CONNECTOR] transient store staging caps: batch=%d pending=%d", + self._store_pipeline.staging_bytes, + self._store_pipeline.pending_staging_limit_bytes, + ) + + if self._layer_names and sample is not None: + num_blocks = sample.shape[1] if sample.dim() >= 2 else 1 + layer_size = sample.nbytes // num_blocks + local_slot_size = layer_size * len(self._layer_names) + tp_size = int(getattr(self, "_tp_size", 1)) + _validate_tp_layout( + local_slot_size, + self._slot_size, + tp_size, + int(getattr(self, "_server_tp_size", tp_size)), + int(getattr(self, "_tp_rank", 0)), + int(getattr(self, "_rank_stride_bytes", 0)), + ) + self._local_slot_size = local_slot_size + if self._slot_size == 0: + self._slot_size = local_slot_size * tp_size + logger.info( + "[CONNECTOR] registered local_slot_size=%d from %d layers", + self._local_slot_size, + len(self._layer_names), + ) + + if sample is not None: + self._store_pipeline.staging_bytes = max( + self._store_pipeline.staging_bytes or DEFAULT_STORE_STAGING_BYTES, + self._local_slot_size, + ) + self._store_pipeline.staging_pool = FixedCudaStagingPool( + device=sample.device, + buffer_bytes=self._store_pipeline.staging_bytes, + depth=_store_staging_pool_depth( + self._store_pipeline.staging_bytes, + self._store_pipeline.pending_staging_limit_bytes, + ), + ) + self._load_pipeline.set_staging_pool( + FixedCudaStagingPool( + device=sample.device, + buffer_bytes=self._store_pipeline.staging_bytes, + depth=_load_staging_pool_depth( + self._store_pipeline.staging_bytes, + self._store_pipeline.pending_staging_limit_bytes, + sample.device, + ), + ) + ) + self._load_pipeline.staging_registered = False + logger.info( + "[CONNECTOR] preallocated staging buffer=%d cap=%d pending=%d " + "load_request_max_inflight=%d load_staging_depth=%d", + self._store_pipeline.staging_bytes, + self._store_pipeline.staging_bytes, + self._store_pipeline.pending_staging_limit_bytes, + _LOAD_REQUEST_MAX_INFLIGHT, + self._load_pipeline.staging_pool.depth, + ) + if sample.dim() >= 5: + _warm_rope_apply_backends( + device=sample.device, + dtype=sample.dtype, + block_tokens=int(sample.shape[-3]), + heads=int(sample.shape[-2]), + head_dim=int(sample.shape[-1]), + rotary_dim=int(getattr(self, "_rope_rotary_dim", 0)), + rope_base=float(getattr(self, "_rope_base", 10000.0)), + is_neox_style=bool(getattr(self, "_rope_is_neox_style", True)), + ) + + self._init_server_transfer() + + def register_cross_layers_kv_cache( + self, + kv_cache: torch.Tensor, + attn_backend: type[Any], + ) -> None: + """Register vLLM's cross-layer KV cache tensor. + + Args: + kv_cache: vLLM tensor whose logical layout starts with + ``[blocks, layers, 2, block_tokens, heads, head_dim]`` for the + NHD layout DaseR requests. + attn_backend: Attention backend that created ``kv_cache``. + + Async/thread-safety: + Called once during worker initialization before request traffic. + """ + kv_cache_config = getattr(self, "_kv_cache_config", None) + layer_names: list[str] = [] + if kv_cache_config is not None: + for group in getattr(kv_cache_config, "kv_cache_groups", []): + layer_names.extend(list(getattr(group, "layer_names", []))) + if not layer_names: + layer_count = int(kv_cache.shape[1]) if kv_cache.dim() >= 2 else 0 + layer_names = [f"layer.{idx}" for idx in range(layer_count)] + self._kv_caches = {CROSS_LAYER_KV_CACHE_KEY: kv_cache} + self._layer_names = layer_names + self._layer_idx_map = {name: idx for idx, name in enumerate(self._layer_names)} + if kv_cache.dim() < 6: + logger.warning( + "[CONNECTOR] cross-layer KV cache has unsupported shape=%s", + tuple(kv_cache.shape), + ) + return + ( + self._store_pipeline.staging_bytes, + self._store_pipeline.pending_staging_limit_bytes, + ) = _derive_staging_limits(kv_cache.device) + layer_size = kv_cache[0, 0].nbytes + local_slot_size = layer_size * len(self._layer_names) + tp_size = int(getattr(self, "_tp_size", 1)) + _validate_tp_layout( + local_slot_size, + self._slot_size, + tp_size, + int(getattr(self, "_server_tp_size", tp_size)), + int(getattr(self, "_tp_rank", 0)), + int(getattr(self, "_rank_stride_bytes", 0)), + ) + self._local_slot_size = local_slot_size + if self._slot_size == 0: + self._slot_size = local_slot_size * tp_size + logger.info( + "[CONNECTOR] registered cross-layer local_slot_size=%d from %d layers", + self._local_slot_size, + len(self._layer_names), + ) + self._store_pipeline.staging_bytes = max( + self._store_pipeline.staging_bytes or DEFAULT_STORE_STAGING_BYTES, + self._local_slot_size, + ) + self._store_pipeline.staging_pool = FixedCudaStagingPool( + device=kv_cache.device, + buffer_bytes=self._store_pipeline.staging_bytes, + depth=_store_staging_pool_depth( + self._store_pipeline.staging_bytes, + self._store_pipeline.pending_staging_limit_bytes, + ), + ) + self._load_pipeline.set_staging_pool( + FixedCudaStagingPool( + device=kv_cache.device, + buffer_bytes=self._store_pipeline.staging_bytes, + depth=_load_staging_pool_depth( + self._store_pipeline.staging_bytes, + self._store_pipeline.pending_staging_limit_bytes, + kv_cache.device, + ), + ) + ) + self._load_pipeline.staging_registered = False + logger.info( + "[CONNECTOR] register_cross_layers_kv_cache: layers=%d shape=%s " + "dtype=%s load_request_max_inflight=%d load_staging_depth=%d", + len(self._layer_names), + tuple(kv_cache.shape), + kv_cache.dtype, + _LOAD_REQUEST_MAX_INFLIGHT, + self._load_pipeline.staging_pool.depth, + ) + _warm_rope_apply_backends( + device=kv_cache.device, + dtype=kv_cache.dtype, + block_tokens=int(kv_cache.shape[-3]), + heads=int(kv_cache.shape[-2]), + head_dim=int(kv_cache.shape[-1]), + rotary_dim=int(getattr(self, "_rope_rotary_dim", 0)), + rope_base=float(getattr(self, "_rope_base", 10000.0)), + is_neox_style=bool(getattr(self, "_rope_is_neox_style", True)), + ) + _warm_cross_layer_restore_backends( + device=kv_cache.device, + dtype=kv_cache.dtype, + layers=int(kv_cache.shape[1]), + block_tokens=int(kv_cache.shape[-3]), + heads=int(kv_cache.shape[-2]), + head_dim=int(kv_cache.shape[-1]), + rotary_dim=int(getattr(self, "_rope_rotary_dim", 0)), + rope_base=float(getattr(self, "_rope_base", 10000.0)), + is_neox_style=bool(getattr(self, "_rope_is_neox_style", True)), + ) + self._init_server_transfer() + + def bind_connector_metadata(self, connector_metadata: DaserConnectorMeta) -> None: + """Receive scheduler metadata before each forward pass. + + Args: + connector_metadata: DaserConnectorMeta from build_connector_meta. + """ + self._meta = connector_metadata + self._reap_save_futures(block=False) + self._pending_commits = set() + for spec in connector_metadata.reqs_to_store.values(): + if spec.block_ids: + self._pending_commits.add(spec.chunk_key) + + def clear_connector_metadata(self) -> None: + """Clear metadata after forward pass completes.""" + self._meta = None + + def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: + """Submit async KV cache loads for cache-hit requests. + + Args: + forward_context: vLLM ForwardContext for this forward pass. + """ + del forward_context, kwargs + if self._meta is None or not self._meta.reqs_to_load: + return + logger.debug( + "[CONNECTOR] start_load_kv: %d reqs to load", + len(self._meta.reqs_to_load), + ) + reqs_to_load = dict(self._meta.reqs_to_load) + if not self._ensure_transfer_ready(): + self._mark_load_start_failed( + reqs_to_load, + "server transfer config is not ready", + ) + return + + num_layers = len(self._layer_names) + if num_layers == 0: + self._mark_load_start_failed(reqs_to_load, "no registered KV cache layers") + return + + sample_tensor = next(iter(self._kv_caches.values()), None) + if sample_tensor is None: + self._mark_load_start_failed(reqs_to_load, "no registered KV cache tensor") + return + + pending_loads = self._load_pipeline.pending + 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() + pending_loads[req_id] = _PendingLoad( + future=request_future, + block_ids=list(spec.block_ids), + lease=None, + ) + self._enqueue_load_request( + load_queue, + _QueuedLoadRequest( + req_id=req_id, + spec_id=spec_id, + spec=spec, + future=request_future, + ), + ) + self._ensure_load_request_dispatcher(sample_tensor) + + def _mark_load_start_failed( + self, + reqs_to_load: dict[str, Any], + reason: str, + ) -> None: + """Record failed load submission so vLLM can release waiting requests. + + Args: + reqs_to_load: Load metadata that could not be submitted. + reason: Human-readable failure reason used in diagnostics. + + Async/thread-safety: + Called on the vLLM worker thread before any background load is + started. Completion is later reported through ``get_finished``. + """ + if not reqs_to_load: + return + block_ids = [ + block_id for spec in reqs_to_load.values() for block_id in spec.block_ids + ] + failed_future = _ImmediateLoadError(reason) + pending_loads = self._load_pipeline.pending + for req_id in {base_req_id(req_id) for req_id in reqs_to_load}: + pending_loads[req_id] = _PendingLoad( + future=failed_future, + block_ids=list(block_ids), + lease=None, + ) + + 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: + 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. + """ + del sample_tensor + return self._load_pipeline.ensure_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. + """ + del load_queue + self._load_pipeline.enqueue(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 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. + """ + self._load_pipeline.ensure_dispatcher( + self._run_load_request_dispatcher(sample_tensor) + ) + + 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 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. + """ + if sample_tensor.device.type == "cuda": + torch.cuda.set_device(sample_tensor.device) + 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, + ) + 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: + try: + item = load_queue.get_nowait() + except asyncio.QueueEmpty: + return + 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 _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 = self._load_pipeline.staging_pool + if pool is not None: + return pool + buffer_bytes = max( + self._store_pipeline.staging_bytes or DEFAULT_STORE_STAGING_BYTES, + _local_slot_bytes(self), + ) + pool = FixedCudaStagingPool( + device=sample_tensor.device, + buffer_bytes=buffer_bytes, + depth=_load_staging_pool_depth( + buffer_bytes, + self._store_pipeline.pending_staging_limit_bytes, + sample_tensor.device, + ), + ) + self._load_pipeline.set_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 + + submitted_at = time.perf_counter() + if self._load_pipeline.staging_registered: + transfer_coro = self._transfer_load_registered_cuda( + buffer_index=buffer_index, + producer_pid=os.getpid(), + 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( + buffer_index=buffer_index, + 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, + ) + future = self._submit_load_coroutine(transfer_coro) + 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 _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) + spec = replace( + item.spec, + file_offset=_rank_lane_offset( + item.spec.start_slot, + _local_slot_bytes(self), + int(getattr(self, "_rank_stride_bytes", 0)), + int(getattr(self, "_tp_rank", 0)), + ), + ) + load_batches = _build_load_read_batches( + {item.spec_id: spec}, + _local_slot_bytes(self), + 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, + 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=_local_slot_bytes(self), + 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. + + Args: + layer_name: ignored. + """ + return + + def save_kv_layer( + self, + layer_name: str, + kv_layer: torch.Tensor, + attn_metadata: "AttentionMetadata", + **kwargs: Any, + ) -> None: + """Submit this layer's KV blocks for server-owned transfer. + + Args: + layer_name: name of the current attention layer. + kv_layer: full KV cache tensor for this layer. + attn_metadata: attention metadata (not directly used). + """ + if self._meta is None or not self._meta.reqs_to_store: + return + if not self._ensure_transfer_ready(): + return + + if layer_name not in self._layer_idx_map: + logger.warning( + "[CONNECTOR] save_kv_layer: unknown layer %s, skipping", layer_name + ) + + def wait_for_save(self) -> None: + """Queue stores until vLLM reports request completion.""" + if self._meta is None: + return + + commit_keys = list(self._pending_commits) + reqs_to_store = dict(self._meta.reqs_to_store) + if commit_keys and reqs_to_store: + pending_finished = self._store_pipeline.pending_finished_saves + for req_id, spec in reqs_to_store.items(): + base_id = base_req_id(req_id) + save = pending_finished.get(base_id) + if save is None: + save = _DeferredFinishedSave(commit_keys=set(), reqs_to_store={}) + pending_finished[base_id] = save + save.reqs_to_store[req_id] = spec + if spec.chunk_key in commit_keys: + save.commit_keys.add(spec.chunk_key) + self._pending_commits.clear() + + def get_finished( + self, finished_req_ids: set[str] + ) -> tuple[set[str] | None, set[str] | None]: + """Collect completed background transfers after a worker step. + + Args: + finished_req_ids: Request IDs that vLLM finished in this step. + + Returns: + Finished-saving request IDs and finished async-loading request IDs. + """ + self._reap_save_futures(block=False) + finished_recving = self._collect_finished_loads() + pending_finished = self._store_pipeline.pending_finished_saves + if not pending_finished: + return None, finished_recving or None + + finished_sending: set[str] = set() + candidates = set(finished_req_ids) + candidates.update( + req_id for req_id, save in pending_finished.items() if save.submitted + ) + for req_id in list(candidates): + save = pending_finished.get(req_id) + if save is None: + continue + if not save.submitted: + save.future = self._submit_finished_save(save) + save.submitted = True + future = save.future + if future is None: + finished_sending.add(req_id) + del pending_finished[req_id] + elif future.done(): + future.result(timeout=120.0) + finished_sending.add(req_id) + del pending_finished[req_id] + return finished_sending or None, finished_recving or None + + def get_block_ids_with_load_errors(self) -> set[int]: + """Return and clear block IDs whose async load failed. + + Returns: + vLLM block IDs that should be treated as invalid. + """ + invalid_blocks = self._load_pipeline.invalid_block_ids + invalid = set(invalid_blocks) + invalid_blocks.clear() + return invalid + + def _collect_finished_loads(self) -> set[str]: + """Poll async load futures without blocking. + + Returns: + Base request IDs whose async load future completed in this poll. + """ + pending_loads = self._load_pipeline.pending + if not pending_loads: + return set() + + finished_recving: set[str] = set() + collected_futures: set[int] = set() + for req_id, load in list(pending_loads.items()): + if not load.future.done(): + continue + future_id = id(load.future) + try: + if future_id not in collected_futures: + load.future.result() + collected_futures.add(future_id) + except Exception as exc: # noqa: BLE001 + logger.warning( + "[CONNECTOR] async load failed req=%s blocks=%s: %s", + req_id, + load.block_ids, + exc, + ) + self._load_pipeline.invalid_block_ids.update(load.block_ids) + collected_futures.add(future_id) + finally: + load.release() + del pending_loads[req_id] + finished_recving.add(req_id) + return finished_recving + + def shutdown(self) -> None: + """Stop the background IO loop.""" + if self._role != KVConnectorRole.WORKER: + return + pending_loads = self._load_pipeline.pending + for load in {id(load.future): load for load in pending_loads.values()}.values(): + if not load.future.done(): + try: + load.future.result(timeout=120.0) + except Exception: # noqa: BLE001 + pass + self._collect_finished_loads() + for req_id in list(self._store_pipeline.pending_finished_saves): + self.get_finished({req_id}) + self._reap_save_futures(block=True) + self._load_pipeline.shutdown() + self._store_pipeline.shutdown() + + def _submit_load_coroutine(self, coro: Any) -> Any: + """Submit foreground load work to the dedicated load event loop. + + Args: + coro: Coroutine object to schedule. + + Returns: + Future returned by ``asyncio.run_coroutine_threadsafe``. + + Async/thread-safety: + Called from vLLM worker threads. A load-only loop prevents cache-hit + reads from queueing behind background store coroutines. + """ + return self._load_pipeline.submit(coro) + + def _submit_store_coroutine(self, coro: Any) -> Any: + """Submit background store and commit work to the store event loop. + + Args: + coro: Coroutine object to schedule. + + Returns: + Future returned by ``asyncio.run_coroutine_threadsafe``. + + Async/thread-safety: + Called from vLLM worker threads. Store work is serialized on the + store loop and does not occupy the foreground load loop. + """ + return self._store_pipeline.submit(coro) + + def _ensure_transfer_ready(self) -> bool: + """Refresh server transfer config and mark worker data plane ready.""" + if getattr(self, "_transfer_ready", False): + return True + + self._refresh_runtime_config() + if not self._slot_size or ( + not getattr(self, "_skip_l2", False) and not self._store_path + ): + logger.warning( + "[CONNECTOR] server transfer config is not ready; start DaseR server " + "before sending requests", + ) + return False + + _validate_tp_layout( + _local_slot_bytes(self), + self._slot_size, + int(getattr(self, "_tp_size", 1)), + int(getattr(self, "_server_tp_size", 1)), + int(getattr(self, "_tp_rank", 0)), + int(getattr(self, "_rank_stride_bytes", 0)), + ) + + self._transfer_ready = True + logger.info("[CONNECTOR] server transfer mode=%s", self._transfer_mode) + return True + + def _refresh_runtime_config(self) -> None: + """Refresh worker-owned storage geometry directly over sync IPC.""" + client = IPCClientSync(self._socket_path) + try: + config = client.get_runtime_config() + except Exception as exc: # noqa: BLE001 + logger.info("[CONNECTOR] runtime config unavailable: %s", exc) + return + finally: + client.close() + self._store_path = str(config.get("store_path", self._store_path)) + self._slot_size = int(config.get("slot_size", self._slot_size)) + self._server_tp_size = int( + config.get("tensor_parallel_size", self._server_tp_size) + ) + self._rank_stride_bytes = int( + config.get("rank_stride_bytes", self._rank_stride_bytes) + ) + self._skip_l2 = bool(config.get("skip_l2", self._skip_l2)) + self._transfer_mode = str(config.get("transfer_mode", self._transfer_mode)) + + def _init_server_transfer(self) -> None: + """Initialize the server-owned transfer layer on both IO loops. + + Async/thread-safety: + Called from a vLLM worker thread after KV-cache registration. Does + nothing until the server transfer config and both async IO loops + are ready. + """ + if not self._ensure_transfer_ready(): + return + 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._store_pipeline.client.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. + """ + return list(self._load_pipeline.clients()) + + 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. + """ + return self._load_pipeline.client(buffer_index) + + 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 self._load_pipeline.staging_registered: + return + pool = self._load_pipeline.staging_pool + if pool is None or not self._load_ipc_clients(): + return + try: + for buffer_index in range(pool.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) + load_client = self._load_ipc_client_for_buffer(buffer_index) + self._submit_load_coroutine( + load_client.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_pipeline.staging_registered = False + return + self._load_pipeline.staging_registered = True + logger.info( + "[CONNECTOR] registered %d fixed load staging buffers", + pool.depth, + ) + + def _reap_save_futures(self, block: bool) -> None: + """Collect completed background save tasks. + + Args: + block: If True, wait for every pending save. If False, collect only + tasks that are already complete. + """ + remaining: list[_SaveFuture] = [] + pending_bytes = self._store_pipeline.pending_staging_bytes + for record in self._store_pipeline.save_futures: + if block or record.future.done(): + try: + record.future.result(timeout=120.0) + finally: + pending_bytes = max(0, pending_bytes - record.staging_bytes) + record.release() + else: + remaining.append(record) + self._store_pipeline.save_futures = remaining + self._store_pipeline.pending_staging_bytes = pending_bytes + + def _track_save_future( + self, + future: Any, + staging_bytes: int, + staging_lease: CudaStagingLease | None, + ) -> None: + """Track one background save future and its live staging bytes. + + Args: + future: Future returned by ``asyncio.run_coroutine_threadsafe``. + staging_bytes: GPU staging bytes kept alive by the future. + staging_lease: Optional reusable staging lease to release after + ``future`` completes. + + Async/thread-safety: + Called on the worker thread. Completion is collected by + ``_reap_save_futures``. + """ + self._store_pipeline.pending_staging_bytes += staging_bytes + self._store_pipeline.save_futures.append( + _SaveFuture(future=future, staging_bytes=staging_bytes, lease=staging_lease) + ) + + def _wait_for_save_staging_capacity(self, nbytes: int) -> None: + """Apply backpressure before allocating another store staging buffer. + + Args: + nbytes: Size of the next staging tensor. + + Async/thread-safety: + Called by vLLM's worker thread. It may wait for already-submitted + background stores when live staging would exceed the configured + cap. + """ + limit = max( + self._store_pipeline.pending_staging_limit_bytes + or DEFAULT_PENDING_STORE_STAGING_BYTES, + nbytes, + ) + while ( + self._store_pipeline.pending_staging_bytes + nbytes > limit + and self._store_pipeline.save_futures + ): + record = self._store_pipeline.save_futures.pop(0) + try: + record.future.result(timeout=120.0) + finally: + self._store_pipeline.pending_staging_bytes = max( + 0, + self._store_pipeline.pending_staging_bytes - record.staging_bytes, + ) + 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 oldest store futures until a lease is released. + """ + pool = self._store_pipeline.staging_pool + self._wait_for_save_staging_capacity(nbytes) + while ( + pool is not None + and pool.available == 0 + and self._store_pipeline.save_futures + ): + record = self._store_pipeline.save_futures.pop(0) + try: + record.future.result(timeout=120.0) + finally: + self._store_pipeline.pending_staging_bytes = max( + 0, + self._store_pipeline.pending_staging_bytes - record.staging_bytes, + ) + record.release() + self._reap_save_futures(block=False) + + def _acquire_staging( + self, + nbytes: int, + device: torch.device, + ) -> CudaStagingLease: + """Acquire a reusable staging buffer for a CUDA IPC transfer. + + Args: + nbytes: Logical byte count needed for the transfer. + device: Device used when the pool has not been initialized yet. + + Returns: + A staging lease whose ``view`` is safe to export through CUDA IPC. + + Async/thread-safety: + Called from the worker thread. Store-path callers must retain the + lease until the background server transfer completes. + """ + pool = self._store_pipeline.staging_pool + if pool is None: + max_bytes = max( + nbytes, + self._store_pipeline.staging_bytes or DEFAULT_STORE_STAGING_BYTES, + ) + pool = FixedCudaStagingPool( + device=device, + buffer_bytes=max_bytes, + depth=_store_staging_pool_depth( + max_bytes, + self._store_pipeline.pending_staging_limit_bytes + or DEFAULT_PENDING_STORE_STAGING_BYTES, + ), + ) + self._store_pipeline.staging_pool = pool + return pool.acquire( + nbytes, + wait_for_release=lambda: self._wait_for_store_staging_release(nbytes), + ) + + def _stage_store_batch( + self, + block_ids: list[int], + spans: list[StoreWriteSpan], + ) -> StagedStoreBatch | None: + """Snapshot one bounded batch of KV blocks into CUDA staging. + + Args: + block_ids: vLLM KV block IDs to snapshot. + spans: Server store spans targeting this staging batch. + + Returns: + A staged batch ready for CUDA IPC transfer, or ``None`` when the + connector has no layer state. + + Async/thread-safety: + Runs on the vLLM worker thread so KV cache reads are launched before + vLLM can recycle the source blocks. The returned tensor is kept + alive by the background transfer future. + """ + num_layers = len(self._layer_names) + if num_layers == 0: + return None + sample_tensor = next(iter(self._kv_caches.values()), None) + if sample_tensor is None: + return None + if not block_ids or not spans: + return None + local_slot_size = _local_slot_bytes(self) + nbytes = len(block_ids) * local_slot_size + self._wait_for_save_staging_capacity(nbytes) + staging_lease = self._acquire_staging(nbytes, sample_tensor.device) + staging = staging_lease.view + block_index = torch.tensor( + block_ids, + dtype=torch.long, + device=sample_tensor.device, + ) + cross_layer_kv_cache = self._kv_caches.get(CROSS_LAYER_KV_CACHE_KEY) + if cross_layer_kv_cache is not None: + _copy_cross_layer_kv_cache_to_staging( + staging=staging, + kv_cache=cross_layer_kv_cache, + block_ids=block_ids, + num_layers=num_layers, + slot_size=local_slot_size, + block_index=block_index, + ) + else: + for layer_name in self._layer_names: + _copy_kv_cache_to_staging( + staging=staging, + kv_layer=self._kv_caches[layer_name], + layer_idx=self._layer_idx_map[layer_name], + block_ids=block_ids, + num_layers=num_layers, + slot_size=local_slot_size, + block_index=block_index, + ) + return StagedStoreBatch( + buffer=staging, + ready_event=_record_cuda_event(staging), + spans=spans, + lease=staging_lease, + ) + + def _submit_finished_save(self, save: _DeferredFinishedSave) -> Any | None: + """Submit one request's deferred KV store after request completion. + + Args: + save: Deferred store plan built during ``wait_for_save``. + + Returns: + Future that completes after all store batches are committed, or + ``None`` when no store batch could be staged. + + Async/thread-safety: + Called by vLLM's worker thread from ``get_finished`` while vLLM is + still holding the finished request's KV blocks. + """ + batch_futures = [] + reqs_to_store = { + req_id: replace( + spec, + file_offset=_rank_lane_offset( + spec.start_slot, + _local_slot_bytes(self), + int(getattr(self, "_rank_stride_bytes", 0)), + int(getattr(self, "_tp_rank", 0)), + ), + ) + for req_id, spec in save.reqs_to_store.items() + } + batches = _build_staging_store_batches( + reqs_to_store, + _local_slot_bytes(self), + max_batch_bytes=( + self._store_pipeline.staging_bytes or DEFAULT_STORE_STAGING_BYTES + ), + ) + for block_ids, spans in batches: + staged = self._stage_store_batch(block_ids, spans) + if staged is None: + continue + future = self._submit_store_coroutine( + self._write_cuda_buffer( + buffer=staged.buffer, + ready_event=staged.ready_event, + spans=staged.spans, + ) + ) + self._track_save_future(future, staged.buffer.nbytes, staged.lease) + batch_futures.append(future) + if not batch_futures: + return None + commit_future = self._submit_store_coroutine( + self._commit_after_store_futures(batch_futures, sorted(save.commit_keys)), + ) + self._track_save_future(commit_future, 0, None) + return commit_future + + async def _commit_after_store_futures( + self, + batch_futures: list[Any], + commit_keys: list[str], + ) -> None: + """Commit chunks after all staged transfer batches finish. + + Args: + batch_futures: Futures for each staged store batch. + commit_keys: Chunk keys to publish after stores complete. + + Async/thread-safety: + Runs on the connector background event loop and does not read vLLM + KV cache tensors. + """ + stored_keys: list[str] = [] + for future in batch_futures: + stored_keys.extend(await asyncio.wrap_future(future)) + await self._commit_stored_keys(stored_keys, commit_keys) + + async def _write_cuda_buffer( + self, + buffer: torch.Tensor, + ready_event: torch.cuda.Event | None, + spans: list[StoreWriteSpan], + ) -> list[str]: + """Write selected spans from one contiguous CUDA buffer. + + Args: + buffer: CUDA tensor exported over CUDA IPC. + ready_event: Producer-stream event for ``buffer``. + spans: Source/destination write spans. + + Returns: + Chunk keys accepted by the server for this buffer. + """ + if buffer.device.type == "cuda": + torch.cuda.set_device(buffer.device) + if ready_event is not None: + ready_event.synchronize() + else: + _synchronize_cuda_tensor(buffer) + cp_buffer = cupy.asarray(buffer) + cuda_ipc_handle = export_cuda_ipc_handle(cp_buffer) + device_id = cuda_array_device_id(cp_buffer) + device_ptr = cuda_array_pointer(cp_buffer) + ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) + stored_keys = await self._transfer_store_cuda( + cuda_ipc_handle=cuda_ipc_handle, + nbytes=buffer.nbytes, + device_id=device_id, + device_ptr=device_ptr, + allocation_base_ptr=ipc_base_ptr, + allocation_offset=ipc_offset, + producer_pid=os.getpid(), + spans=[ + { + "source_offset": span.source_offset, + "nbytes": span.nbytes, + "file_offset": span.file_offset, + "chunk_key": span.chunk_key, + "start_slot": span.start_slot, + "num_slots": span.num_slots, + } + for span in spans + ], + ) + return stored_keys + + async def _commit_stored_keys( + self, + stored_keys: list[str], + commit_keys: list[str], + ) -> None: + """Commit requested chunks whose store spans were accepted.""" + requested = set(commit_keys) + candidate_keys = [key for key in stored_keys if key in requested] + keys_to_commit = list(dict.fromkeys(candidate_keys)) + await self._store_pipeline.client.commit_chunks( + keys_to_commit, + tp_rank=int(getattr(self, "_tp_rank", 0)), + tp_size=int(getattr(self, "_tp_size", 1)), + ) + + async def _transfer_load_cuda(self, **kwargs: Any) -> dict[str, Any]: + """Load through the dedicated worker load IPC client. + + Args: + **kwargs: forwarded CUDA transfer payload fields. + + Returns: + Server load response with timing counters. + + Async/thread-safety: + Runs on the worker load event loop. A dedicated client keeps + cache-hit loads from queueing behind store RPCs. + """ + 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. + + 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. + """ + 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. + + Args: + **kwargs: forwarded CUDA transfer payload fields. + + Returns: + Chunk keys accepted by the server. + + Async/thread-safety: + Runs on the worker store event loop and serializes only with other + store/commit traffic. + """ + return await self._store_pipeline.client.transfer_store_cuda(**kwargs) diff --git a/docs/design/architecture.md b/docs/design/architecture.md index fce7760..81408a1 100644 --- a/docs/design/architecture.md +++ b/docs/design/architecture.md @@ -63,18 +63,27 @@ graph TB DC["DaserConnector
KVConnectorBase_V1"] SCHED["scheduler.py
SCHEDULER role"] WORKER["worker.py
WORKER role"] + LIFE["RequestLifecycle
pending state + IPC orchestration"] + RUNTIME["WorkerRuntime
KV/meta/completion ownership"] + LOAD["LoadPipeline
queue + load loop"] + STORE["StorePipeline
deferred save + store loop"] VAPI --> DC DC --> SCHED DC --> WORKER + SCHED --> LIFE + WORKER --> RUNTIME + RUNTIME --> LOAD + RUNTIME --> STORE end NVMe[("NVMe
daser.store / daser.index")] User -- "HTTP" --> HTTP HTTP -- "prefill / completion HTTP" --> VAPI - SCHED -- "lookup / alloc / runtime config" --> IPC - WORKER -- "CUDA IPC handle + transfer ops" --> IPC + LIFE -- "lookup / alloc / runtime config" --> IPC + LOAD -- "CUDA IPC handle + load ops" --> IPC + STORE -- "CUDA IPC handle + store / commit ops" --> IPC GDS -- "GDS IO" --> NVMe L1 -- "L2 daser.store IO" --> NVMe CM -- "save / load metadata" --> NVMe @@ -233,11 +242,22 @@ batch，导出 CUDA IPC handle，请求 server 读回 spans，再按层批量拷 KV cache。load 和 store 使用同一套 worker-side staging 抽象。 `wait_for_layer_load` 是 no-op，以兼容 vLLM FULL CUDA graph 模式。 +`WorkerConnectorMixin` 只适配 vLLM hooks；`WorkerRuntime` 集中 KV layout、step +metadata、completion 和 shutdown，内部的 load/store pipeline 各自拥有 queue、 +event loop、IPC client、staging lease 和 future。固定 staging pool 由两个 pipeline +共享的 worker memory module 管理；slot-major tensor layout 和 RoPE restore 留在 +staging module。 + +Scheduler 侧同样把 hook adapter 与状态 implementation 分开：request lifecycle +集中持有 pending load/store/alloc/async-save 状态并编排同步 IPC，chunk/prefix reuse +strategy 只生成 store intent，不反向修改 connector 私有状态。 + ### 后台 asyncio IO loop -vLLM worker 线程不直接运行可重入 event loop。WORKER role 在初始化时创建 -`daser-io` 后台线程，所有 transfer IPC 和 async IPC commit 都通过 -`run_coroutine_threadsafe` 提交。 +vLLM worker 线程不直接运行可重入 event loop。Load/store pipeline 分别拥有 +`daser-load-io` 和 `daser-store-io` 后台线程，transfer IPC 和 async IPC commit +通过 `run_coroutine_threadsafe` 提交。独立 loop 避免 cache-hit load 排在后台 +store coroutine 后面。 ### Transfer backend 启动后不可切换 diff --git a/docs/design/components.md b/docs/design/components.md index 00a9824..44edce4 100644 --- a/docs/design/components.md +++ b/docs/design/components.md @@ -5,9 +5,13 @@ | 组件 | 进程 | 职责 | |------|------|------| | `DaserConnector` | vLLM | vLLM `KVConnectorBase_V1` 入口；保留在 `daser/connector/daser_connector.py` 供 `kv_connector_module_path` 加载 | -| `SchedulerConnectorMixin` | vLLM scheduler | `daser/connector/scheduler.py`；负责 lookup、pending load/store 跟踪、slot 分配和 connector metadata 构造 | -| `WorkerConnectorMixin` | vLLM worker | `daser/connector/worker.py`；负责 KV cache 注册、CUDA IPC handle 导出、后台 IPC loop | -| `CudaStagingPool` | vLLM worker | `daser/connector/staging.py`；负责 GDS 和 iouring 共享的 bounded slot-major GPU staging 复用 | +| `SchedulerConnectorMixin` | vLLM scheduler | `daser/connector/scheduler.py`；适配 vLLM scheduler hooks，不拥有请求 lifecycle 状态 | +| `RequestLifecycle` | vLLM scheduler | 集中 lookup、pending load/store/alloc/async-save、slot 分配、preemption、completion 和 connector metadata 构造 | +| `WorkerConnectorMixin` | vLLM worker | `daser/connector/worker.py`；适配 vLLM worker hooks 和启动期 kernel warmup，不拥有 pipeline 状态 | +| `WorkerRuntime` | vLLM worker | 集中 KV layout、step metadata、load/store completion 和 shutdown，组合两个独立 pipeline | +| `LoadPipeline` / `StorePipeline` | vLLM worker | 分别拥有 load/store event loop、IPC client、staging lease、future、backpressure 和 transfer plan | +| worker memory module | vLLM worker | 负责 load/store 共用的 bounded CUDA staging pool、indexed lease 和 capacity derivation | +| staging tensor module | vLLM worker | 负责 slot-major KV copy、CUDA producer synchronization 和 RoPE restore，不拥有 request/IPC 状态 | | `IPCClientSync` | vLLM scheduler | 阻塞式 Unix socket 客户端，用于 `get_runtime_config`、`lookup`、`alloc_chunk` | | `IPCClientAsync` | vLLM worker | asyncio Unix socket 客户端，用于 `transfer_store`、`transfer_load`、`commit_chunks` | | `TransferLayer` | DaseR | `daser/transfer/base.py`；server-owned KV 数据传输抽象，由 `IPCServer` 按 runtime config 初始化 | diff --git a/docs/design/flows.md b/docs/design/flows.md index aaf047d..343a046 100644 --- a/docs/design/flows.md +++ b/docs/design/flows.md @@ -71,28 +71,30 @@ chunk 已在上传时缓存，task suffix 通常是一次性的，因此推理 ```mermaid sequenceDiagram - participant S as vLLM Scheduler + participant S as Scheduler hook + participant SL as RequestLifecycle participant IPC as IPC server participant C as ServerCore participant CM as ChunkManager - S->>S: get_num_new_matched_tokens(request) - S->>S: full_aligned = floor(len(tokens) / block_tokens) * block_tokens - S->>S: store_key = xxh3_128(tokens[:full_aligned]) - S->>IPC: match_and_alloc(prefix, "", model_id) + S->>SL: get_num_new_matched_tokens(request) + SL->>SL: full_aligned = floor(len(tokens) / block_tokens) * block_tokens + SL->>SL: reuse strategy builds store intent + SL->>IPC: match_and_alloc(prefix, "", model_id) IPC->>C: match_and_alloc(...) C-->>IPC: chunks=[] / alloc=null - IPC-->>S: miss - S->>S: track PendingStore(chunk_key, token_count) + IPC-->>SL: miss + SL->>SL: track PendingStore(chunk_key, token_count) - S->>S: update_state_after_alloc(block_ids) - S->>IPC: alloc_chunk(chunk_key, token_count, model_id) + S->>SL: update_state_after_alloc(block_ids) + SL->>IPC: alloc_chunk(chunk_key, token_count, model_id) IPC->>C: alloc_chunk(...) C->>CM: allocate slots, evict old chunks if needed CM-->>C: start_slot, num_slots C-->>IPC: file_offset, pos_offset - IPC-->>S: allocation - S->>S: build_connector_meta(reqs_to_store) + IPC-->>SL: allocation + S->>SL: build_connector_meta(scheduler_output) + SL-->>S: reqs_to_store ``` `alloc_chunk` 只预留 metadata，chunk 还不会进入 `RetrievalIndex`。 @@ -101,22 +103,23 @@ sequenceDiagram ```mermaid sequenceDiagram - participant W as vLLM Worker - participant BG as daser-io loop + participant W as Worker hook + participant WR as WorkerRuntime + participant BG as StorePipeline / daser-store-io participant IPC as IPC server participant TL as TransferLayer - W->>W: bind_connector_metadata(reqs_to_store) + W->>WR: bind_connector_metadata(reqs_to_store) loop each attention layer - W->>W: save_kv_layer(layer_name, kv_layer) + W->>WR: save_kv_layer(layer_name, kv_layer) end - W->>W: wait_for_save() - W->>W: split reqs into bounded staging batches + W->>WR: wait_for_save() + WR->>BG: defer/submit reqs_to_store + BG->>BG: split reqs into bounded staging batches loop each staging batch - W->>W: lease bounded GPU staging view - W->>W: copy selected block KV into slot-major staging - W->>W: record producer CUDA event - W->>BG: run_coroutine_threadsafe(_write_cuda_buffer) + BG->>BG: lease bounded GPU staging view + BG->>BG: copy selected block KV into slot-major staging + BG->>BG: record producer CUDA event BG->>BG: wait producer event BG->>IPC: transfer_store(cuda_ipc_handle, spans) IPC->>TL: store_bytes_grouped(staging slices, file_offset) @@ -127,7 +130,7 @@ sequenceDiagram ``` `wait_for_save` 在当前 worker step 内完成 KV -> staging snapshot，然后把 -server transfer 交给后台 `daser-io` loop。未完成 batch 的 staging lease 由 +server transfer 交给后台 `daser-store-io` loop。未完成 batch 的 staging lease 由 future 持有，完成后归还 `CudaStagingPool`；`shutdown` 会阻塞等待所有 pending store future 完成。 @@ -154,22 +157,24 @@ commit 完成后 chunk 才能被 lookup 命中。 ```mermaid sequenceDiagram - participant S as vLLM Scheduler + participant S as Scheduler hook + participant SL as RequestLifecycle participant IPC as IPC server participant C as ServerCore participant RI as RetrievalIndex - S->>S: get_num_new_matched_tokens(request) - S->>IPC: match_and_alloc(prefix, "", model_id) + S->>SL: get_num_new_matched_tokens(request) + SL->>IPC: match_and_alloc(prefix, "", model_id) IPC->>C: lookup(prefix, model_id) C->>RI: lookup(tokens, model_id) RI-->>C: matched chunks C-->>IPC: chunks - IPC-->>S: chunks - S->>S: compute extra_tokens and pending_loads - S->>S: update_state_after_alloc(block_ids) - S->>S: map chunk target ranges to vLLM block ids - S->>S: build_connector_meta(reqs_to_load) + IPC-->>SL: chunks + SL->>SL: compute extra_tokens and pending_loads + S->>SL: update_state_after_alloc(block_ids) + SL->>SL: map chunk target ranges to vLLM block ids + S->>SL: build_connector_meta(scheduler_output) + SL-->>S: reqs_to_load ``` `PrefixHashIndex` 返回 rolling-prefix 的连续 slot 命中； @@ -180,17 +185,18 @@ vLLM 的 external tokens 是连续可用的前缀范围。 ```mermaid sequenceDiagram - participant W as vLLM Worker - participant BG as daser-io loop + participant W as Worker hook + participant WR as WorkerRuntime + participant BG as LoadPipeline / daser-load-io participant IPC as IPC server participant TL as TransferLayer - W->>W: start_load_kv(forward_context) - W->>W: split spans into bounded staging batches + W->>WR: start_load_kv(forward_context) + WR->>BG: submit reqs_to_load + BG->>BG: split spans into bounded staging batches loop each load batch - W->>W: lease GPU uint8 staging view - W->>W: export CUDA IPC handle for staging - W->>BG: transfer_load(cuda_ipc_handle, spans).result(timeout=120s) + BG->>BG: lease GPU uint8 staging view + BG->>BG: export CUDA IPC handle for staging BG->>IPC: transfer_load(cuda_ipc_handle, spans) IPC->>TL: load_bytes_grouped(staging slices, file_offset) TL-->>IPC: bytes read diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index ea18b19..2b46b90 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -19,56 +19,58 @@ rolling_prefix_key, rolling_prefix_keys, ) +from daser.connector.load_pipeline import ( + build_load_copy_runs as _build_load_copy_runs, +) +from daser.connector.load_pipeline import ( + build_load_read_batches as _build_load_read_batches, +) +from daser.connector.load_pipeline import ( + build_load_read_plan as _build_load_read_plan, +) from daser.connector.metadata import ( DaserConnectorMeta, ReqLoadSpec, ReqStoreSpec, StoreWriteSpan, ) +from daser.connector.request_lifecycle import RequestLifecycle from daser.connector.reuse import PrefixReuseStrategy from daser.connector.scheduler import ( - SchedulerConnectorMixin, _block_ids_for_chunk, _contiguous_prefix_tokens, _trim_chunk_to_external_window, ) from daser.connector.staging import ( - DEFAULT_PENDING_STORE_STAGING_BYTES, DEFAULT_ROPE_DELTA_SCALE, - DEFAULT_STORE_STAGING_BYTES, - MIN_STORE_STAGING_BYTES, - StoreCudaStagingPool, ) from daser.connector.staging import ( apply_rope_delta_to_key_block as _apply_rope_delta_to_key_block, ) from daser.connector.staging import ( - build_load_copy_runs as _build_load_copy_runs, + copy_staging_to_kv_cache as _copy_staging_to_kv_cache, ) from daser.connector.staging import ( - build_load_read_batches as _build_load_read_batches, + record_cuda_event as _record_cuda_event, ) from daser.connector.staging import ( - build_load_read_plan as _build_load_read_plan, + synchronize_cuda_tensor as _synchronize_cuda_tensor, ) -from daser.connector.staging import ( +from daser.connector.store_pipeline import ( build_staging_store_batches as _build_staging_store_batches, ) -from daser.connector.staging import ( - copy_staging_to_kv_cache as _copy_staging_to_kv_cache, -) -from daser.connector.staging import ( - derive_store_staging_limits as _derive_store_staging_limits, +from daser.connector.worker_memory import ( + DEFAULT_PENDING_STORE_STAGING_BYTES, + DEFAULT_STORE_STAGING_BYTES, + MIN_STORE_STAGING_BYTES, + FixedCudaStagingPool, ) -from daser.connector.staging import ( - record_cuda_event as _record_cuda_event, +from daser.connector.worker_memory import ( + derive_staging_limits as _derive_staging_limits, ) -from daser.connector.staging import ( - synchronize_cuda_tensor as _synchronize_cuda_tensor, -) -from daser.connector.worker import ( +from daser.connector.worker_runtime import ( LoadRequestDispatcher, - WorkerConnectorMixin, + WorkerRuntime, _DeferredFinishedSave, _InflightRequestLoad, _load_staging_pool_depth, @@ -118,7 +120,75 @@ def runtime_state(self): ) -class _WorkerProbe(DaserConnector): +class _LoadPipelineProbe: + """In-memory load pipeline implementing the worker-facing interface.""" + + def __init__( + self, + client: object | None = None, + clients: list[object] | None = None, + ) -> None: + self.pending: dict[str, object] = {} + self.invalid_block_ids: set[int] = set() + self.staging_pool = None + self.staging_registered = False + self._clients = clients if clients is not None else ([client] if client else []) + self.queue: asyncio.Queue[object] = asyncio.Queue() + self.shutdown_called = False + self.submit_calls = 0 + + def set_staging_pool(self, pool: object) -> None: + self.staging_pool = pool + + def clients(self) -> tuple[object, ...]: + return tuple(self._clients) + + def client(self, buffer_index: int | None = None) -> object: + index = 0 if buffer_index is None else buffer_index + return self._clients[index % len(self._clients)] + + def ensure_queue(self) -> asyncio.Queue[object]: + return self.queue + + def enqueue(self, item: object) -> None: + self.queue.put_nowait(item) + + def ensure_dispatcher(self, coro: object) -> None: + coro.close() + + def submit(self, coro: object) -> object: + self.submit_calls += 1 + coro.close() + return SimpleNamespace(result=lambda timeout=None: None) + + def shutdown(self) -> None: + self.shutdown_called = True + + +class _StorePipelineProbe: + """In-memory store pipeline implementing the worker-facing interface.""" + + def __init__(self, client: object | None = None) -> None: + self.client = client + self.save_futures: list[object] = [] + self.pending_staging_bytes = 0 + self.staging_bytes = 0 + self.pending_staging_limit_bytes = 0 + self.staging_pool = None + self.pending_finished_saves: dict[str, object] = {} + self.shutdown_called = False + self.submit_calls = 0 + + def submit(self, coro: object) -> object: + self.submit_calls += 1 + coro.close() + return SimpleNamespace(result=lambda timeout=None: None) + + def shutdown(self) -> None: + self.shutdown_called = True + + +class _WorkerProbe(WorkerRuntime): """Worker-side probe with minimal state for transfer readiness tests.""" def __init__(self, store_path: str) -> None: @@ -141,11 +211,8 @@ def __init__(self, store_path: str) -> None: self._layer_names = [] self._transfer_mode = "gds" self._skip_l2 = False - self._pending_loads = {} - self._invalid_load_block_ids = set() - self._pending_finished_saves = {} - self._save_futures = [] - self._pending_save_staging_bytes = 0 + self._load_pipeline = _LoadPipelineProbe(self) + self._store_pipeline = _StorePipelineProbe(self) def _refresh_runtime_config(self) -> None: return @@ -154,8 +221,8 @@ def _ensure_load_request_dispatcher(self, sample_tensor: torch.Tensor) -> None: """Drain queued requests synchronously for worker probe tests.""" dispatcher = LoadRequestDispatcher(max_inflight=8, staging_depth=1) queued = [] - while not self._load_request_queue.empty(): - item = self._load_request_queue.get() + while not self._load_pipeline.queue.empty(): + item = self._load_pipeline.queue.get_nowait() if item is None: return queued.append(item) @@ -171,7 +238,7 @@ def transfer_ready(self): return self._transfer_ready -class _SchedulerProbe(DaserConnector): +class _SchedulerProbe(RequestLifecycle): """Scheduler-side probe that can emulate deferred runtime config.""" def __init__(self, ipc_client) -> None: @@ -196,7 +263,7 @@ def mark_runtime_ready(self, model_id: str) -> None: self._model_id = model_id -class _AllocatingSchedulerProbe(SchedulerConnectorMixin): +class _AllocatingSchedulerProbe(RequestLifecycle): """Minimal scheduler probe that records allocation RPCs.""" def __init__(self) -> None: @@ -342,14 +409,13 @@ def alloc_chunk( return self._owner.alloc_chunk(chunk_key, token_count, model_id) -class _CommitProbe(WorkerConnectorMixin): +class _CommitProbe(WorkerRuntime): """Minimal worker probe exposing commit filtering behavior.""" def __init__(self) -> None: self.committed: list[list[str]] = [] self.commit_metadata: list[tuple[int, int]] = [] - self._ipc_async = self - self._ipc_store_async = self + self._store_pipeline = _StorePipelineProbe(self) async def commit_chunks( self, chunk_keys: list[str], tp_rank: int = 0, tp_size: int = 1 @@ -365,15 +431,14 @@ async def commit_stored_keys( await self._commit_stored_keys(stored_keys, commit_keys) -class _QueueProbe(WorkerConnectorMixin): +class _QueueProbe(WorkerRuntime): """Minimal worker probe exposing independent load/store IPC clients.""" def __init__(self) -> None: self.load_calls: list[str] = [] self.store_calls: list[str] = [] - self._ipc_load_async = self - self._ipc_store_async = self - self._ipc_async = self + self._load_pipeline = _LoadPipelineProbe(self) + self._store_pipeline = _StorePipelineProbe(self) async def transfer_load_cuda(self, **_kwargs) -> dict[str, int]: """Record load client usage.""" @@ -394,18 +459,17 @@ async def store_via_public_helper(self) -> None: await self._transfer_store_cuda() -class _FinishedSaveProbe(WorkerConnectorMixin): +class _FinishedSaveProbe(WorkerRuntime): """Worker probe for finished-request save scheduling.""" def __init__(self) -> None: self._meta = None self._pending_commits = set() - self._pending_finished_saves = {} - self._save_futures = [] - self._pending_save_staging_bytes = 0 + self._load_pipeline = _LoadPipelineProbe() + self._store_pipeline = _StorePipelineProbe() self._slot_size = 32 - self._store_staging_bytes = 128 - self._pending_store_staging_limit_bytes = 128 + self._store_pipeline.staging_bytes = 128 + self._store_pipeline.pending_staging_limit_bytes = 128 self._layer_names = ["layer.0"] self._kv_caches = {"layer.0": torch.empty(1)} self.staged_batches: list[tuple[list[int], list[StoreWriteSpan]]] = [] @@ -418,11 +482,11 @@ def _clear_save_state(self) -> None: def _reap_save_futures(self, block: bool) -> None: if block: - for future, _bytes, lease in self._save_futures: + for future, _bytes, lease in self._store_pipeline.save_futures: future.result(timeout=120.0) if lease is not None: lease.release() - self._save_futures = [] + self._store_pipeline.save_futures = [] def _stage_store_batch( self, @@ -470,7 +534,7 @@ def _track_save_future( staging_lease, ) -> None: self.tracked.append((staging_bytes, staging_lease)) - self._save_futures.append((future, staging_bytes, staging_lease)) + self._store_pipeline.save_futures.append((future, staging_bytes, staging_lease)) async def _commit_after_store_futures(self, batch_futures, commit_keys): self.committed_after.append((len(batch_futures), list(commit_keys))) @@ -491,14 +555,14 @@ def seed_finished_save( commit_keys: list[str], ) -> None: """Seed a deferred finished save for worker completion tests.""" - self._pending_finished_saves[req_id] = _DeferredFinishedSave( + self._store_pipeline.pending_finished_saves[req_id] = _DeferredFinishedSave( commit_keys=set(commit_keys), reqs_to_store=reqs_to_store, ) def pending_finished_save_ids(self) -> set[str]: """Return request IDs with deferred save work.""" - return set(self._pending_finished_saves) + return set(self._store_pipeline.pending_finished_saves) def pending_commit_keys(self) -> set[str]: """Return pending worker commit keys.""" @@ -537,15 +601,13 @@ def mark_done(self) -> None: self._done = True -class _AsyncLoadProbe(WorkerConnectorMixin): +class _AsyncLoadProbe(WorkerRuntime): """Worker probe with pending async load futures.""" def __init__(self) -> None: self._role = KVConnectorRole.WORKER - self._pending_loads = {} - self._pending_finished_saves = {} - self._save_futures = [] - self._invalid_load_block_ids = set() + self._load_pipeline = _LoadPipelineProbe() + self._store_pipeline = _StorePipelineProbe() self.released: list[object] = [] self.load_loop_stopped = False self.store_loop_stopped = False @@ -571,17 +633,14 @@ def join(self, timeout: float | None = None) -> None: del timeout self._callback() - self._load_loop = _Loop(lambda: setattr(self, "load_loop_stopped", True)) - self._store_loop = _Loop(lambda: setattr(self, "store_loop_stopped", True)) - self._load_thread = _Thread(lambda: setattr(self, "load_thread_joined", True)) - self._store_thread = _Thread(lambda: setattr(self, "store_thread_joined", True)) + del _Loop, _Thread def _reap_save_futures(self, block: bool) -> None: del block def get_finished(self, finished_req_ids: set[str]): """Expose worker completion polling without shutdown test doubles.""" - return WorkerConnectorMixin.get_finished(self, finished_req_ids) + return WorkerRuntime.get_finished(self, finished_req_ids) def _submit_load_coroutine(self, coro): coro.close() @@ -610,24 +669,24 @@ def seed_load( lease: object | None = None, ) -> None: """Seed one pending async load.""" - self._pending_loads[req_id] = _PendingLoad( + self._load_pipeline.pending[req_id] = _PendingLoad( future=future, block_ids=block_ids, lease=lease, ) -class _LoopProbe(WorkerConnectorMixin): +class _LoopProbe(WorkerRuntime): """Minimal worker probe exposing background loop selection.""" def __init__(self) -> None: - self._load_loop = object() - self._store_loop = object() + self._load_pipeline = _LoadPipelineProbe() + self._store_pipeline = _StorePipelineProbe() @property def loop_pair(self) -> tuple[object, object]: """Return the load and store loops configured on this probe.""" - return self._load_loop, self._store_loop + return self._load_pipeline, self._store_pipeline def submit_load(self) -> None: """Submit a load coroutine through the worker helper.""" @@ -752,6 +811,10 @@ class DummyVLLMConfig: ) assert isinstance(connector._cache_reuse_strategy, PrefixReuseStrategy) # noqa: SLF001 + assert isinstance( # noqa: SLF001 + connector._request_lifecycle._reuse_strategy(), # noqa: SLF001 + PrefixReuseStrategy, + ) def test_start_load_kv_initializes_gds_after_server_creates_store( @@ -827,24 +890,32 @@ def fake_submit_load_coroutine(coro): monkeypatch.setattr(connector, "_acquire_staging", lambda *args: _Lease()) monkeypatch.setattr(connector, "_submit_load_coroutine", fake_submit_load_coroutine) - monkeypatch.setattr("daser.connector.worker.cupy.asarray", lambda tensor: tensor) - monkeypatch.setattr("daser.connector.worker.export_cuda_ipc_handle", lambda _: b"0") - monkeypatch.setattr("daser.connector.worker.cuda_array_device_id", lambda _: 0) - monkeypatch.setattr("daser.connector.worker.cuda_array_pointer", lambda _: 1) monkeypatch.setattr( - "daser.connector.worker._cuda_allocation_base_and_offset", + "daser.connector.worker_runtime.cupy.asarray", lambda tensor: tensor + ) + monkeypatch.setattr( + "daser.connector.worker_runtime.export_cuda_ipc_handle", lambda _: b"0" + ) + monkeypatch.setattr( + "daser.connector.worker_runtime.cuda_array_device_id", lambda _: 0 + ) + monkeypatch.setattr( + "daser.connector.worker_runtime.cuda_array_pointer", lambda _: 1 + ) + monkeypatch.setattr( + "daser.connector.worker_runtime._cuda_allocation_base_and_offset", lambda _: (1, 0), ) monkeypatch.setattr( - "daser.connector.worker._copy_staging_to_kv_cache", + "daser.connector.worker_runtime._copy_staging_to_kv_cache", lambda **kwargs: 1, ) monkeypatch.setattr( - "daser.connector.worker._synchronize_cuda_tensor", + "daser.connector.worker_runtime._synchronize_cuda_tensor", lambda tensor: None, ) - with caplog.at_level("INFO", logger="daser.connector.worker"): + with caplog.at_level("INFO", logger="daser.connector.worker_runtime"): connector.start_load_kv(forward_context=object()) assert "start_load_kv timing" not in caplog.text @@ -862,10 +933,12 @@ async def transfer_load_registered_cuda(self, **kwargs): self.calls.append(kwargs) return {"client": self.name} - class Probe(WorkerConnectorMixin): + load0 = _Client("load0") + load1 = _Client("load1") + + class Probe(WorkerRuntime): def __init__(self) -> None: - self._ipc_load_async = _Client("fallback") - self._ipc_load_async_pool = [_Client("load0"), _Client("load1")] + self._load_pipeline = _LoadPipelineProbe(clients=[load0, load1]) probe = Probe() @@ -882,7 +955,7 @@ def __init__(self) -> None: assert asyncio.run(coro0) == {"client": "load0"} assert asyncio.run(coro1) == {"client": "load1"} - assert probe._ipc_load_async.calls == [] # noqa: SLF001 + assert load0.calls and load1.calls def test_get_finished_releases_each_request_load_independently() -> None: @@ -899,7 +972,7 @@ def test_get_finished_releases_each_request_load_independently() -> 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 + assert set(connector._load_pipeline.pending) == {"req-b"} # noqa: SLF001 def test_start_load_kv_releases_waiting_request_when_load_cannot_start() -> None: @@ -1114,7 +1187,7 @@ def test_get_finished_does_not_block_on_pending_async_load() -> None: assert finished_sending is None assert finished_recving is None assert future.result_calls == 0 - assert "req" in connector._pending_loads # noqa: SLF001 + assert "req" in connector._load_pipeline.pending # noqa: SLF001 def test_get_finished_reports_completed_async_load() -> None: @@ -1128,8 +1201,8 @@ def test_get_finished_reports_completed_async_load() -> None: assert finished_sending is None assert finished_recving == {"req"} assert future.result_calls == 1 - assert connector._pending_loads == {} # noqa: SLF001 - assert connector._invalid_load_block_ids == set() # noqa: SLF001 + assert connector._load_pipeline.pending == {} # noqa: SLF001 + assert connector._load_pipeline.invalid_block_ids == set() # noqa: SLF001 def test_load_request_dispatcher_limits_active_requests_by_staging_depth() -> None: @@ -1242,8 +1315,8 @@ def test_get_finished_releases_request_after_async_load_failure() -> None: assert finished_sending is None assert finished_recving == {"req"} - assert connector._pending_loads == {} # noqa: SLF001 - assert connector._invalid_load_block_ids == {4, 5} # noqa: SLF001 + assert connector._load_pipeline.pending == {} # noqa: SLF001 + assert connector._load_pipeline.invalid_block_ids == {4, 5} # noqa: SLF001 def test_shutdown_collects_failed_async_load_without_raising() -> None: @@ -1254,12 +1327,10 @@ def test_shutdown_collects_failed_async_load_without_raising() -> None: connector.shutdown() - assert connector._pending_loads == {} # noqa: SLF001 - assert connector._invalid_load_block_ids == {4, 5} # noqa: SLF001 - assert connector.load_loop_stopped - assert connector.store_loop_stopped - assert connector.load_thread_joined - assert connector.store_thread_joined + assert connector._load_pipeline.pending == {} # noqa: SLF001 + assert connector._load_pipeline.invalid_block_ids == {4, 5} # noqa: SLF001 + assert connector._load_pipeline.shutdown_called # noqa: SLF001 + assert connector._store_pipeline.shutdown_called # noqa: SLF001 def test_worker_transfer_ready_allows_skip_l2_without_store_path() -> None: @@ -1272,6 +1343,83 @@ def test_worker_transfer_ready_allows_skip_l2_without_store_path() -> None: assert connector.transfer_ready is True +def test_worker_runtime_refreshes_l1_only_transfer_config(monkeypatch) -> None: + """Deferred worker config refresh must propagate L1-only transfer settings.""" + + class DummyIPCClient: + def __init__(self, socket_path: str) -> None: + self.socket_path = socket_path + + def get_runtime_config(self) -> dict[str, object]: + return { + "store_path": "", + "slot_size": 1024, + "tensor_parallel_size": 2, + "rank_stride_bytes": 512, + "transfer_mode": "iouring", + "skip_l2": True, + } + + def close(self) -> None: + return + + monkeypatch.setattr( + "daser.connector.worker_runtime.IPCClientSync", + DummyIPCClient, + ) + connector = WorkerRuntime.__new__(WorkerRuntime) + connector._socket_path = "/unused/daser.sock" # noqa: SLF001 + connector._store_path = "" # noqa: SLF001 + connector._slot_size = 0 # noqa: SLF001 + connector._server_tp_size = 1 # noqa: SLF001 + connector._rank_stride_bytes = 0 # noqa: SLF001 + connector._transfer_mode = "gds" # noqa: SLF001 + connector._skip_l2 = False # noqa: SLF001 + + connector._refresh_runtime_config() # noqa: SLF001 + + assert connector._slot_size == 1024 # noqa: SLF001 + assert connector._server_tp_size == 2 # noqa: SLF001 + assert connector._rank_stride_bytes == 512 # noqa: SLF001 + assert connector._transfer_mode == "iouring" # noqa: SLF001 + assert connector._skip_l2 is True # noqa: SLF001 + + +def test_request_lifecycle_rebuilds_prefix_keys_after_block_size_refresh() -> None: + """Deferred geometry refresh must rebuild same-mode rolling-prefix keys.""" + + class DummyIPCClient: + def get_runtime_config(self) -> dict[str, object]: + return { + "slot_size": 1024, + "block_tokens": 128, + "model_id": "served-model", + "cache_reuse_mode": "prefix", + } + + lifecycle = RequestLifecycle( + ipc_client=DummyIPCClient(), + block_tokens=16, + slot_size=0, + model_id="default", + cache_reuse_mode="prefix", + runtime_config_ready=False, + ) + lifecycle._refresh_runtime_config() # noqa: SLF001 + tokens = list(range(256)) + pending = lifecycle._reuse_strategy().prepare_store(tokens, 256) # noqa: SLF001 + + assert pending is not None + pending.block_ids = [7, 8] + plan = lifecycle._reuse_strategy().plan_store( # noqa: SLF001 + "req", + pending, + tokens, + set(), + ) + assert [intent.chunk_key for intent in plan.intents] == rolling_keys(tokens, 128) + + def test_runtime_config_ready_allows_skip_l2_without_store_path(monkeypatch): """Scheduler runtime config should be ready when skip_l2 omits store_path.""" @@ -1542,7 +1690,7 @@ def test_trim_chunk_to_external_window_skips_local_prefix_slots(): def test_update_state_after_alloc_single_hit_uses_external_window(): """Single-prefix hit maps the external suffix onto absolute request blocks.""" - class MockConnector(SchedulerConnectorMixin): + class MockConnector(RequestLifecycle): def __init__(self) -> None: self._block_tokens = BLOCK_TOKENS self._slot_size = 32 @@ -1575,7 +1723,7 @@ class MockBlocks: connector = MockConnector() - DaserConnector.update_state_after_alloc( + RequestLifecycle.update_state_after_alloc( connector, MockRequest(), MockBlocks(), @@ -1623,7 +1771,7 @@ def lookup( } ] - class MockConnector(SchedulerConnectorMixin): + class MockConnector(RequestLifecycle): def __init__(self) -> None: self._runtime_config_ready = True self._block_tokens = BLOCK_TOKENS @@ -1664,7 +1812,7 @@ class MockRequest: def test_update_state_after_alloc_multi_hit_trims_each_chunk_to_external_window(): """Multi-chunk hits map onto absolute request block positions.""" - class MockConnector(SchedulerConnectorMixin): + class MockConnector(RequestLifecycle): def __init__(self) -> None: self._block_tokens = BLOCK_TOKENS self._slot_size = 32 @@ -1708,7 +1856,7 @@ class MockBlocks: connector = MockConnector() - DaserConnector.update_state_after_alloc( + RequestLifecycle.update_state_after_alloc( connector, MockRequest(), MockBlocks(), @@ -2175,13 +2323,13 @@ def test_restore_cross_layer_kv_cache_table_tilelang_matches_reference_cuda(): def test_register_kv_caches_warms_dynamic_rope_apply_once(monkeypatch): - from daser.connector import worker + from daser.connector import worker_runtime as worker - class Probe(WorkerConnectorMixin): + class Probe(WorkerRuntime): def __init__(self) -> None: self._slot_size = 0 - self._store_staging_bytes = 0 - self._pending_store_staging_limit_bytes = 0 + self._load_pipeline = _LoadPipelineProbe() + self._store_pipeline = _StorePipelineProbe() self._rope_rotary_dim = 8 self._rope_base = 10000.0 self._rope_is_neox_style = True @@ -2194,7 +2342,7 @@ def _ensure_transfer_ready(self) -> bool: calls = [] monkeypatch.setattr( worker, - "_derive_store_staging_limits", + "_derive_staging_limits", lambda device: (4096, 8192), ) monkeypatch.setattr( @@ -2222,7 +2370,7 @@ def _ensure_transfer_ready(self) -> bool: def test_register_cross_layers_kv_cache_preserves_layer_order(monkeypatch): """Worker registration keeps vLLM layer names for slot-major staging.""" - from daser.connector import worker + from daser.connector import worker_runtime as worker class Group: layer_names = ["layer.0", "layer.1"] @@ -2230,12 +2378,12 @@ class Group: class Config: kv_cache_groups = [Group()] - class Probe(WorkerConnectorMixin): + class Probe(WorkerRuntime): def __init__(self) -> None: self._kv_cache_config = Config() self._slot_size = 0 - self._store_staging_bytes = 0 - self._pending_store_staging_limit_bytes = 0 + self._load_pipeline = _LoadPipelineProbe() + self._store_pipeline = _StorePipelineProbe() self._rope_rotary_dim = 8 self._rope_base = 10000.0 self._rope_is_neox_style = True @@ -2244,6 +2392,9 @@ def __init__(self) -> None: def _ensure_transfer_ready(self) -> bool: return True + def _init_server_transfer(self) -> None: + return + @property def registration_state(self): return ( @@ -2255,7 +2406,7 @@ def registration_state(self): monkeypatch.setattr( worker, - "_derive_store_staging_limits", + "_derive_staging_limits", lambda device: (4096, 8192), ) monkeypatch.setattr(worker, "_warm_rope_apply_backends", lambda **kwargs: None) @@ -2272,9 +2423,9 @@ def registration_state(self): def test_stage_store_batch_does_not_warm_dynamic_rope_again(monkeypatch): - from daser.connector import worker + from daser.connector import worker_runtime as worker - class Probe(WorkerConnectorMixin): + class Probe(WorkerRuntime): def __init__(self) -> None: self.kv_cache = torch.zeros((2, 8, 4, 2, 8), dtype=torch.float32) self._layer_names = ["layer.0"] @@ -2282,11 +2433,9 @@ def __init__(self) -> None: self._kv_caches = {"layer.0": self.kv_cache} self.slot_size = self.kv_cache[:, 0].nbytes self._slot_size = self.slot_size - self._store_staging_bytes = 4096 - self._pending_store_staging_limit_bytes = 8192 - self._store_staging_pool = None - self._pending_save_staging_bytes = 0 - self._save_futures = [] + self._store_pipeline = _StorePipelineProbe() + self._store_pipeline.staging_bytes = 4096 + self._store_pipeline.pending_staging_limit_bytes = 8192 self._rope_rotary_dim = 8 self._rope_base = 10000.0 self._rope_is_neox_style = True @@ -2326,17 +2475,18 @@ def result(self, timeout: float) -> None: assert timeout > 0 completed.append(self.name) - class Probe(WorkerConnectorMixin): + class Probe(WorkerRuntime): def __init__(self) -> None: - self._store_staging_pool = StoreCudaStagingPool( + self._store_pipeline = _StorePipelineProbe() + self._store_pipeline.staging_pool = FixedCudaStagingPool( device=torch.device("cpu"), buffer_bytes=16, depth=1, ) - lease = self._store_staging_pool.acquire(8) - self._pending_store_staging_limit_bytes = 64 - self._pending_save_staging_bytes = 8 - self._save_futures = [ + lease = self._store_pipeline.staging_pool.acquire(8) + self._store_pipeline.pending_staging_limit_bytes = 64 + self._store_pipeline.pending_staging_bytes = 8 + self._store_pipeline.save_futures = [ _SaveFuture(_Future("commit"), 0, None), _SaveFuture(_Future("store"), 8, lease), ] @@ -2349,13 +2499,13 @@ def wait_for_staging(self) -> None: probe.wait_for_staging() assert completed == ["commit", "store"] - assert probe._store_staging_pool.available == 1 # noqa: SLF001 + assert probe._store_pipeline.staging_pool.available == 1 # noqa: SLF001 def test_stage_store_batch_records_ready_event_without_synchronizing(monkeypatch): - from daser.connector import worker + from daser.connector import worker_runtime as worker - class Probe(WorkerConnectorMixin): + class Probe(WorkerRuntime): def __init__(self) -> None: self.kv_cache = torch.zeros((2, 8, 4, 2, 8), dtype=torch.float32) self._layer_names = ["layer.0"] @@ -2363,11 +2513,9 @@ def __init__(self) -> None: self._kv_caches = {"layer.0": self.kv_cache} self.slot_size = self.kv_cache[:, 0].nbytes self._slot_size = self.slot_size - self._store_staging_bytes = 4096 - self._pending_store_staging_limit_bytes = 8192 - self._store_staging_pool = None - self._pending_save_staging_bytes = 0 - self._save_futures = [] + self._store_pipeline = _StorePipelineProbe() + self._store_pipeline.staging_bytes = 4096 + self._store_pipeline.pending_staging_limit_bytes = 8192 self._rope_rotary_dim = 8 self._rope_base = 10000.0 self._rope_is_neox_style = True @@ -2397,9 +2545,9 @@ def stage_store_batch(self, block_ids: list[int], spans: list[StoreWriteSpan]): def test_stage_store_batch_keeps_dynamic_rope_warmup_out_of_store_path(monkeypatch): - from daser.connector import worker + from daser.connector import worker_runtime as worker - class Probe(WorkerConnectorMixin): + class Probe(WorkerRuntime): def __init__(self) -> None: self.kv_cache = torch.zeros((2, 8, 4, 2, 8), dtype=torch.float32) self._layer_names = ["layer.0"] @@ -2407,11 +2555,9 @@ def __init__(self) -> None: self._kv_caches = {"layer.0": self.kv_cache} self.slot_size = self.kv_cache[:, 0].nbytes self._slot_size = self.slot_size - self._store_staging_bytes = 4096 - self._pending_store_staging_limit_bytes = 8192 - self._store_staging_pool = None - self._pending_save_staging_bytes = 0 - self._save_futures = [] + self._store_pipeline = _StorePipelineProbe() + self._store_pipeline.staging_bytes = 4096 + self._store_pipeline.pending_staging_limit_bytes = 8192 self._rope_rotary_dim = 8 self._rope_base = 10000.0 self._rope_is_neox_style = True @@ -2436,7 +2582,7 @@ def stage_store_batch(self, block_ids: list[int], spans: list[StoreWriteSpan]): def test_update_state_after_alloc_skips_chunks_beyond_external_prefix(): - class MockConnector(SchedulerConnectorMixin): + class MockConnector(RequestLifecycle): def __init__(self) -> None: self._block_tokens = BLOCK_TOKENS self._slot_size = 32 @@ -2486,7 +2632,7 @@ class MockBlocks: connector = MockConnector() - DaserConnector.update_state_after_alloc( + RequestLifecycle.update_state_after_alloc( connector, MockRequest(), MockBlocks(), num_external_tokens=8 ) @@ -3096,7 +3242,7 @@ def lookup(self, tokens, model_id): } ] - class MockConnector(SchedulerConnectorMixin): + class MockConnector(RequestLifecycle): def __init__(self) -> None: self._runtime_config_ready = True self._block_tokens = BLOCK_TOKENS @@ -3824,29 +3970,20 @@ async def test_worker_load_and_store_use_separate_ipc_clients(): assert connector.store_calls == ["store"] -def test_worker_load_and_store_use_separate_background_loops(monkeypatch): +def test_worker_load_and_store_use_separate_background_loops(): """Foreground loads should not queue behind background store loop work.""" connector = _LoopProbe() - submitted_loops = [] - - class Future: - def result(self, timeout: float | None = None) -> None: - return None - - def run_threadsafe(coro, loop): - submitted_loops.append(loop) - coro.close() - return Future() - - monkeypatch.setattr(asyncio, "run_coroutine_threadsafe", run_threadsafe) connector.submit_load() connector.submit_store() - assert submitted_loops == list(connector.loop_pair) + load_pipeline, store_pipeline = connector.loop_pair + assert load_pipeline.submit_calls == 1 + assert store_pipeline.submit_calls == 1 + assert load_pipeline is not store_pipeline -def test_derive_store_staging_limits_scale_with_vram(monkeypatch): +def test_derive_staging_limits_scale_with_vram(monkeypatch): """GPU staging caps consider device size and currently free VRAM.""" class Props: @@ -3863,7 +4000,7 @@ def __init__(self, total_memory: int) -> None: "mem_get_info", lambda device=None: (12 << 30, 24 << 30), ) - small_batch, small_pending = _derive_store_staging_limits(torch.device("cuda")) + small_batch, small_pending = _derive_staging_limits(torch.device("cuda")) assert small_batch == max( MIN_STORE_STAGING_BYTES, min((24 << 30) // 50, (12 << 30) // 10), @@ -3883,7 +4020,7 @@ def __init__(self, total_memory: int) -> None: "mem_get_info", lambda device=None: (64 << 30, 80 << 30), ) - large_batch, large_pending = _derive_store_staging_limits(torch.device("cuda")) + large_batch, large_pending = _derive_staging_limits(torch.device("cuda")) assert large_batch == DEFAULT_STORE_STAGING_BYTES assert large_pending == DEFAULT_PENDING_STORE_STAGING_BYTES @@ -3892,14 +4029,14 @@ def __init__(self, total_memory: int) -> None: "mem_get_info", lambda device=None: (8 << 30, 80 << 30), ) - tight_batch, tight_pending = _derive_store_staging_limits(torch.device("cuda")) + tight_batch, tight_pending = _derive_staging_limits(torch.device("cuda")) assert tight_batch == (8 << 30) // 10 assert tight_pending == (8 << 30) // 5 def test_store_cuda_staging_pool_reuses_preallocated_buffer(): """Store staging pool reuses its init-time allocation after release.""" - pool = StoreCudaStagingPool( + pool = FixedCudaStagingPool( device=torch.device("cpu"), buffer_bytes=128, depth=1, diff --git a/tests/connector/test_gds_transfer.py b/tests/connector/test_gds_transfer.py index 6627521..1378028 100644 --- a/tests/connector/test_gds_transfer.py +++ b/tests/connector/test_gds_transfer.py @@ -89,7 +89,7 @@ def test_missing_file_raises(tmp_path): def test_fixed_staging_pool_reuses_two_preallocated_buffers() -> None: """Fixed staging pool reuses bounded preallocated buffers.""" # First Party - from daser.connector.staging import FixedCudaStagingPool + from daser.connector.worker_memory import FixedCudaStagingPool pool = FixedCudaStagingPool( device=torch.device("cpu"), @@ -114,7 +114,7 @@ def test_fixed_staging_pool_reuses_two_preallocated_buffers() -> None: 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 + from daser.connector.worker_memory import FixedCudaStagingPool pool = FixedCudaStagingPool( device=torch.device("cpu"), @@ -139,7 +139,7 @@ def release_first() -> None: def test_fixed_staging_pool_rejects_oversized_request() -> None: """Fixed staging pool rejects requests larger than one buffer.""" # First Party - from daser.connector.staging import FixedCudaStagingPool + from daser.connector.worker_memory import FixedCudaStagingPool pool = FixedCudaStagingPool( device=torch.device("cpu"), diff --git a/tests/unit/test_block_aligned_bug.py b/tests/unit/test_block_aligned_bug.py index da03ed9..f4320de 100644 --- a/tests/unit/test_block_aligned_bug.py +++ b/tests/unit/test_block_aligned_bug.py @@ -8,10 +8,11 @@ # Third Party import pytest -# First Party -from daser.connector.daser_connector import DaserConnector from daser.connector.helpers import PendingStore, hash_tokens +# First Party +from daser.connector.request_lifecycle import RequestLifecycle + class TestTokeniseAndTruncateBug: def test_exact_block_aligned_length_not_multiple_of_block_tokens(self): @@ -73,7 +74,7 @@ def lookup(self, prefix, model_id): } ] - class MockDaserConnector(DaserConnector): + class MockDaserConnector(RequestLifecycle): def __init__(self): self._block_tokens = BLOCK_TOKENS self._socket_path = "/tmp/test.sock" @@ -121,7 +122,7 @@ def lookup(self, prefix, model_id): } ] - class MockDaserConnector(DaserConnector): + class MockDaserConnector(RequestLifecycle): def __init__(self): self._block_tokens = BLOCK_TOKENS self._socket_path = "/tmp/test.sock" @@ -169,7 +170,7 @@ def alloc_chunk(self, chunk_key, token_count, model_id): mock_ipc = MockIPCClientSync() - class MockDaserConnector(DaserConnector): + class MockDaserConnector(RequestLifecycle): def __init__(self): self._block_tokens = BLOCK_TOKENS self._socket_path = "/tmp/test.sock" @@ -230,7 +231,7 @@ def alloc_chunk(self, chunk_key, token_count, model_id): mock_ipc = MockIPCClientSync() - class MockDaserConnector(DaserConnector): + class MockDaserConnector(RequestLifecycle): def __init__(self): self._block_tokens = BLOCK_TOKENS self._socket_path = "/tmp/test.sock" @@ -299,7 +300,7 @@ def alloc_chunk(self, chunk_key, token_count, model_id): mock_ipc = MockIPCClientSync() tokens = list(range(630)) - class MockDaserConnector(DaserConnector): + class MockDaserConnector(RequestLifecycle): def __init__(self): self._block_tokens = BLOCK_TOKENS self._socket_path = "/tmp/test.sock" @@ -400,7 +401,7 @@ def lookup(self, prefix, model_id): } ] - class MockDaserConnector(DaserConnector): + class MockDaserConnector(RequestLifecycle): def __init__(self): self._block_tokens = BLOCK_TOKENS self._socket_path = "/tmp/test.sock" diff --git a/tests/unit/test_chunk_reuse_scheduler.py b/tests/unit/test_chunk_reuse_scheduler.py index 18ee782..79ad68e 100644 --- a/tests/unit/test_chunk_reuse_scheduler.py +++ b/tests/unit/test_chunk_reuse_scheduler.py @@ -4,12 +4,12 @@ from typing import Any # First Party -from daser.connector.scheduler import SchedulerConnectorMixin +from daser.connector.request_lifecycle import RequestLifecycle BLOCK_TOKENS = 4 -class _SchedulerProbe(SchedulerConnectorMixin): +class _SchedulerProbe(RequestLifecycle): """Minimal scheduler-role connector for chunk reuse credit tests.""" def __init__(self, chunks: list[dict[str, Any]]) -> None: diff --git a/tests/unit/test_skip_save_flag.py b/tests/unit/test_skip_save_flag.py index 7d9977d..6751153 100644 --- a/tests/unit/test_skip_save_flag.py +++ b/tests/unit/test_skip_save_flag.py @@ -14,10 +14,11 @@ # Standard from typing import Any, Optional -# First Party -from daser.connector.daser_connector import DaserConnector from daser.connector.helpers import PendingStore, hash_tokens +# First Party +from daser.connector.request_lifecycle import RequestLifecycle + BLOCK_TOKENS = 16 @@ -57,7 +58,7 @@ def record_external_prefix_cache(self, queries: int, hits: int) -> None: self.external_prefix_records.append((queries, hits)) -class _MockDaserConnector(DaserConnector): +class _MockDaserConnector(RequestLifecycle): """Test connector that bypasses vLLM init for scheduler-path testing. Mirrors the subclass pattern used by ``test_block_aligned_bug`` so From 00e1a6359c66428b07c73f03a7ac259113181f14 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Sun, 12 Jul 2026 18:32:01 +0800 Subject: [PATCH 2/4] refactor(connector): deepen worker pipeline ownership - move load and store state machines behind pipeline operation APIs - retain worker and scheduler adapters while narrowing runtime coordination - propagate delayed tensor-parallel geometry into pipeline ownership - select the staged CUDA device before background IPC export - remove duplicated config state and private compatibility surfaces --- daser/connector/daser_connector.py | 231 +-- daser/connector/load_pipeline.py | 586 +++++++- daser/connector/request_lifecycle.py | 12 + daser/connector/scheduler.py | 22 - daser/connector/store_pipeline.py | 393 ++++- daser/connector/worker_memory.py | 65 + daser/connector/worker_runtime.py | 1777 ++--------------------- daser/transfer/cuda_ipc.py | 27 + tests/connector/test_daser_connector.py | 1306 ++--------------- tests/unit/test_block_aligned_bug.py | 4 +- 10 files changed, 1320 insertions(+), 3103 deletions(-) diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index e43cbaa..d15e252 100644 --- a/daser/connector/daser_connector.py +++ b/daser/connector/daser_connector.py @@ -16,27 +16,12 @@ # First Party from daser.connector.ipc_client import IPCClientSync -from daser.connector.load_pipeline import ( - build_load_read_plan as _build_load_read_plan, -) from daser.connector.metadata import DaserConnectorMeta, ReqLoadSpec, ReqStoreSpec from daser.connector.request_lifecycle import RequestLifecycle -from daser.connector.reuse import CacheReuseStrategy, build_cache_reuse_strategy -from daser.connector.scheduler import ( - SchedulerConnectorMixin, - _block_ids_for_chunk, - _contiguous_prefix_tokens, - _trim_chunk_to_external_window, -) +from daser.connector.scheduler import SchedulerConnectorMixin from daser.connector.staging import ( DEFAULT_ROPE_DELTA_SCALE, ) -from daser.connector.staging import ( - apply_rope_delta_to_key_block as _apply_rope_delta_to_key_block, -) -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_runtime import WorkerRuntime from daser.logging import init_logger @@ -49,15 +34,33 @@ "DaserConnectorMeta", "ReqLoadSpec", "ReqStoreSpec", - "_apply_rope_delta_to_key_block", - "_build_load_read_plan", - "_block_ids_for_chunk", - "_contiguous_prefix_tokens", - "_copy_staging_to_kv_cache", - "_trim_chunk_to_external_window", ] +def _extract_rope_config(vllm_config: "VllmConfig") -> tuple[float, int, bool]: + """Extract worker RoPE geometry from the vLLM model config.""" + model_config = getattr(vllm_config, "model_config", None) + if model_config is None: + return 10000.0, 0, True + try: + head_size = int(model_config.get_head_size()) + except Exception: # noqa: BLE001 + logger.warning("[CONNECTOR] could not infer RoPE head size") + return 10000.0, 0, True + hf_text_config = getattr(model_config, "hf_text_config", None) + rope_parameters = getattr(hf_text_config, "rope_parameters", None) or {} + if not isinstance(rope_parameters, dict): + rope_parameters = {} + model_type = str(getattr(hf_text_config, "model_type", "")) + rope_base = ( + 1000000.0 + if "qwen" in model_type and "rope_theta" not in rope_parameters + else float(rope_parameters.get("rope_theta", 10000.0)) + ) + partial = float(rope_parameters.get("partial_rotary_factor", 1.0)) + return rope_base, int(head_size * partial), True + + class DaserConnector( SchedulerConnectorMixin, WorkerConnectorMixin, @@ -91,67 +94,46 @@ def __init__( ): extra = vllm_config.kv_transfer_config.kv_connector_extra_config or {} - self._socket_path: str = extra.get("socket_path", "/tmp/daser.sock") - self._store_path: str = "" - self._slot_size: int = 0 - parallel_config = getattr(vllm_config, "parallel_config", None) - self._tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1) or 1) - self._tp_rank = ( - get_tensor_model_parallel_rank() - if role == KVConnectorRole.WORKER and self._tp_size > 1 - else 0 - ) - self._server_tp_size = 1 - self._local_slot_size = 0 - self._rank_stride_bytes = 0 - self._block_tokens: int = 16 - self._model_id: str = "default" - self._skip_l2: bool = bool(extra.get("skip_l2", False)) - cache_reuse_mode = str(extra.get("cache_reuse_mode", "chunk")) - self._cache_reuse_strategy: CacheReuseStrategy - self._set_cache_reuse_strategy(cache_reuse_mode) - self._runtime_config_ready = False - self._rope_base: float = 10000.0 - self._rope_rotary_dim: int = 0 - self._rope_is_neox_style: bool = True - self._rope_delta_scale: float = float( - extra.get("rope_delta_scale", DEFAULT_ROPE_DELTA_SCALE) - ) - self._load_key_scale: float = float(extra.get("load_key_scale", 1.0)) - self._load_value_scale: float = float(extra.get("load_value_scale", 1.0)) - self._init_rope_config(vllm_config) - + socket_path = str(extra.get("socket_path", "/tmp/daser.sock")) if role == KVConnectorRole.SCHEDULER: - self._refresh_runtime_config() self._request_lifecycle = RequestLifecycle( - ipc_client=IPCClientSync(self._socket_path), - block_tokens=self._block_tokens, - slot_size=self._slot_size, - model_id=self._model_id, - cache_reuse_mode=self._cache_reuse_mode, - runtime_config_ready=self._runtime_config_ready, + ipc_client=IPCClientSync(socket_path), + block_tokens=16, + slot_size=0, + model_id="default", + cache_reuse_mode=str(extra.get("cache_reuse_mode", "chunk")), + runtime_config_ready=False, ) + self._request_lifecycle.refresh_runtime_config() else: + parallel_config = getattr(vllm_config, "parallel_config", None) + tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1) or 1) + tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0 + rope_base, rope_rotary_dim, rope_is_neox_style = _extract_rope_config( + vllm_config + ) self._worker_runtime = WorkerRuntime( - socket_path=self._socket_path, + socket_path=socket_path, transfer_mode=str(extra.get("transfer_mode", "iouring")), - skip_l2=self._skip_l2, - tp_size=self._tp_size, - tp_rank=self._tp_rank, - server_tp_size=self._server_tp_size, - slot_size=self._slot_size, - store_path=self._store_path, - rank_stride_bytes=self._rank_stride_bytes, - rope_base=self._rope_base, - rope_rotary_dim=self._rope_rotary_dim, - rope_is_neox_style=self._rope_is_neox_style, - rope_delta_scale=self._rope_delta_scale, - load_key_scale=self._load_key_scale, - load_value_scale=self._load_value_scale, + skip_l2=bool(extra.get("skip_l2", False)), + tp_size=tp_size, + tp_rank=tp_rank, + server_tp_size=1, + slot_size=0, + store_path="", + rank_stride_bytes=0, + rope_base=rope_base, + rope_rotary_dim=rope_rotary_dim, + rope_is_neox_style=rope_is_neox_style, + rope_delta_scale=float( + extra.get("rope_delta_scale", DEFAULT_ROPE_DELTA_SCALE) + ), + load_key_scale=float(extra.get("load_key_scale", 1.0)), + load_value_scale=float(extra.get("load_value_scale", 1.0)), kv_cache_config=kv_cache_config, ) - logger.info("[CONNECTOR] role=%s socket=%s", role.name, self._socket_path) + logger.info("[CONNECTOR] role=%s socket=%s", role.name, socket_path) @property def prefer_cross_layer_blocks(self) -> bool: @@ -182,104 +164,3 @@ def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: Class-level config helper with no mutable state. """ return "NHD" - - def _refresh_runtime_config(self) -> None: - """Refresh server-owned runtime config over IPC when available.""" - client = getattr(self, "_ipc_sync", None) - owns_client = client is None - if client is None: - client = IPCClientSync(self._socket_path) - try: - config = client.get_runtime_config() - except Exception as exc: # noqa: BLE001 - logger.info("[CONNECTOR] runtime config unavailable: %s", exc) - return - finally: - if owns_client: - close = getattr(client, "close", None) - if close is not None: - close() - - self._store_path = str(config.get("store_path", self._store_path)) - self._slot_size = int(config.get("slot_size", self._slot_size)) - self._server_tp_size = int( - config.get("tensor_parallel_size", self._server_tp_size) - ) - self._rank_stride_bytes = int( - config.get("rank_stride_bytes", self._rank_stride_bytes) - ) - self._block_tokens = int(config.get("block_tokens", self._block_tokens)) - self._model_id = str(config.get("model_id", self._model_id)) - self._set_cache_reuse_strategy(str(config["cache_reuse_mode"])) - self._skip_l2 = bool(config.get("skip_l2", self._skip_l2)) - self._runtime_config_ready = bool( - self._slot_size and (self._store_path or self._skip_l2) - ) - self._transfer_mode = str( - config.get("transfer_mode", getattr(self, "_transfer_mode", "iouring")) - ) - logger.info( - "[CONNECTOR] runtime config store=%s slot_size=%d block_tokens=%d " - "model=%s transfer=%s skip_l2=%s", - self._store_path, - self._slot_size, - self._block_tokens, - self._model_id, - getattr(self, "_transfer_mode", "iouring"), - self._skip_l2, - ) - - def _set_cache_reuse_strategy(self, cache_reuse_mode: str) -> None: - """Set scheduler cache reuse strategy. - - Args: - cache_reuse_mode: either ``"chunk"`` or ``"prefix"``. - """ - self._cache_reuse_mode = cache_reuse_mode - self._cache_reuse_strategy = build_cache_reuse_strategy( - cache_reuse_mode, - self._block_tokens, - ) - - def _init_rope_config(self, vllm_config: "VllmConfig") -> None: - """Extract default RoPE settings from vLLM model config. - - Args: - vllm_config: vLLM runtime config passed to the connector. - """ - model_config = getattr(vllm_config, "model_config", None) - if model_config is None: - return - try: - head_size = int(model_config.get_head_size()) - except Exception: # noqa: BLE001 - logger.warning("[CONNECTOR] could not infer RoPE head size") - return - - hf_text_config = getattr(model_config, "hf_text_config", None) - rope_parameters = getattr(hf_text_config, "rope_parameters", None) or {} - if not isinstance(rope_parameters, dict): - rope_parameters = {} - model_type = str(getattr(hf_text_config, "model_type", "")) - if "qwen" in model_type and "rope_theta" not in rope_parameters: - rope_base = 1000000.0 - else: - rope_base = float(rope_parameters.get("rope_theta", 10000.0)) - partial = float(rope_parameters.get("partial_rotary_factor", 1.0)) - rotary_dim = int(head_size * partial) - - self._rope_base = rope_base - self._rope_rotary_dim = rotary_dim - self._rope_is_neox_style = True - logger.info( - "[CONNECTOR] rope base=%s rotary_dim=%d neox=%s", - self._rope_base, - self._rope_rotary_dim, - self._rope_is_neox_style, - ) - logger.info( - "[CONNECTOR] load tuning rope_delta_scale=%s key_scale=%s value_scale=%s", - self._rope_delta_scale, - self._load_key_scale, - self._load_value_scale, - ) diff --git a/daser/connector/load_pipeline.py b/daser/connector/load_pipeline.py index 5c9ebb9..7c03b8d 100644 --- a/daser/connector/load_pipeline.py +++ b/daser/connector/load_pipeline.py @@ -3,25 +3,88 @@ from __future__ import annotations import asyncio +from collections import deque +from concurrent.futures import Future from dataclasses import dataclass, replace +import os import threading from typing import Any +# Third Party +import cupy +import torch + +from daser.connector.helpers import base_req_id from daser.connector.ipc_client import IPCClientAsync from daser.connector.metadata import ReqLoadSpec -from daser.connector.worker_memory import FixedCudaStagingPool +from daser.connector.staging import copy_staging_to_kv_cache +from daser.connector.worker_memory import ( + CudaStagingLease, + FixedCudaStagingPool, +) +from daser.logging import init_logger +from daser.transfer.cuda_ipc import ( + cuda_allocation_base_and_offset, + cuda_array_device_id, + cuda_array_pointer, + export_cuda_ipc_handle, +) + +logger = init_logger(__name__) + +_LOAD_DISPATCH_WAIT_TIMEOUT_S = 0.001 +_LoadBatch = tuple[int, list[dict[str, int]], list[Any]] + + +@dataclass +class _PendingLoad: + """Track one request load until vLLM can resume it.""" + + future: Future[None] + block_ids: list[int] + + +@dataclass +class _InflightLoadBatch: + """Hold one submitted load batch and its fixed staging lease.""" + + buffer_index: int + per_req_ranges: list[Any] + staging_lease: CudaStagingLease + future: Any + + +@dataclass +class _QueuedLoadRequest: + """Represent one request waiting for a load staging buffer.""" + + req_id: str + spec_id: str + spec: ReqLoadSpec + future: Future[None] + + +@dataclass +class _InflightRequestLoad: + """Track active and remaining load batches for one request.""" + + item: _QueuedLoadRequest + buffer_index: int + batches: list[_LoadBatch] + next_batch: int + active: _InflightLoadBatch | None class LoadPipeline: - """Own load-side asyncio, IPC, staging, queue, and completion state. + """Own the complete worker load state machine. Args: socket_path: DaseR server Unix socket path. - client_count: Number of independent async IPC connections. + client_count: Independent load IPC lanes and maximum inflight requests. Async/thread-safety: - Construction and public methods run on the vLLM worker thread. Async - IPC runs exclusively on the private ``daser-load-io`` thread. + Public methods are called on the vLLM worker thread. Queue dispatch, + IPC, and CUDA restore execute on the private load thread. """ def __init__(self, socket_path: str, client_count: int) -> None: @@ -39,84 +102,231 @@ def __init__(self, socket_path: str, client_count: int) -> None: self._invalid_block_ids: set[int] = set() self._staging_pool: FixedCudaStagingPool | None = None self._staging_registered = False + self._kv_caches: dict[str, torch.Tensor] = {} + self._layer_names: list[str] = [] + self._local_slot_size = 0 + self._rank_stride_bytes = 0 + self._tp_rank = 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._cuda_stream: torch.cuda.Stream | None = None self._thread.start() - @property - def pending(self) -> dict[str, Any]: - """Return worker-thread-owned request completion records.""" - return self._pending - - @property - def invalid_block_ids(self) -> set[int]: - """Return load-error block IDs accumulated by completion polling.""" - return self._invalid_block_ids - - @property - def staging_pool(self) -> FixedCudaStagingPool | None: - """Return the configured fixed staging pool, if registered.""" - return self._staging_pool - - @property - def staging_registered(self) -> bool: - """Return whether every fixed staging buffer has CUDA IPC registration.""" - return self._staging_registered - - @staging_registered.setter - def staging_registered(self, value: bool) -> None: - self._staging_registered = value - - def set_staging_pool(self, pool: FixedCudaStagingPool) -> None: - """Install the fixed load staging pool before request traffic.""" - self._staging_pool = pool + def configure( + self, + *, + kv_caches: dict[str, torch.Tensor], + layer_names: list[str], + local_slot_size: int, + rank_stride_bytes: int, + tp_rank: int, + staging_pool: FixedCudaStagingPool, + load_key_scale: float, + load_value_scale: float, + rope_delta_scale: float, + rope_base: float, + rope_rotary_dim: int, + rope_is_neox_style: bool, + ) -> None: + """Configure immutable KV layout and transform state. + + Args: + kv_caches: Registered vLLM KV tensors. + layer_names: Stable storage layer order. + local_slot_size: Bytes stored per slot by this TP rank. + rank_stride_bytes: Byte distance between rank lanes. + tp_rank: Current tensor-parallel rank. + staging_pool: Fixed load staging buffers. + load_key_scale: Load-time key scaling factor. + load_value_scale: Load-time value scaling factor. + rope_delta_scale: Position-offset scaling factor. + rope_base: RoPE theta/base. + rope_rotary_dim: Number of dimensions covered by RoPE. + rope_is_neox_style: Whether RoPE uses split-half rotation. + + Async/thread-safety: + Called once on the worker thread before request traffic. + """ + self._kv_caches = kv_caches + self._layer_names = list(layer_names) + self._local_slot_size = local_slot_size + self._rank_stride_bytes = rank_stride_bytes + self._tp_rank = tp_rank + self._staging_pool = staging_pool self._staging_registered = False + self._load_key_scale = load_key_scale + self._load_value_scale = load_value_scale + self._rope_delta_scale = rope_delta_scale + self._rope_base = rope_base + self._rope_rotary_dim = rope_rotary_dim + self._rope_is_neox_style = rope_is_neox_style + + def initialize_transfer(self) -> None: + """Initialize load IPC lanes and register staging buffers. + + Async/thread-safety: + Called on the worker thread during startup. IPC runs on the load + loop and is joined before this method returns. + """ + for client in self._clients: + self._submit(client.init_transfer()).result(timeout=120.0) + self._register_staging_buffers() + + def configure_rank_geometry(self, rank_stride_bytes: int, tp_rank: int) -> None: + """Apply server-finalized tensor-parallel lane geometry. + + Args: + rank_stride_bytes: Byte distance between server-owned rank lanes. + tp_rank: Current tensor-parallel rank. + + Async/thread-safety: + Called on the worker thread after runtime-config refresh and before + any load is submitted. + """ + self._rank_stride_bytes = rank_stride_bytes + self._tp_rank = tp_rank + + def start(self, reqs_to_load: dict[str, ReqLoadSpec]) -> None: + """Queue request loads for background transfer and restore. + + Args: + reqs_to_load: Scheduler load metadata keyed by request/spec ID. + + Async/thread-safety: + Called from the worker thread. Queue dispatch, IPC, and restore run + on the private load thread. + """ + if not reqs_to_load: + return + if not self._layer_names or not self._kv_caches: + self.mark_failed(reqs_to_load, "no registered KV cache layout") + return + queue = self._ensure_queue() + for spec_id, spec in reqs_to_load.items(): + req_id = base_req_id(spec_id) + request_future: Future[None] = Future() + self._pending[req_id] = _PendingLoad( + future=request_future, + block_ids=list(spec.block_ids), + ) + self._loop.call_soon_threadsafe( + queue.put_nowait, + _QueuedLoadRequest(req_id, spec_id, spec, request_future), + ) + self._ensure_dispatcher() + + def mark_failed( + self, + reqs_to_load: dict[str, ReqLoadSpec], + reason: str, + ) -> None: + """Record submission failures for completion polling. + + Args: + reqs_to_load: Load specs that could not be submitted. + reason: Diagnostic failure reason. + + Async/thread-safety: + Called on the worker thread before background submission. + """ + block_ids = [ + block_id for spec in reqs_to_load.values() for block_id in spec.block_ids + ] + failed_future: Future[None] = Future() + failed_future.set_exception(RuntimeError(reason)) + for req_id in {base_req_id(req_id) for req_id in reqs_to_load}: + self._pending[req_id] = _PendingLoad(failed_future, list(block_ids)) + + def collect_finished(self) -> set[str]: + """Collect completed loads without blocking the worker thread. + + Returns: + Base request IDs whose load lifecycle completed in this poll. + + Async/thread-safety: + Called on the worker thread; request futures provide cross-thread + visibility from the load thread. + """ + finished: set[str] = set() + collected_futures: set[int] = set() + for req_id, load in list(self._pending.items()): + if not load.future.done(): + continue + future_id = id(load.future) + try: + if future_id not in collected_futures: + load.future.result() + collected_futures.add(future_id) + except Exception as exc: # noqa: BLE001 + logger.warning( + "[CONNECTOR] async load failed req=%s blocks=%s: %s", + req_id, + load.block_ids, + exc, + ) + self._invalid_block_ids.update(load.block_ids) + collected_futures.add(future_id) + finally: + del self._pending[req_id] + finished.add(req_id) + return finished + + def take_invalid_block_ids(self) -> set[int]: + """Return and clear block IDs targeted by failed loads. + + Returns: + vLLM block IDs that must be invalidated. + + Async/thread-safety: + Called on the worker thread after ``collect_finished``. + """ + invalid = set(self._invalid_block_ids) + self._invalid_block_ids.clear() + return invalid + + def shutdown(self) -> None: + """Drain queue ownership, close IPC clients, and stop the load loop.""" + for load in {id(item.future): item for item in self._pending.values()}.values(): + if not load.future.done(): + try: + load.future.result(timeout=120.0) + except Exception: # noqa: BLE001 + pass + self.collect_finished() + if self._queue is not None: + self._loop.call_soon_threadsafe(self._queue.put_nowait, None) + if self._dispatcher_future is not None: + self._dispatcher_future.result(timeout=120.0) + for client in self._clients: + self._submit(client.close()).result(timeout=5.0) + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=5.0) - def submit(self, coro: Any) -> Any: - """Submit a coroutine to the private load event loop.""" + def _submit(self, coro: Any) -> Any: return asyncio.run_coroutine_threadsafe(coro, self._loop) - def client(self, buffer_index: int | None = None) -> IPCClientAsync: - """Return the IPC client assigned to a staging buffer index.""" + def _client(self, buffer_index: int | None = None) -> IPCClientAsync: index = 0 if buffer_index is None else int(buffer_index) return self._clients[index % len(self._clients)] - def clients(self) -> tuple[IPCClientAsync, ...]: - """Return all load IPC clients for startup and shutdown orchestration.""" - return tuple(self._clients) - - def ensure_queue(self) -> asyncio.Queue[Any]: - """Create the load request queue on its owning event loop once.""" + def _ensure_queue(self) -> asyncio.Queue[Any]: with self._queue_lock: if self._queue is None: - future = self.submit(self._create_queue()) - self._queue = future.result(timeout=5.0) + self._queue = self._submit(self._create_queue()).result(timeout=5.0) return self._queue - def enqueue(self, item: Any) -> None: - """Append one request from the worker thread without blocking.""" - queue = self.ensure_queue() - self._loop.call_soon_threadsafe(queue.put_nowait, item) - - def ensure_dispatcher(self, coro: Any) -> None: - """Start the single queue dispatcher when none is active.""" + def _ensure_dispatcher(self) -> None: with self._queue_lock: if ( self._dispatcher_future is not None and not self._dispatcher_future.done() ): - coro.close() return - self._dispatcher_future = self.submit(coro) - - def shutdown(self) -> None: - """Drain queue ownership, close IPC clients, and stop the load loop.""" - if self._queue is not None: - self._loop.call_soon_threadsafe(self._queue.put_nowait, None) - if self._dispatcher_future is not None: - self._dispatcher_future.result(timeout=120.0) - for client in self._clients: - self.submit(client.close()).result(timeout=5.0) - self._loop.call_soon_threadsafe(self._loop.stop) - self._thread.join(timeout=5.0) + self._dispatcher_future = self._submit(self._run_dispatcher()) async def _create_queue(self) -> asyncio.Queue[Any]: return asyncio.Queue() @@ -125,6 +335,254 @@ def _run_loop(self) -> None: asyncio.set_event_loop(self._loop) self._loop.run_forever() + async def _run_dispatcher(self) -> None: + sample_tensor = next(iter(self._kv_caches.values())) + if sample_tensor.device.type == "cuda": + torch.cuda.set_device(sample_tensor.device) + if self._cuda_stream is None: + self._cuda_stream = torch.cuda.Stream(device=sample_tensor.device) + queue = self._ensure_queue() + if self._staging_pool is None: + raise RuntimeError("load staging pool is not configured") + free_buffers = deque( + range(max(1, min(len(self._clients), self._staging_pool.depth))) + ) + queued: deque[_QueuedLoadRequest] = deque() + active: list[_InflightRequestLoad] = [] + try: + while True: + if not queued: + if active: + await self._drain_queue(queue, queued) + else: + item = await queue.get() + if item is None: + return + queued.append(item) + while queued and free_buffers: + state = self._submit_request( + queued.popleft(), free_buffers.popleft() + ) + if state.active is None: + free_buffers.append(state.buffer_index) + else: + active.append(state) + consumed = False + for state in list(active): + active_batch = state.active + if active_batch is not None and not active_batch.future.done(): + continue + reusable_buffer, request_done = self._consume_request(state) + if request_done: + active.remove(state) + free_buffers.append(reusable_buffer) + consumed = True + if consumed: + continue + if active: + await self._wait_for_completion(active) + except BaseException as exc: + for state in active: + if not state.item.future.done(): + state.item.future.set_exception(exc) + if state.active is not None: + state.active.staging_lease.release() + for item in queued: + if not item.future.done(): + item.future.set_exception(exc) + raise + + async def _drain_queue( + self, + queue: asyncio.Queue[Any], + queued: deque[_QueuedLoadRequest], + ) -> None: + while True: + try: + item = queue.get_nowait() + except asyncio.QueueEmpty: + return + if item is None: + await queue.put(None) + return + queued.append(item) + + async def _wait_for_completion( + self, + active: list[_InflightRequestLoad], + ) -> None: + wrapped = { + asyncio.wrap_future(state.active.future) + for state in active + if state.active is not None + } + if not wrapped: + await asyncio.sleep(0) + return + done, _pending = await asyncio.wait( + wrapped, + timeout=_LOAD_DISPATCH_WAIT_TIMEOUT_S, + return_when=asyncio.FIRST_COMPLETED, + ) + if not done: + await asyncio.sleep(0) + + def _submit_request( + self, + item: _QueuedLoadRequest, + buffer_index: int, + ) -> _InflightRequestLoad: + if self._staging_pool is None: + raise RuntimeError("load staging pool is not configured") + spec = replace( + item.spec, + file_offset=( + self._tp_rank * self._rank_stride_bytes + + item.spec.start_slot * self._local_slot_size + ), + ) + batches = build_load_read_batches( + {item.spec_id: spec}, + self._local_slot_size, + max_batch_bytes=self._staging_pool.buffer_bytes, + include_req_ids=True, + ) + if not batches: + item.future.set_result(None) + return _InflightRequestLoad(item, buffer_index, [], 0, 0, None) + return _InflightRequestLoad( + item=item, + buffer_index=buffer_index, + batches=batches, + next_batch=1, + active=self._submit_batch(batches[0], buffer_index), + ) + + def _submit_batch( + self, + batch: _LoadBatch, + buffer_index: int, + ) -> _InflightLoadBatch: + if self._staging_pool is None: + raise RuntimeError("load staging pool is not configured") + total_bytes, spans, per_req_ranges = batch + lease = self._staging_pool.acquire_index(buffer_index, total_bytes) + staging = lease.view + if self._staging_registered: + transfer = self._client(buffer_index).transfer_load_registered_cuda( + buffer_index=buffer_index, + producer_pid=os.getpid(), + nbytes=total_bytes, + spans=spans, + ) + else: + cp_staging = cupy.asarray(staging) + device_ptr = cuda_array_pointer(cp_staging) + allocation_base, allocation_offset = cuda_allocation_base_and_offset( + device_ptr + ) + transfer = self._client(buffer_index).transfer_load_cuda( + cuda_ipc_handle=export_cuda_ipc_handle(cp_staging), + nbytes=total_bytes, + device_id=cuda_array_device_id(cp_staging), + device_ptr=device_ptr, + allocation_base_ptr=allocation_base, + allocation_offset=allocation_offset, + producer_pid=os.getpid(), + spans=spans, + ) + return _InflightLoadBatch( + buffer_index=buffer_index, + per_req_ranges=per_req_ranges, + staging_lease=lease, + future=self._submit(transfer), + ) + + def _consume_request( + self, + state: _InflightRequestLoad, + ) -> tuple[int, bool]: + active = state.active + if active is None: + return state.buffer_index, True + self._consume_batch(active) + state.buffer_index = active.buffer_index + if state.next_batch >= len(state.batches): + if not state.item.future.done(): + state.item.future.set_result(None) + return state.buffer_index, True + state.active = self._submit_batch( + state.batches[state.next_batch], + state.buffer_index, + ) + state.next_batch += 1 + return state.buffer_index, False + + def _consume_batch(self, state: _InflightLoadBatch) -> None: + try: + state.future.result(timeout=120.0) + if self._cuda_stream is None: + self._restore_batch(state) + else: + with torch.cuda.stream(self._cuda_stream): + self._restore_batch(state) + self._cuda_stream.synchronize() + finally: + state.staging_lease.release() + + def _restore_batch(self, state: _InflightLoadBatch) -> None: + staging = state.staging_lease.view + for run in build_load_copy_runs(state.per_req_ranges): + 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._local_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, + ) + + def _register_staging_buffers(self) -> None: + pool = self._staging_pool + if self._staging_registered or pool is None: + return + try: + for buffer_index in range(pool.depth): + tensor = pool.buffer(buffer_index) + cp_tensor = cupy.asarray(tensor) + device_ptr = cuda_array_pointer(cp_tensor) + allocation_base, allocation_offset = cuda_allocation_base_and_offset( + device_ptr + ) + self._submit( + self._client(buffer_index).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=allocation_base, + allocation_offset=allocation_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._staging_registered = False + return + self._staging_registered = True + logger.info("[CONNECTOR] registered %d load staging buffers", pool.depth) + @dataclass(frozen=True) class LoadCopyRun: diff --git a/daser/connector/request_lifecycle.py b/daser/connector/request_lifecycle.py index 338d611..71e94c0 100644 --- a/daser/connector/request_lifecycle.py +++ b/daser/connector/request_lifecycle.py @@ -331,6 +331,18 @@ def build_connector_meta( ) return meta + def refresh_runtime_config(self) -> None: + """Refresh scheduler geometry and reuse policy from the DaseR server. + + Returns: + None. + + Async/thread-safety: + Called on the scheduler thread and performs synchronous control-plane + IPC during connector initialization or recovery, never worker IO. + """ + self._refresh_runtime_config() + def _drop_preempted_pending_state( self, scheduler_output: "SchedulerOutput", diff --git a/daser/connector/scheduler.py b/daser/connector/scheduler.py index 074d79b..b0f009f 100644 --- a/daser/connector/scheduler.py +++ b/daser/connector/scheduler.py @@ -10,28 +10,6 @@ from vllm.v1.request import Request from daser.connector.metadata import DaserConnectorMeta -from daser.connector.scheduler_planning import ( - _base_req_id, - _block_ids_for_chunk, - _computed_tokens_after_step, - _contiguous_prefix_tokens, - _get_kv_transfer_flag, - _matches_request_or_store_id, - _store_slot_index, - _trim_chunk_to_external_window, -) - -__all__ = [ - "SchedulerConnectorMixin", - "_base_req_id", - "_block_ids_for_chunk", - "_computed_tokens_after_step", - "_contiguous_prefix_tokens", - "_get_kv_transfer_flag", - "_matches_request_or_store_id", - "_store_slot_index", - "_trim_chunk_to_external_window", -] class SchedulerConnectorMixin: diff --git a/daser/connector/store_pipeline.py b/daser/connector/store_pipeline.py index 42d2080..13d9782 100644 --- a/daser/connector/store_pipeline.py +++ b/daser/connector/store_pipeline.py @@ -3,29 +3,58 @@ from __future__ import annotations import asyncio -from dataclasses import dataclass +from dataclasses import dataclass, replace +import os import threading from typing import Any +import cupy import torch +from daser.connector.helpers import base_req_id from daser.connector.ipc_client import IPCClientAsync from daser.connector.metadata import ReqStoreSpec, StoreWriteSpan +from daser.connector.staging import ( + CROSS_LAYER_KV_CACHE_KEY, + copy_cross_layer_kv_cache_to_staging, + copy_kv_cache_to_staging, + record_cuda_event, +) from daser.connector.worker_memory import ( DEFAULT_STORE_STAGING_BYTES, CudaStagingLease, + FixedCudaStagingPool, +) +from daser.logging import init_logger +from daser.transfer.cuda_ipc import ( + cuda_allocation_base_and_offset, + cuda_array_device_id, + cuda_array_pointer, + export_cuda_ipc_handle, ) +logger = init_logger(__name__) + + +@dataclass +class _DeferredFinishedSave: + """Hold request store work until vLLM reports it finished.""" + + commit_keys: set[str] + reqs_to_store: dict[str, ReqStoreSpec] + finished: bool = False + future: Any | None = None + class StorePipeline: - """Own store-side asyncio, IPC, staging, and future state. + """Own the complete worker store state machine. Args: socket_path: DaseR server Unix socket path. Async/thread-safety: - Construction and public methods run on the vLLM worker thread. Async - IPC runs exclusively on the private ``daser-store-io`` thread. + Public methods are called on the vLLM worker thread. Snapshot, IPC, and + commit execute on the private store thread and its fixed CUDA stream. """ def __init__(self, socket_path: str) -> None: @@ -36,42 +65,360 @@ def __init__(self, socket_path: str) -> None: daemon=True, name="daser-store-io", ) - self.save_futures: list[Any] = [] - self.pending_staging_bytes = 0 - self.staging_bytes = 0 - self.pending_staging_limit_bytes = 0 - self.staging_pool: Any | None = None - self.pending_finished_saves: dict[str, Any] = {} + self._staging_bytes = 0 + self._staging_pool: FixedCudaStagingPool | None = None + self._pending_finished_saves: dict[str, _DeferredFinishedSave] = {} + self._kv_caches: dict[str, torch.Tensor] = {} + self._layer_names: list[str] = [] + self._layer_idx_map: dict[str, int] = {} + self._local_slot_size = 0 + self._rank_stride_bytes = 0 + self._tp_rank = 0 + self._tp_size = 1 + self._cuda_stream: torch.cuda.Stream | None = None self._thread.start() - @property - def client(self) -> IPCClientAsync: - """Return the store IPC client for pipeline-owned coroutines.""" - return self._client + def configure( + self, + *, + kv_caches: dict[str, torch.Tensor], + layer_names: list[str], + layer_idx_map: dict[str, int], + local_slot_size: int, + rank_stride_bytes: int, + tp_rank: int, + tp_size: int, + staging_bytes: int, + staging_pool: FixedCudaStagingPool, + ) -> None: + """Configure immutable KV layout and staging state. - def submit(self, coro: Any) -> Any: - """Submit a coroutine to the private store event loop.""" - return asyncio.run_coroutine_threadsafe(coro, self._loop) + Args: + kv_caches: Registered vLLM KV tensors. + layer_names: Stable storage layer order. + layer_idx_map: Layer names mapped to storage indices. + local_slot_size: Bytes stored per slot by this TP rank. + rank_stride_bytes: Byte distance between rank lanes. + tp_rank: Current tensor-parallel rank. + tp_size: Tensor-parallel world size. + staging_bytes: Maximum bytes per store batch. + staging_pool: Fixed store staging buffers. + + Async/thread-safety: + Called once on the worker thread before request traffic. + """ + self._kv_caches = kv_caches + self._layer_names = list(layer_names) + self._layer_idx_map = dict(layer_idx_map) + self._local_slot_size = local_slot_size + self._rank_stride_bytes = rank_stride_bytes + self._tp_rank = tp_rank + self._tp_size = tp_size + self._staging_bytes = staging_bytes + self._staging_pool = staging_pool + + def initialize_transfer(self) -> None: + """Initialize the store IPC transfer client on its event loop. + + Async/thread-safety: + Called on the worker thread during startup and waits only for the + store loop's initialization future. + """ + self._submit(self._client.init_transfer()).result(timeout=120.0) + + def configure_rank_geometry( + self, + rank_stride_bytes: int, + tp_rank: int, + tp_size: int, + ) -> None: + """Apply server-finalized tensor-parallel lane geometry. + + Args: + rank_stride_bytes: Byte distance between server-owned rank lanes. + tp_rank: Current tensor-parallel rank. + tp_size: Tensor-parallel world size used for commit coordination. + + Async/thread-safety: + Called on the worker thread after runtime-config refresh and before + any store is submitted. + """ + self._rank_stride_bytes = rank_stride_bytes + self._tp_rank = tp_rank + self._tp_size = tp_size + + def queue_finished( + self, + reqs_to_store: dict[str, ReqStoreSpec], + commit_keys: set[str], + ) -> None: + """Queue stores until vLLM reports their requests finished. + + Args: + reqs_to_store: Store metadata for the current worker step. + commit_keys: Chunk keys eligible for commit after transfer. + + Async/thread-safety: + Called on the worker thread to accumulate immutable store intent. + CUDA ordering is captured later, when vLLM reports completion. + """ + for req_id, spec in reqs_to_store.items(): + base_id = base_req_id(req_id) + save = self._pending_finished_saves.get(base_id) + if save is None: + save = _DeferredFinishedSave(set(), {}) + self._pending_finished_saves[base_id] = save + save.reqs_to_store[req_id] = spec + if spec.chunk_key in commit_keys: + save.commit_keys.add(spec.chunk_key) + + def collect_finished(self, finished_req_ids: set[str]) -> set[str]: + """Submit newly finished stores and collect completed requests. + + Args: + finished_req_ids: Requests vLLM finished in this step. + + Returns: + Request IDs whose store and commit lifecycle has completed. + + Async/thread-safety: + Called on the worker thread. Store work runs on the private loop. + """ + finished: set[str] = set() + for req_id in finished_req_ids: + save = self._pending_finished_saves.get(req_id) + if save is not None: + save.finished = True + for req_id, save in list(self._pending_finished_saves.items()): + if not save.finished and save.future is None: + continue + if save.future is not None and save.future.done(): + save.future.result(timeout=120.0) + finished.add(req_id) + del self._pending_finished_saves[req_id] + has_inflight = any( + save.future is not None for save in self._pending_finished_saves.values() + ) + if not has_inflight: + next_save = next( + ( + save + for save in self._pending_finished_saves.values() + if save.finished and save.future is None + ), + None, + ) + if next_save is not None: + self._submit_save(next_save) + return finished def shutdown(self) -> None: - """Close the IPC client and stop the private store event loop.""" - self.submit(self._client.close()).result(timeout=5.0) - self._loop.call_soon_threadsafe(self._loop.stop) - self._thread.join(timeout=5.0) + """Finish queued stores, close IPC, and stop the store loop. + + Async/thread-safety: + Called once on the worker thread after request traffic stops. + """ + first_error: BaseException | None = None + try: + submitted_ids = [ + req_id + for req_id, save in self._pending_finished_saves.items() + if save.future is not None + ] + for req_id in submitted_ids: + save = self._pending_finished_saves[req_id] + try: + if save.future is not None: + save.future.result(timeout=120.0) + except BaseException as exc: # preserve cleanup during shutdown + if first_error is None: + first_error = exc + finally: + del self._pending_finished_saves[req_id] + for req_id in list(self._pending_finished_saves): + save = self._pending_finished_saves[req_id] + try: + self._submit_save(save) + save.future.result(timeout=120.0) + except BaseException as exc: + if first_error is None: + first_error = exc + finally: + del self._pending_finished_saves[req_id] + try: + self._submit(self._client.close()).result(timeout=5.0) + except BaseException as exc: + if first_error is None: + first_error = exc + finally: + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=5.0) + if first_error is not None: + raise first_error + + def _submit(self, coro: Any) -> Any: + return asyncio.run_coroutine_threadsafe(coro, self._loop) def _run_loop(self) -> None: asyncio.set_event_loop(self._loop) self._loop.run_forever() + def _submit_save(self, save: _DeferredFinishedSave) -> None: + """Capture producer ordering and submit one save to the store thread.""" + sample = next(iter(self._kv_caches.values()), None) + producer_event = record_cuda_event(sample) if sample is not None else None + save.future = self._submit(self._store_finished_save(save, producer_event)) + + def _stage_finished_save( + self, + save: _DeferredFinishedSave, + producer_event: torch.cuda.Event | None, + ) -> list["StagedStoreBatch"] | None: + if not self._kv_caches or self._staging_pool is None: + return [] + reqs_to_store = { + req_id: replace( + spec, + file_offset=( + self._tp_rank * self._rank_stride_bytes + + spec.start_slot * self._local_slot_size + ), + ) + for req_id, spec in save.reqs_to_store.items() + } + batches = build_staging_store_batches( + reqs_to_store, + self._local_slot_size, + max_batch_bytes=self._staging_bytes, + ) + if len(batches) > self._staging_pool.available: + return None + staged_batches: list[StagedStoreBatch] = [] + try: + for block_ids, spans in batches: + staged_batches.append( + self._stage_batch(block_ids, spans, producer_event) + ) + except BaseException: + for staged in staged_batches: + staged.lease.release() + raise + return staged_batches + + async def _store_finished_save( + self, + save: _DeferredFinishedSave, + producer_event: torch.cuda.Event | None, + ) -> None: + staged_batches = self._stage_finished_save(save, producer_event) + if staged_batches is None: + raise RuntimeError("store staging capacity cannot satisfy one request") + stored_keys: list[str] = [] + for staged in staged_batches: + try: + stored_keys.extend(await self._write_cuda_buffer(staged)) + finally: + staged.lease.release() + requested = save.commit_keys + keys_to_commit = list( + dict.fromkeys(key for key in stored_keys if key in requested) + ) + await self._client.commit_chunks( + keys_to_commit, + tp_rank=self._tp_rank, + tp_size=self._tp_size, + ) + + def _stage_batch( + self, + block_ids: list[int], + spans: list[StoreWriteSpan], + producer_event: torch.cuda.Event | None, + ) -> "StagedStoreBatch": + if self._staging_pool is None: + raise RuntimeError("store staging pool is not configured") + sample = next(iter(self._kv_caches.values())) + nbytes = len(block_ids) * self._local_slot_size + lease = self._staging_pool.acquire(nbytes) + stream = self._cuda_stream + if sample.device.type == "cuda" and stream is None: + torch.cuda.set_device(sample.device) + stream = torch.cuda.Stream(device=sample.device) + self._cuda_stream = stream + if stream is not None: + if producer_event is not None: + stream.wait_event(producer_event) + with torch.cuda.stream(stream): + self._copy_blocks(lease.view, block_ids, sample) + stream.synchronize() + else: + self._copy_blocks(lease.view, block_ids, sample) + return StagedStoreBatch(lease.view, spans, lease) + + def _copy_blocks( + self, + staging: torch.Tensor, + block_ids: list[int], + sample: torch.Tensor, + ) -> None: + block_index = torch.tensor(block_ids, dtype=torch.long, device=sample.device) + cross_layer = self._kv_caches.get(CROSS_LAYER_KV_CACHE_KEY) + if cross_layer is not None: + copy_cross_layer_kv_cache_to_staging( + staging, + cross_layer, + block_ids, + len(self._layer_names), + self._local_slot_size, + block_index, + ) + return + for layer_name in self._layer_names: + copy_kv_cache_to_staging( + staging, + self._kv_caches[layer_name], + self._layer_idx_map[layer_name], + block_ids, + len(self._layer_names), + self._local_slot_size, + block_index, + ) + + async def _write_cuda_buffer(self, staged: "StagedStoreBatch") -> list[str]: + if staged.buffer.device.type == "cuda": + torch.cuda.set_device(staged.buffer.device) + cp_buffer = cupy.asarray(staged.buffer) + device_ptr = cuda_array_pointer(cp_buffer) + allocation_base, allocation_offset = cuda_allocation_base_and_offset(device_ptr) + return await self._client.transfer_store_cuda( + cuda_ipc_handle=export_cuda_ipc_handle(cp_buffer), + nbytes=staged.buffer.nbytes, + device_id=cuda_array_device_id(cp_buffer), + device_ptr=device_ptr, + allocation_base_ptr=allocation_base, + allocation_offset=allocation_offset, + producer_pid=os.getpid(), + spans=[ + { + "source_offset": span.source_offset, + "nbytes": span.nbytes, + "file_offset": span.file_offset, + "chunk_key": span.chunk_key, + "start_slot": span.start_slot, + "num_slots": span.num_slots, + } + for span in staged.spans + ], + ) + @dataclass(frozen=True) class StagedStoreBatch: """Hold a worker CUDA snapshot until its async store completes.""" buffer: torch.Tensor - ready_event: torch.cuda.Event | None spans: list[StoreWriteSpan] - lease: CudaStagingLease | None + lease: CudaStagingLease def build_staging_store_batches( diff --git a/daser/connector/worker_memory.py b/daser/connector/worker_memory.py index 2c0e008..5b9f64b 100644 --- a/daser/connector/worker_memory.py +++ b/daser/connector/worker_memory.py @@ -14,6 +14,7 @@ DEFAULT_STORE_STAGING_BYTES = 1536 << 20 DEFAULT_PENDING_STORE_STAGING_BYTES = 3072 << 20 MIN_STORE_STAGING_BYTES = 64 << 20 +MIN_STORE_STAGING_POOL_DEPTH = 1 class _CudaStagingLeaseOwner(Protocol): @@ -233,3 +234,67 @@ def derive_staging_limits(device: torch.device) -> tuple[int, int]: max(batch, min(total // 25, free // 5)), ) return batch, pending + + +def store_staging_pool_depth( + buffer_bytes: int, + pending_limit_bytes: int, +) -> int: + """Return fixed store staging pool depth for a 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 startup calculation. + """ + 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, + ) + + +def load_staging_pool_depth( + buffer_bytes: int, + pending_limit_bytes: int, + device: torch.device, + max_inflight: int, + reserve_bytes: int, +) -> int: + """Return fixed load staging depth under memory and inflight limits. + + Args: + buffer_bytes: Capacity of one fixed staging buffer. + pending_limit_bytes: Existing worker staging byte budget. + device: Device used for staging allocation. + max_inflight: Maximum request loads allowed in flight. + reserve_bytes: Free CUDA memory to leave unallocated. + + Returns: + Number of fixed load staging buffers to preallocate. + + Async/thread-safety: + Called during worker initialization before request traffic. It may query + CUDA free memory but does not allocate. + """ + depth = min( + 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) - reserve_bytes) + memory_depth = max(1, usable_bytes // buffer_bytes) + return max(1, min(depth, memory_depth)) diff --git a/daser/connector/worker_runtime.py b/daser/connector/worker_runtime.py index 14c8e6f..d49ec7d 100644 --- a/daser/connector/worker_runtime.py +++ b/daser/connector/worker_runtime.py @@ -3,16 +3,9 @@ from __future__ import annotations # Standard -import asyncio -from collections import deque -from dataclasses import dataclass, replace -import os -import threading -import time from typing import TYPE_CHECKING, Any # Third Party -import cupy import torch from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole @@ -22,45 +15,18 @@ from vllm.forward_context import ForwardContext # First Party -from daser.connector.helpers import base_req_id from daser.connector.ipc_client import IPCClientSync from daser.connector.load_pipeline import LoadPipeline -from daser.connector.load_pipeline import ( - build_load_copy_runs as _build_load_copy_runs, -) -from daser.connector.load_pipeline import ( - build_load_read_batches as _build_load_read_batches, -) from daser.connector.metadata import ( DaserConnectorMeta, - ReqStoreSpec, - StoreWriteSpan, ) from daser.connector.staging import CROSS_LAYER_KV_CACHE_KEY, FUSED_RESTORE_MIN_SLOTS -from daser.connector.staging import ( - copy_cross_layer_kv_cache_to_staging as _copy_cross_layer_kv_cache_to_staging, -) -from daser.connector.staging import ( - copy_kv_cache_to_staging as _copy_kv_cache_to_staging, -) -from daser.connector.staging import ( - copy_staging_to_kv_cache as _copy_staging_to_kv_cache, -) -from daser.connector.staging import ( - record_cuda_event as _record_cuda_event, -) -from daser.connector.staging import ( - synchronize_cuda_tensor as _synchronize_cuda_tensor, -) -from daser.connector.store_pipeline import StagedStoreBatch, StorePipeline -from daser.connector.store_pipeline import ( - build_staging_store_batches as _build_staging_store_batches, -) +from daser.connector.store_pipeline import StorePipeline from daser.connector.worker_memory import ( - DEFAULT_PENDING_STORE_STAGING_BYTES, DEFAULT_STORE_STAGING_BYTES, - CudaStagingLease, FixedCudaStagingPool, + load_staging_pool_depth, + store_staging_pool_depth, ) from daser.connector.worker_memory import ( derive_staging_limits as _derive_staging_limits, @@ -73,43 +39,12 @@ apply_rope_delta_to_kv_key_block_table, restore_cross_layer_kv_cache_table, ) -from daser.transfer.cuda_ipc import ( - cuda_array_device_id, - cuda_array_pointer, - export_cuda_ipc_handle, -) logger = init_logger(__name__) _ROPE_WARMUP_BLOCKS = 1 _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]] - - -def _rank_lane_offset( - start_slot: int, - local_slot_size: int, - rank_stride_bytes: int, - tp_rank: int, -) -> int: - """Return the physical offset for one rank-local logical slot range. - - Args: - start_slot: First server-owned logical slot. - local_slot_size: Bytes stored per logical slot by one TP rank. - rank_stride_bytes: Byte distance between adjacent rank lanes. - tp_rank: Current vLLM tensor-parallel rank. - - Returns: - Physical store offset for ``start_slot`` in ``tp_rank``'s lane. - - Async/thread-safety: - Pure arithmetic used on the worker thread before IPC submission. - """ - return tp_rank * rank_stride_bytes + start_slot * local_slot_size def _local_slot_bytes(connector: Any) -> int: @@ -159,379 +94,6 @@ def _validate_tp_layout( raise ValueError("DaseR runtime config is missing TP rank stride") -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) - - -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)) - - -@dataclass -class _DeferredFinishedSave: - """Store work held until vLLM reports a request as finished.""" - - commit_keys: set[str] - reqs_to_store: dict[str, ReqStoreSpec] - submitted: bool = False - future: Any | None = None - - -@dataclass -class _SaveFuture: - """One background save future and the staging it keeps alive. - - Attributes: - future: future returned by ``asyncio.run_coroutine_threadsafe``. - staging_bytes: GPU staging bytes held alive until completion. - lease: optional reusable staging lease released after completion. - """ - - future: Any - staging_bytes: int - lease: CudaStagingLease | None - - def release(self) -> None: - """Release the reusable staging lease, if any.""" - if self.lease is not None: - self.lease.release() - self.lease = None - - -@dataclass -class _PendingLoad: - """One background load future tracked until vLLM can resume the request. - - Attributes: - future: Future running the cross-layer load work. - block_ids: vLLM KV block IDs targeted by the load. - lease: optional staging lease held until the load completes. - """ - - future: Any - block_ids: list[int] - lease: CudaStagingLease | None - - def release(self) -> None: - """Release the reusable staging lease, if any.""" - if self.lease is not None: - self.lease.release() - 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[Any] - staging_lease: CudaStagingLease - future: Any - 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.""" - - 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. - - Args: - message: Error message raised when the future is collected. - - Async/thread-safety: - Immutable testable stand-in for a failed background future. It does not - spawn threads or perform IO. - """ - - def __init__(self, message: str) -> None: - self._message = message - - def done(self) -> bool: - """Return True because this failed future is already complete.""" - return True - - def result(self, timeout: float | None = None) -> None: - """Raise the seeded load-start failure. - - Args: - timeout: Ignored timeout for ``Future`` API compatibility. - """ - del timeout - 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() - - -@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 LoadRequestDispatcher: - """Bound request-level load concurrency by max in-flight and staging depth. - - Args: - max_inflight: Maximum request loads allowed in flight. - staging_depth: Number of fixed staging buffers available to request - loads. - - Async/thread-safety: - The dispatcher object is owned by the connector load asyncio loop. Test - helpers may call its pure synchronous scheduling helpers directly. - """ - - def __init__( - self, - max_inflight: int, - staging_depth: int, - ) -> None: - self._effective_inflight = max(1, min(max_inflight, staging_depth)) - self._free_buffers: deque[int] = deque(range(self._effective_inflight)) - - @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: - 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. - """ - 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 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. - - 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]: - """Return CUDA allocation base pointer and byte offset for ``device_ptr``. - - Args: - device_ptr: CUDA device pointer exported through IPC. - - Returns: - Tuple of ``(allocation_base_ptr, byte_offset)``. - """ - try: - from cuda.bindings import driver as cuda_driver - - result, base_ptr, _allocation_size = cuda_driver.cuMemGetAddressRange( - device_ptr - ) - if result == cuda_driver.CUresult.CUDA_SUCCESS: - base = int(base_ptr) - return base, int(device_ptr) - base - except Exception as exc: # noqa: BLE001 - logger.debug("[CONNECTOR] cuMemGetAddressRange failed: %s", exc) - return int(device_ptr), 0 - - def _warm_rope_apply_backends( device: torch.device, dtype: torch.dtype, @@ -699,13 +261,13 @@ def __init__( self._kv_cache_config = kv_cache_config self._role = KVConnectorRole.WORKER self._transfer_ready = False + self._pipelines_initialized = False self._load_pipeline = LoadPipeline(socket_path, _LOAD_REQUEST_MAX_INFLIGHT) self._store_pipeline = StorePipeline(socket_path) self._kv_caches: dict[str, torch.Tensor] = {} self._layer_names: list[str] = [] self._layer_idx_map: dict[str, int] = {} self._meta: DaserConnectorMeta | None = None - self._pending_commits: set[str] = set() def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: """Register the per-layer KV cache tensors. @@ -717,35 +279,26 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: self._layer_names = list(kv_caches.keys()) self._layer_idx_map = {name: idx for idx, name in enumerate(self._layer_names)} sample = next(iter(kv_caches.values()), None) - if kv_caches: - ( - self._store_pipeline.staging_bytes, - self._store_pipeline.pending_staging_limit_bytes, - ) = _derive_staging_limits(sample.device) + if sample is not None: logger.info( "[CONNECTOR] register_kv_caches: %d layers, first shape=%s dtype=%s", len(kv_caches), sample.shape, sample.dtype, ) - logger.info( - "[CONNECTOR] transient store staging caps: batch=%d pending=%d", - self._store_pipeline.staging_bytes, - self._store_pipeline.pending_staging_limit_bytes, - ) if self._layer_names and sample is not None: num_blocks = sample.shape[1] if sample.dim() >= 2 else 1 layer_size = sample.nbytes // num_blocks local_slot_size = layer_size * len(self._layer_names) - tp_size = int(getattr(self, "_tp_size", 1)) + tp_size = self._tp_size _validate_tp_layout( local_slot_size, self._slot_size, tp_size, - int(getattr(self, "_server_tp_size", tp_size)), - int(getattr(self, "_tp_rank", 0)), - int(getattr(self, "_rank_stride_bytes", 0)), + self._server_tp_size, + self._tp_rank, + self._rank_stride_bytes, ) self._local_slot_size = local_slot_size if self._slot_size == 0: @@ -757,39 +310,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: ) if sample is not None: - self._store_pipeline.staging_bytes = max( - self._store_pipeline.staging_bytes or DEFAULT_STORE_STAGING_BYTES, - self._local_slot_size, - ) - self._store_pipeline.staging_pool = FixedCudaStagingPool( - device=sample.device, - buffer_bytes=self._store_pipeline.staging_bytes, - depth=_store_staging_pool_depth( - self._store_pipeline.staging_bytes, - self._store_pipeline.pending_staging_limit_bytes, - ), - ) - self._load_pipeline.set_staging_pool( - FixedCudaStagingPool( - device=sample.device, - buffer_bytes=self._store_pipeline.staging_bytes, - depth=_load_staging_pool_depth( - self._store_pipeline.staging_bytes, - self._store_pipeline.pending_staging_limit_bytes, - sample.device, - ), - ) - ) - self._load_pipeline.staging_registered = False - logger.info( - "[CONNECTOR] preallocated staging buffer=%d cap=%d pending=%d " - "load_request_max_inflight=%d load_staging_depth=%d", - self._store_pipeline.staging_bytes, - self._store_pipeline.staging_bytes, - self._store_pipeline.pending_staging_limit_bytes, - _LOAD_REQUEST_MAX_INFLIGHT, - self._load_pipeline.staging_pool.depth, - ) + self._configure_pipelines(sample) if sample.dim() >= 5: _warm_rope_apply_backends( device=sample.device, @@ -797,9 +318,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: block_tokens=int(sample.shape[-3]), heads=int(sample.shape[-2]), head_dim=int(sample.shape[-1]), - rotary_dim=int(getattr(self, "_rope_rotary_dim", 0)), - rope_base=float(getattr(self, "_rope_base", 10000.0)), - is_neox_style=bool(getattr(self, "_rope_is_neox_style", True)), + rotary_dim=self._rope_rotary_dim, + rope_base=self._rope_base, + is_neox_style=self._rope_is_neox_style, ) self._init_server_transfer() @@ -820,7 +341,7 @@ def register_cross_layers_kv_cache( Async/thread-safety: Called once during worker initialization before request traffic. """ - kv_cache_config = getattr(self, "_kv_cache_config", None) + kv_cache_config = self._kv_cache_config layer_names: list[str] = [] if kv_cache_config is not None: for group in getattr(kv_cache_config, "kv_cache_groups", []): @@ -837,20 +358,16 @@ def register_cross_layers_kv_cache( tuple(kv_cache.shape), ) return - ( - self._store_pipeline.staging_bytes, - self._store_pipeline.pending_staging_limit_bytes, - ) = _derive_staging_limits(kv_cache.device) layer_size = kv_cache[0, 0].nbytes local_slot_size = layer_size * len(self._layer_names) - tp_size = int(getattr(self, "_tp_size", 1)) + tp_size = self._tp_size _validate_tp_layout( local_slot_size, self._slot_size, tp_size, - int(getattr(self, "_server_tp_size", tp_size)), - int(getattr(self, "_tp_rank", 0)), - int(getattr(self, "_rank_stride_bytes", 0)), + self._server_tp_size, + self._tp_rank, + self._rank_stride_bytes, ) self._local_slot_size = local_slot_size if self._slot_size == 0: @@ -860,30 +377,7 @@ def register_cross_layers_kv_cache( self._local_slot_size, len(self._layer_names), ) - self._store_pipeline.staging_bytes = max( - self._store_pipeline.staging_bytes or DEFAULT_STORE_STAGING_BYTES, - self._local_slot_size, - ) - self._store_pipeline.staging_pool = FixedCudaStagingPool( - device=kv_cache.device, - buffer_bytes=self._store_pipeline.staging_bytes, - depth=_store_staging_pool_depth( - self._store_pipeline.staging_bytes, - self._store_pipeline.pending_staging_limit_bytes, - ), - ) - self._load_pipeline.set_staging_pool( - FixedCudaStagingPool( - device=kv_cache.device, - buffer_bytes=self._store_pipeline.staging_bytes, - depth=_load_staging_pool_depth( - self._store_pipeline.staging_bytes, - self._store_pipeline.pending_staging_limit_bytes, - kv_cache.device, - ), - ) - ) - self._load_pipeline.staging_registered = False + load_staging_depth = self._configure_pipelines(kv_cache) logger.info( "[CONNECTOR] register_cross_layers_kv_cache: layers=%d shape=%s " "dtype=%s load_request_max_inflight=%d load_staging_depth=%d", @@ -891,7 +385,7 @@ def register_cross_layers_kv_cache( tuple(kv_cache.shape), kv_cache.dtype, _LOAD_REQUEST_MAX_INFLIGHT, - self._load_pipeline.staging_pool.depth, + load_staging_depth, ) _warm_rope_apply_backends( device=kv_cache.device, @@ -899,9 +393,9 @@ def register_cross_layers_kv_cache( block_tokens=int(kv_cache.shape[-3]), heads=int(kv_cache.shape[-2]), head_dim=int(kv_cache.shape[-1]), - rotary_dim=int(getattr(self, "_rope_rotary_dim", 0)), - rope_base=float(getattr(self, "_rope_base", 10000.0)), - is_neox_style=bool(getattr(self, "_rope_is_neox_style", True)), + rotary_dim=self._rope_rotary_dim, + rope_base=self._rope_base, + is_neox_style=self._rope_is_neox_style, ) _warm_cross_layer_restore_backends( device=kv_cache.device, @@ -910,9 +404,9 @@ def register_cross_layers_kv_cache( block_tokens=int(kv_cache.shape[-3]), heads=int(kv_cache.shape[-2]), head_dim=int(kv_cache.shape[-1]), - rotary_dim=int(getattr(self, "_rope_rotary_dim", 0)), - rope_base=float(getattr(self, "_rope_base", 10000.0)), - is_neox_style=bool(getattr(self, "_rope_is_neox_style", True)), + rotary_dim=self._rope_rotary_dim, + rope_base=self._rope_base, + is_neox_style=self._rope_is_neox_style, ) self._init_server_transfer() @@ -923,519 +417,44 @@ def bind_connector_metadata(self, connector_metadata: DaserConnectorMeta) -> Non connector_metadata: DaserConnectorMeta from build_connector_meta. """ self._meta = connector_metadata - self._reap_save_futures(block=False) - self._pending_commits = set() - for spec in connector_metadata.reqs_to_store.values(): - if spec.block_ids: - self._pending_commits.add(spec.chunk_key) def clear_connector_metadata(self) -> None: """Clear metadata after forward pass completes.""" self._meta = None def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: - """Submit async KV cache loads for cache-hit requests. + """Submit cache-hit requests to the load pipeline. Args: - forward_context: vLLM ForwardContext for this forward pass. + forward_context: vLLM forward context for this step. + **kwargs: Additional vLLM hook arguments, currently unused. + + Async/thread-safety: + Called on the vLLM worker thread. Transfer and restore execute on + the load pipeline thread. """ del forward_context, kwargs if self._meta is None or not self._meta.reqs_to_load: return - logger.debug( - "[CONNECTOR] start_load_kv: %d reqs to load", - len(self._meta.reqs_to_load), - ) reqs_to_load = dict(self._meta.reqs_to_load) if not self._ensure_transfer_ready(): - self._mark_load_start_failed( + self._load_pipeline.mark_failed( reqs_to_load, "server transfer config is not ready", ) return + self._load_pipeline.start(reqs_to_load) - num_layers = len(self._layer_names) - if num_layers == 0: - self._mark_load_start_failed(reqs_to_load, "no registered KV cache layers") - return - - sample_tensor = next(iter(self._kv_caches.values()), None) - if sample_tensor is None: - self._mark_load_start_failed(reqs_to_load, "no registered KV cache tensor") - return - - pending_loads = self._load_pipeline.pending - 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() - pending_loads[req_id] = _PendingLoad( - future=request_future, - block_ids=list(spec.block_ids), - lease=None, - ) - self._enqueue_load_request( - load_queue, - _QueuedLoadRequest( - req_id=req_id, - spec_id=spec_id, - spec=spec, - future=request_future, - ), - ) - self._ensure_load_request_dispatcher(sample_tensor) - - def _mark_load_start_failed( - self, - reqs_to_load: dict[str, Any], - reason: str, - ) -> None: - """Record failed load submission so vLLM can release waiting requests. - - Args: - reqs_to_load: Load metadata that could not be submitted. - reason: Human-readable failure reason used in diagnostics. - - Async/thread-safety: - Called on the vLLM worker thread before any background load is - started. Completion is later reported through ``get_finished``. - """ - if not reqs_to_load: - return - block_ids = [ - block_id for spec in reqs_to_load.values() for block_id in spec.block_ids - ] - failed_future = _ImmediateLoadError(reason) - pending_loads = self._load_pipeline.pending - for req_id in {base_req_id(req_id) for req_id in reqs_to_load}: - pending_loads[req_id] = _PendingLoad( - future=failed_future, - block_ids=list(block_ids), - lease=None, - ) - - 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: - 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. - """ - del sample_tensor - return self._load_pipeline.ensure_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. - """ - del load_queue - self._load_pipeline.enqueue(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 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. - """ - self._load_pipeline.ensure_dispatcher( - self._run_load_request_dispatcher(sample_tensor) - ) - - 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 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. - """ - if sample_tensor.device.type == "cuda": - torch.cuda.set_device(sample_tensor.device) - 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, - ) - 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: - try: - item = load_queue.get_nowait() - except asyncio.QueueEmpty: - return - 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 _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 = self._load_pipeline.staging_pool - if pool is not None: - return pool - buffer_bytes = max( - self._store_pipeline.staging_bytes or DEFAULT_STORE_STAGING_BYTES, - _local_slot_bytes(self), - ) - pool = FixedCudaStagingPool( - device=sample_tensor.device, - buffer_bytes=buffer_bytes, - depth=_load_staging_pool_depth( - buffer_bytes, - self._store_pipeline.pending_staging_limit_bytes, - sample_tensor.device, - ), - ) - self._load_pipeline.set_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 - - submitted_at = time.perf_counter() - if self._load_pipeline.staging_registered: - transfer_coro = self._transfer_load_registered_cuda( - buffer_index=buffer_index, - producer_pid=os.getpid(), - 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( - buffer_index=buffer_index, - 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, - ) - future = self._submit_load_coroutine(transfer_coro) - 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 _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) - spec = replace( - item.spec, - file_offset=_rank_lane_offset( - item.spec.start_slot, - _local_slot_bytes(self), - int(getattr(self, "_rank_stride_bytes", 0)), - int(getattr(self, "_tp_rank", 0)), - ), - ) - load_batches = _build_load_read_batches( - {item.spec_id: spec}, - _local_slot_bytes(self), - 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, - sample_tensor: torch.Tensor, - ) -> _ConsumedLoadBatch: - """Wait for one submitted load batch and restore it into vLLM KV cache. + def wait_for_layer_load(self, layer_name: str) -> None: + """Return after request-level load completion restored every layer. 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. + layer_name: Layer reported by vLLM; no per-layer wait is required. 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=_local_slot_bytes(self), - 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. - - Args: - layer_name: ignored. + Called on the vLLM worker thread. """ - return + del layer_name def save_kv_layer( self, @@ -1465,21 +484,12 @@ def wait_for_save(self) -> None: """Queue stores until vLLM reports request completion.""" if self._meta is None: return - - commit_keys = list(self._pending_commits) reqs_to_store = dict(self._meta.reqs_to_store) - if commit_keys and reqs_to_store: - pending_finished = self._store_pipeline.pending_finished_saves - for req_id, spec in reqs_to_store.items(): - base_id = base_req_id(req_id) - save = pending_finished.get(base_id) - if save is None: - save = _DeferredFinishedSave(commit_keys=set(), reqs_to_store={}) - pending_finished[base_id] = save - save.reqs_to_store[req_id] = spec - if spec.chunk_key in commit_keys: - save.commit_keys.add(spec.chunk_key) - self._pending_commits.clear() + commit_keys = { + spec.chunk_key for spec in reqs_to_store.values() if spec.block_ids + } + if commit_keys: + self._store_pipeline.queue_finished(reqs_to_store, commit_keys) def get_finished( self, finished_req_ids: set[str] @@ -1492,32 +502,8 @@ def get_finished( Returns: Finished-saving request IDs and finished async-loading request IDs. """ - self._reap_save_futures(block=False) - finished_recving = self._collect_finished_loads() - pending_finished = self._store_pipeline.pending_finished_saves - if not pending_finished: - return None, finished_recving or None - - finished_sending: set[str] = set() - candidates = set(finished_req_ids) - candidates.update( - req_id for req_id, save in pending_finished.items() if save.submitted - ) - for req_id in list(candidates): - save = pending_finished.get(req_id) - if save is None: - continue - if not save.submitted: - save.future = self._submit_finished_save(save) - save.submitted = True - future = save.future - if future is None: - finished_sending.add(req_id) - del pending_finished[req_id] - elif future.done(): - future.result(timeout=120.0) - finished_sending.add(req_id) - del pending_finished[req_id] + finished_recving = self._load_pipeline.collect_finished() + finished_sending = self._store_pipeline.collect_finished(finished_req_ids) return finished_sending or None, finished_recving or None def get_block_ids_with_load_errors(self) -> set[int]: @@ -1526,120 +512,50 @@ def get_block_ids_with_load_errors(self) -> set[int]: Returns: vLLM block IDs that should be treated as invalid. """ - invalid_blocks = self._load_pipeline.invalid_block_ids - invalid = set(invalid_blocks) - invalid_blocks.clear() - return invalid - - def _collect_finished_loads(self) -> set[str]: - """Poll async load futures without blocking. - - Returns: - Base request IDs whose async load future completed in this poll. - """ - pending_loads = self._load_pipeline.pending - if not pending_loads: - return set() - - finished_recving: set[str] = set() - collected_futures: set[int] = set() - for req_id, load in list(pending_loads.items()): - if not load.future.done(): - continue - future_id = id(load.future) - try: - if future_id not in collected_futures: - load.future.result() - collected_futures.add(future_id) - except Exception as exc: # noqa: BLE001 - logger.warning( - "[CONNECTOR] async load failed req=%s blocks=%s: %s", - req_id, - load.block_ids, - exc, - ) - self._load_pipeline.invalid_block_ids.update(load.block_ids) - collected_futures.add(future_id) - finally: - load.release() - del pending_loads[req_id] - finished_recving.add(req_id) - return finished_recving + return self._load_pipeline.take_invalid_block_ids() def shutdown(self) -> None: """Stop the background IO loop.""" if self._role != KVConnectorRole.WORKER: return - pending_loads = self._load_pipeline.pending - for load in {id(load.future): load for load in pending_loads.values()}.values(): - if not load.future.done(): - try: - load.future.result(timeout=120.0) - except Exception: # noqa: BLE001 - pass - self._collect_finished_loads() - for req_id in list(self._store_pipeline.pending_finished_saves): - self.get_finished({req_id}) - self._reap_save_futures(block=True) self._load_pipeline.shutdown() self._store_pipeline.shutdown() - def _submit_load_coroutine(self, coro: Any) -> Any: - """Submit foreground load work to the dedicated load event loop. - - Args: - coro: Coroutine object to schedule. - - Returns: - Future returned by ``asyncio.run_coroutine_threadsafe``. - - Async/thread-safety: - Called from vLLM worker threads. A load-only loop prevents cache-hit - reads from queueing behind background store coroutines. - """ - return self._load_pipeline.submit(coro) - - def _submit_store_coroutine(self, coro: Any) -> Any: - """Submit background store and commit work to the store event loop. - - Args: - coro: Coroutine object to schedule. - - Returns: - Future returned by ``asyncio.run_coroutine_threadsafe``. - - Async/thread-safety: - Called from vLLM worker threads. Store work is serialized on the - store loop and does not occupy the foreground load loop. - """ - return self._store_pipeline.submit(coro) - def _ensure_transfer_ready(self) -> bool: - """Refresh server transfer config and mark worker data plane ready.""" - if getattr(self, "_transfer_ready", False): - return True + """Refresh config and initialize both pipeline transfer clients.""" + if not self._transfer_ready: + self._refresh_runtime_config() + if not self._slot_size or (not self._skip_l2 and not self._store_path): + logger.warning( + "[CONNECTOR] server transfer config is not ready; start DaseR " + "server before sending requests", + ) + return False - self._refresh_runtime_config() - if not self._slot_size or ( - not getattr(self, "_skip_l2", False) and not self._store_path - ): - logger.warning( - "[CONNECTOR] server transfer config is not ready; start DaseR server " - "before sending requests", + _validate_tp_layout( + _local_slot_bytes(self), + self._slot_size, + self._tp_size, + self._server_tp_size, + self._tp_rank, + self._rank_stride_bytes, ) - return False - - _validate_tp_layout( - _local_slot_bytes(self), - self._slot_size, - int(getattr(self, "_tp_size", 1)), - int(getattr(self, "_server_tp_size", 1)), - int(getattr(self, "_tp_rank", 0)), - int(getattr(self, "_rank_stride_bytes", 0)), - ) + self._load_pipeline.configure_rank_geometry( + self._rank_stride_bytes, + self._tp_rank, + ) + self._store_pipeline.configure_rank_geometry( + self._rank_stride_bytes, + self._tp_rank, + self._tp_size, + ) + self._transfer_ready = True + logger.info("[CONNECTOR] server transfer mode=%s", self._transfer_mode) - self._transfer_ready = True - logger.info("[CONNECTOR] server transfer mode=%s", self._transfer_mode) + if not self._pipelines_initialized: + self._load_pipeline.initialize_transfer() + self._store_pipeline.initialize_transfer() + self._pipelines_initialized = True return True def _refresh_runtime_config(self) -> None: @@ -1664,493 +580,78 @@ def _refresh_runtime_config(self) -> None: self._transfer_mode = str(config.get("transfer_mode", self._transfer_mode)) def _init_server_transfer(self) -> None: - """Initialize the server-owned transfer layer on both IO loops. - - Async/thread-safety: - Called from a vLLM worker thread after KV-cache registration. Does - nothing until the server transfer config and both async IO loops - are ready. - """ - if not self._ensure_transfer_ready(): - return - 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._store_pipeline.client.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. - """ - return list(self._load_pipeline.clients()) - - 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. - """ - return self._load_pipeline.client(buffer_index) - - def _register_load_staging_buffers(self) -> None: - """Register fixed load staging buffers with the server. + """Initialize both pipeline-owned transfer clients. 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. + Called on the worker thread after KV-cache registration. Each + pipeline performs its initialization on its private event loop. """ - if self._load_pipeline.staging_registered: - return - pool = self._load_pipeline.staging_pool - if pool is None or not self._load_ipc_clients(): - return - try: - for buffer_index in range(pool.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) - load_client = self._load_ipc_client_for_buffer(buffer_index) - self._submit_load_coroutine( - load_client.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_pipeline.staging_registered = False - return - self._load_pipeline.staging_registered = True - logger.info( - "[CONNECTOR] registered %d fixed load staging buffers", - pool.depth, - ) - - def _reap_save_futures(self, block: bool) -> None: - """Collect completed background save tasks. + self._ensure_transfer_ready() - Args: - block: If True, wait for every pending save. If False, collect only - tasks that are already complete. - """ - remaining: list[_SaveFuture] = [] - pending_bytes = self._store_pipeline.pending_staging_bytes - for record in self._store_pipeline.save_futures: - if block or record.future.done(): - try: - record.future.result(timeout=120.0) - finally: - pending_bytes = max(0, pending_bytes - record.staging_bytes) - record.release() - else: - remaining.append(record) - self._store_pipeline.save_futures = remaining - self._store_pipeline.pending_staging_bytes = pending_bytes - - def _track_save_future( - self, - future: Any, - staging_bytes: int, - staging_lease: CudaStagingLease | None, - ) -> None: - """Track one background save future and its live staging bytes. + def _configure_pipelines(self, sample: torch.Tensor) -> int: + """Configure load and store pipelines from one finalized KV layout. Args: - future: Future returned by ``asyncio.run_coroutine_threadsafe``. - staging_bytes: GPU staging bytes kept alive by the future. - staging_lease: Optional reusable staging lease to release after - ``future`` completes. - - Async/thread-safety: - Called on the worker thread. Completion is collected by - ``_reap_save_futures``. - """ - self._store_pipeline.pending_staging_bytes += staging_bytes - self._store_pipeline.save_futures.append( - _SaveFuture(future=future, staging_bytes=staging_bytes, lease=staging_lease) - ) - - def _wait_for_save_staging_capacity(self, nbytes: int) -> None: - """Apply backpressure before allocating another store staging buffer. - - Args: - nbytes: Size of the next staging tensor. - - Async/thread-safety: - Called by vLLM's worker thread. It may wait for already-submitted - background stores when live staging would exceed the configured - cap. - """ - limit = max( - self._store_pipeline.pending_staging_limit_bytes - or DEFAULT_PENDING_STORE_STAGING_BYTES, - nbytes, - ) - while ( - self._store_pipeline.pending_staging_bytes + nbytes > limit - and self._store_pipeline.save_futures - ): - record = self._store_pipeline.save_futures.pop(0) - try: - record.future.result(timeout=120.0) - finally: - self._store_pipeline.pending_staging_bytes = max( - 0, - self._store_pipeline.pending_staging_bytes - record.staging_bytes, - ) - 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 oldest store futures until a lease is released. - """ - pool = self._store_pipeline.staging_pool - self._wait_for_save_staging_capacity(nbytes) - while ( - pool is not None - and pool.available == 0 - and self._store_pipeline.save_futures - ): - record = self._store_pipeline.save_futures.pop(0) - try: - record.future.result(timeout=120.0) - finally: - self._store_pipeline.pending_staging_bytes = max( - 0, - self._store_pipeline.pending_staging_bytes - record.staging_bytes, - ) - record.release() - self._reap_save_futures(block=False) - - def _acquire_staging( - self, - nbytes: int, - device: torch.device, - ) -> CudaStagingLease: - """Acquire a reusable staging buffer for a CUDA IPC transfer. - - Args: - nbytes: Logical byte count needed for the transfer. - device: Device used when the pool has not been initialized yet. - - Returns: - A staging lease whose ``view`` is safe to export through CUDA IPC. - - Async/thread-safety: - Called from the worker thread. Store-path callers must retain the - lease until the background server transfer completes. - """ - pool = self._store_pipeline.staging_pool - if pool is None: - max_bytes = max( - nbytes, - self._store_pipeline.staging_bytes or DEFAULT_STORE_STAGING_BYTES, - ) - pool = FixedCudaStagingPool( - device=device, - buffer_bytes=max_bytes, - depth=_store_staging_pool_depth( - max_bytes, - self._store_pipeline.pending_staging_limit_bytes - or DEFAULT_PENDING_STORE_STAGING_BYTES, - ), - ) - self._store_pipeline.staging_pool = pool - return pool.acquire( - nbytes, - wait_for_release=lambda: self._wait_for_store_staging_release(nbytes), - ) - - def _stage_store_batch( - self, - block_ids: list[int], - spans: list[StoreWriteSpan], - ) -> StagedStoreBatch | None: - """Snapshot one bounded batch of KV blocks into CUDA staging. - - Args: - block_ids: vLLM KV block IDs to snapshot. - spans: Server store spans targeting this staging batch. + sample: Representative registered KV-cache tensor. Returns: - A staged batch ready for CUDA IPC transfer, or ``None`` when the - connector has no layer state. + Number of preallocated load staging buffers. Async/thread-safety: - Runs on the vLLM worker thread so KV cache reads are launched before - vLLM can recycle the source blocks. The returned tensor is kept - alive by the background transfer future. + Called once on the worker thread during KV-cache registration. """ - num_layers = len(self._layer_names) - if num_layers == 0: - return None - sample_tensor = next(iter(self._kv_caches.values()), None) - if sample_tensor is None: - return None - if not block_ids or not spans: - return None - local_slot_size = _local_slot_bytes(self) - nbytes = len(block_ids) * local_slot_size - self._wait_for_save_staging_capacity(nbytes) - staging_lease = self._acquire_staging(nbytes, sample_tensor.device) - staging = staging_lease.view - block_index = torch.tensor( - block_ids, - dtype=torch.long, - device=sample_tensor.device, - ) - cross_layer_kv_cache = self._kv_caches.get(CROSS_LAYER_KV_CACHE_KEY) - if cross_layer_kv_cache is not None: - _copy_cross_layer_kv_cache_to_staging( - staging=staging, - kv_cache=cross_layer_kv_cache, - block_ids=block_ids, - num_layers=num_layers, - slot_size=local_slot_size, - block_index=block_index, - ) - else: - for layer_name in self._layer_names: - _copy_kv_cache_to_staging( - staging=staging, - kv_layer=self._kv_caches[layer_name], - layer_idx=self._layer_idx_map[layer_name], - block_ids=block_ids, - num_layers=num_layers, - slot_size=local_slot_size, - block_index=block_index, - ) - return StagedStoreBatch( - buffer=staging, - ready_event=_record_cuda_event(staging), - spans=spans, - lease=staging_lease, + staging_bytes, pending_limit_bytes = _derive_staging_limits(sample.device) + staging_bytes = max( + staging_bytes or DEFAULT_STORE_STAGING_BYTES, + self._local_slot_size, ) - - def _submit_finished_save(self, save: _DeferredFinishedSave) -> Any | None: - """Submit one request's deferred KV store after request completion. - - Args: - save: Deferred store plan built during ``wait_for_save``. - - Returns: - Future that completes after all store batches are committed, or - ``None`` when no store batch could be staged. - - Async/thread-safety: - Called by vLLM's worker thread from ``get_finished`` while vLLM is - still holding the finished request's KV blocks. - """ - batch_futures = [] - reqs_to_store = { - req_id: replace( - spec, - file_offset=_rank_lane_offset( - spec.start_slot, - _local_slot_bytes(self), - int(getattr(self, "_rank_stride_bytes", 0)), - int(getattr(self, "_tp_rank", 0)), - ), - ) - for req_id, spec in save.reqs_to_store.items() - } - batches = _build_staging_store_batches( - reqs_to_store, - _local_slot_bytes(self), - max_batch_bytes=( - self._store_pipeline.staging_bytes or DEFAULT_STORE_STAGING_BYTES + store_pool = FixedCudaStagingPool( + device=sample.device, + buffer_bytes=staging_bytes, + depth=store_staging_pool_depth(staging_bytes, pending_limit_bytes), + ) + load_pool = FixedCudaStagingPool( + device=sample.device, + buffer_bytes=staging_bytes, + depth=load_staging_pool_depth( + staging_bytes, + pending_limit_bytes, + sample.device, + _LOAD_REQUEST_MAX_INFLIGHT, + _LOAD_STAGING_RESERVE_BYTES, ), ) - for block_ids, spans in batches: - staged = self._stage_store_batch(block_ids, spans) - if staged is None: - continue - future = self._submit_store_coroutine( - self._write_cuda_buffer( - buffer=staged.buffer, - ready_event=staged.ready_event, - spans=staged.spans, - ) - ) - self._track_save_future(future, staged.buffer.nbytes, staged.lease) - batch_futures.append(future) - if not batch_futures: - return None - commit_future = self._submit_store_coroutine( - self._commit_after_store_futures(batch_futures, sorted(save.commit_keys)), + self._store_pipeline.configure( + kv_caches=self._kv_caches, + layer_names=self._layer_names, + layer_idx_map=self._layer_idx_map, + local_slot_size=self._local_slot_size, + rank_stride_bytes=self._rank_stride_bytes, + tp_rank=self._tp_rank, + tp_size=self._tp_size, + staging_bytes=staging_bytes, + staging_pool=store_pool, + ) + self._load_pipeline.configure( + kv_caches=self._kv_caches, + layer_names=self._layer_names, + local_slot_size=self._local_slot_size, + rank_stride_bytes=self._rank_stride_bytes, + tp_rank=self._tp_rank, + staging_pool=load_pool, + load_key_scale=self._load_key_scale, + load_value_scale=self._load_value_scale, + 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, ) - self._track_save_future(commit_future, 0, None) - return commit_future - - async def _commit_after_store_futures( - self, - batch_futures: list[Any], - commit_keys: list[str], - ) -> None: - """Commit chunks after all staged transfer batches finish. - - Args: - batch_futures: Futures for each staged store batch. - commit_keys: Chunk keys to publish after stores complete. - - Async/thread-safety: - Runs on the connector background event loop and does not read vLLM - KV cache tensors. - """ - stored_keys: list[str] = [] - for future in batch_futures: - stored_keys.extend(await asyncio.wrap_future(future)) - await self._commit_stored_keys(stored_keys, commit_keys) - - async def _write_cuda_buffer( - self, - buffer: torch.Tensor, - ready_event: torch.cuda.Event | None, - spans: list[StoreWriteSpan], - ) -> list[str]: - """Write selected spans from one contiguous CUDA buffer. - - Args: - buffer: CUDA tensor exported over CUDA IPC. - ready_event: Producer-stream event for ``buffer``. - spans: Source/destination write spans. - - Returns: - Chunk keys accepted by the server for this buffer. - """ - if buffer.device.type == "cuda": - torch.cuda.set_device(buffer.device) - if ready_event is not None: - ready_event.synchronize() - else: - _synchronize_cuda_tensor(buffer) - cp_buffer = cupy.asarray(buffer) - cuda_ipc_handle = export_cuda_ipc_handle(cp_buffer) - device_id = cuda_array_device_id(cp_buffer) - device_ptr = cuda_array_pointer(cp_buffer) - ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) - stored_keys = await self._transfer_store_cuda( - cuda_ipc_handle=cuda_ipc_handle, - nbytes=buffer.nbytes, - device_id=device_id, - device_ptr=device_ptr, - allocation_base_ptr=ipc_base_ptr, - allocation_offset=ipc_offset, - producer_pid=os.getpid(), - spans=[ - { - "source_offset": span.source_offset, - "nbytes": span.nbytes, - "file_offset": span.file_offset, - "chunk_key": span.chunk_key, - "start_slot": span.start_slot, - "num_slots": span.num_slots, - } - for span in spans - ], - ) - return stored_keys - - async def _commit_stored_keys( - self, - stored_keys: list[str], - commit_keys: list[str], - ) -> None: - """Commit requested chunks whose store spans were accepted.""" - requested = set(commit_keys) - candidate_keys = [key for key in stored_keys if key in requested] - keys_to_commit = list(dict.fromkeys(candidate_keys)) - await self._store_pipeline.client.commit_chunks( - keys_to_commit, - tp_rank=int(getattr(self, "_tp_rank", 0)), - tp_size=int(getattr(self, "_tp_size", 1)), - ) - - async def _transfer_load_cuda(self, **kwargs: Any) -> dict[str, Any]: - """Load through the dedicated worker load IPC client. - - Args: - **kwargs: forwarded CUDA transfer payload fields. - - Returns: - Server load response with timing counters. - - Async/thread-safety: - Runs on the worker load event loop. A dedicated client keeps - cache-hit loads from queueing behind store RPCs. - """ - 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. - - 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. - """ - 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. - - Args: - **kwargs: forwarded CUDA transfer payload fields. - - Returns: - Chunk keys accepted by the server. - - Async/thread-safety: - Runs on the worker store event loop and serializes only with other - store/commit traffic. - """ - return await self._store_pipeline.client.transfer_store_cuda(**kwargs) + logger.info( + "[CONNECTOR] preallocated staging bytes=%d pending=%d " + "store_depth=%d load_depth=%d", + staging_bytes, + pending_limit_bytes, + store_pool.depth, + load_pool.depth, + ) + return load_pool.depth diff --git a/daser/transfer/cuda_ipc.py b/daser/transfer/cuda_ipc.py index 26868b0..0c5e8dd 100644 --- a/daser/transfer/cuda_ipc.py +++ b/daser/transfer/cuda_ipc.py @@ -110,3 +110,30 @@ def cuda_array_device_id(array: Any) -> int: CUDA device ordinal. """ return int(array.device.id) + + +def cuda_allocation_base_and_offset(device_ptr: int) -> tuple[int, int]: + """Return the CUDA allocation base and byte offset for a tensor pointer. + + Args: + device_ptr: CUDA device pointer exported through IPC. + + Returns: + Tuple of allocation base pointer and byte offset. When the CUDA driver + query is unavailable, the pointer itself is used as the base. + + Async/thread-safety: + Read-only CUDA driver query safe during worker transfer preparation. + """ + try: + from cuda.bindings import driver as cuda_driver + + result, base_ptr, _allocation_size = cuda_driver.cuMemGetAddressRange( + device_ptr + ) + if result == cuda_driver.CUresult.CUDA_SUCCESS: + base = int(base_ptr) + return base, int(device_ptr) - base + except Exception: # noqa: BLE001 + pass + return int(device_ptr), 0 diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index 2b46b90..955f605 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Standard -import asyncio from types import SimpleNamespace # Third Party @@ -36,7 +35,7 @@ ) from daser.connector.request_lifecycle import RequestLifecycle from daser.connector.reuse import PrefixReuseStrategy -from daser.connector.scheduler import ( +from daser.connector.scheduler_planning import ( _block_ids_for_chunk, _contiguous_prefix_tokens, _trim_chunk_to_external_window, @@ -56,6 +55,10 @@ from daser.connector.staging import ( synchronize_cuda_tensor as _synchronize_cuda_tensor, ) +from daser.connector.store_pipeline import ( + StagedStoreBatch, + StorePipeline, +) from daser.connector.store_pipeline import ( build_staging_store_batches as _build_staging_store_batches, ) @@ -64,21 +67,12 @@ DEFAULT_STORE_STAGING_BYTES, MIN_STORE_STAGING_BYTES, FixedCudaStagingPool, + load_staging_pool_depth, ) from daser.connector.worker_memory import ( derive_staging_limits as _derive_staging_limits, ) -from daser.connector.worker_runtime import ( - LoadRequestDispatcher, - WorkerRuntime, - _DeferredFinishedSave, - _InflightRequestLoad, - _load_staging_pool_depth, - _PendingLoad, - _rank_lane_offset, - _RequestLoadFuture, - _SaveFuture, -) +from daser.connector.worker_runtime import WorkerRuntime BLOCK_TOKENS = 4 NUM_LAYERS = 2 @@ -107,137 +101,6 @@ def test_rolling_prefix_keys_match_single_step_helper() -> None: ) -class _RuntimeConfigProbe(DaserConnector): - """Test connector exposing runtime config state through public properties.""" - - @property - def runtime_state(self): - return ( - self._store_path, - self._slot_size, - self._block_tokens, - self._model_id, - ) - - -class _LoadPipelineProbe: - """In-memory load pipeline implementing the worker-facing interface.""" - - def __init__( - self, - client: object | None = None, - clients: list[object] | None = None, - ) -> None: - self.pending: dict[str, object] = {} - self.invalid_block_ids: set[int] = set() - self.staging_pool = None - self.staging_registered = False - self._clients = clients if clients is not None else ([client] if client else []) - self.queue: asyncio.Queue[object] = asyncio.Queue() - self.shutdown_called = False - self.submit_calls = 0 - - def set_staging_pool(self, pool: object) -> None: - self.staging_pool = pool - - def clients(self) -> tuple[object, ...]: - return tuple(self._clients) - - def client(self, buffer_index: int | None = None) -> object: - index = 0 if buffer_index is None else buffer_index - return self._clients[index % len(self._clients)] - - def ensure_queue(self) -> asyncio.Queue[object]: - return self.queue - - def enqueue(self, item: object) -> None: - self.queue.put_nowait(item) - - def ensure_dispatcher(self, coro: object) -> None: - coro.close() - - def submit(self, coro: object) -> object: - self.submit_calls += 1 - coro.close() - return SimpleNamespace(result=lambda timeout=None: None) - - def shutdown(self) -> None: - self.shutdown_called = True - - -class _StorePipelineProbe: - """In-memory store pipeline implementing the worker-facing interface.""" - - def __init__(self, client: object | None = None) -> None: - self.client = client - self.save_futures: list[object] = [] - self.pending_staging_bytes = 0 - self.staging_bytes = 0 - self.pending_staging_limit_bytes = 0 - self.staging_pool = None - self.pending_finished_saves: dict[str, object] = {} - self.shutdown_called = False - self.submit_calls = 0 - - def submit(self, coro: object) -> object: - self.submit_calls += 1 - coro.close() - return SimpleNamespace(result=lambda timeout=None: None) - - def shutdown(self) -> None: - self.shutdown_called = True - - -class _WorkerProbe(WorkerRuntime): - """Worker-side probe with minimal state for transfer readiness tests.""" - - def __init__(self, store_path: str) -> None: - self._meta = DaserConnectorMeta( - reqs_to_load={ - "req": ReqLoadSpec( - chunk_key="hit", - start_slot=0, - num_slots=1, - block_ids=[0], - file_offset=0, - token_count=BLOCK_TOKENS, - ) - } - ) - self._transfer_ready = False - self._store_path = store_path - self._slot_size = 1024 - self._block_tokens = 4 - self._layer_names = [] - self._transfer_mode = "gds" - self._skip_l2 = False - self._load_pipeline = _LoadPipelineProbe(self) - self._store_pipeline = _StorePipelineProbe(self) - - def _refresh_runtime_config(self) -> None: - return - - def _ensure_load_request_dispatcher(self, sample_tensor: torch.Tensor) -> None: - """Drain queued requests synchronously for worker probe tests.""" - dispatcher = LoadRequestDispatcher(max_inflight=8, staging_depth=1) - queued = [] - while not self._load_pipeline.queue.empty(): - item = self._load_pipeline.queue.get_nowait() - if item is None: - return - queued.append(item) - active = dispatcher.submit_ready(self, queued, sample_tensor) - while active: - consumed = dispatcher.consume_ready(self, active, sample_tensor) - if not consumed: - return - active.extend(dispatcher.submit_ready(self, queued, sample_tensor)) - - @property - def transfer_ready(self): - return self._transfer_ready - - class _SchedulerProbe(RequestLifecycle): """Scheduler-side probe that can emulate deferred runtime config.""" @@ -409,298 +272,6 @@ def alloc_chunk( return self._owner.alloc_chunk(chunk_key, token_count, model_id) -class _CommitProbe(WorkerRuntime): - """Minimal worker probe exposing commit filtering behavior.""" - - def __init__(self) -> None: - self.committed: list[list[str]] = [] - self.commit_metadata: list[tuple[int, int]] = [] - self._store_pipeline = _StorePipelineProbe(self) - - async def commit_chunks( - self, chunk_keys: list[str], tp_rank: int = 0, tp_size: int = 1 - ) -> None: - """Record chunk keys submitted to the async IPC client.""" - self.committed.append(list(chunk_keys)) - self.commit_metadata.append((tp_rank, tp_size)) - - async def commit_stored_keys( - self, stored_keys: list[str], commit_keys: list[str] - ) -> None: - """Expose worker commit filtering through a public test helper.""" - await self._commit_stored_keys(stored_keys, commit_keys) - - -class _QueueProbe(WorkerRuntime): - """Minimal worker probe exposing independent load/store IPC clients.""" - - def __init__(self) -> None: - self.load_calls: list[str] = [] - self.store_calls: list[str] = [] - self._load_pipeline = _LoadPipelineProbe(self) - self._store_pipeline = _StorePipelineProbe(self) - - async def transfer_load_cuda(self, **_kwargs) -> dict[str, int]: - """Record load client usage.""" - self.load_calls.append("load") - return {} - - async def transfer_store_cuda(self, **_kwargs) -> list[str]: - """Record store client usage.""" - self.store_calls.append("store") - return [] - - async def load_via_public_helper(self) -> None: - """Issue one load through the worker helper under test.""" - await self._transfer_load_cuda() - - async def store_via_public_helper(self) -> None: - """Issue one store through the worker helper under test.""" - await self._transfer_store_cuda() - - -class _FinishedSaveProbe(WorkerRuntime): - """Worker probe for finished-request save scheduling.""" - - def __init__(self) -> None: - self._meta = None - self._pending_commits = set() - self._load_pipeline = _LoadPipelineProbe() - self._store_pipeline = _StorePipelineProbe() - self._slot_size = 32 - self._store_pipeline.staging_bytes = 128 - self._store_pipeline.pending_staging_limit_bytes = 128 - self._layer_names = ["layer.0"] - self._kv_caches = {"layer.0": torch.empty(1)} - self.staged_batches: list[tuple[list[int], list[StoreWriteSpan]]] = [] - self.submitted = 0 - self.tracked: list[tuple[int, object | None]] = [] - self.committed_after: list[tuple[int, list[str]]] = [] - - def _clear_save_state(self) -> None: - return - - def _reap_save_futures(self, block: bool) -> None: - if block: - for future, _bytes, lease in self._store_pipeline.save_futures: - future.result(timeout=120.0) - if lease is not None: - lease.release() - self._store_pipeline.save_futures = [] - - def _stage_store_batch( - self, - block_ids: list[int], - spans: list[StoreWriteSpan], - ): - self.staged_batches.append((list(block_ids), list(spans))) - - class _Staged: - buffer = torch.empty(len(block_ids) * 32, dtype=torch.uint8) - ready_event = None - lease = object() - - def __init__(self, spans): - self.spans = spans - - return _Staged(spans) - - def _submit_store_coroutine(self, coro): - self.submitted += 1 - if getattr(getattr(coro, "cr_code", None), "co_name", "") == ( - "_commit_after_store_futures" - ): - asyncio.get_event_loop().run_until_complete(coro) - else: - coro.close() - - class _Future: - def __init__(self, value): - self._value = value - - def done(self) -> bool: - return True - - def result(self, timeout: float): - del timeout - return self._value - - return _Future(["stored"]) - - def _track_save_future( - self, - future, - staging_bytes: int, - staging_lease, - ) -> None: - self.tracked.append((staging_bytes, staging_lease)) - self._store_pipeline.save_futures.append((future, staging_bytes, staging_lease)) - - async def _commit_after_store_futures(self, batch_futures, commit_keys): - self.committed_after.append((len(batch_futures), list(commit_keys))) - - def set_pending_meta( - self, - reqs_to_store: dict[str, ReqStoreSpec], - commit_keys: set[str], - ) -> None: - """Seed worker metadata and commit keys for save tests.""" - self._meta = DaserConnectorMeta(reqs_to_store=reqs_to_store) - self._pending_commits = commit_keys - - def seed_finished_save( - self, - req_id: str, - reqs_to_store: dict[str, ReqStoreSpec], - commit_keys: list[str], - ) -> None: - """Seed a deferred finished save for worker completion tests.""" - self._store_pipeline.pending_finished_saves[req_id] = _DeferredFinishedSave( - commit_keys=set(commit_keys), - reqs_to_store=reqs_to_store, - ) - - def pending_finished_save_ids(self) -> set[str]: - """Return request IDs with deferred save work.""" - return set(self._store_pipeline.pending_finished_saves) - - def pending_commit_keys(self) -> set[str]: - """Return pending worker commit keys.""" - return set(self._pending_commits) - - def set_submit_store_coroutine(self, submitter) -> None: - """Replace the store coroutine submitter for worker tests.""" - self._submit_store_coroutine = submitter - - def disable_store_staging(self) -> None: - """Make staging fail for store lifecycle regression tests.""" - self._stage_store_batch = lambda _block_ids, _spans: None - - -class _AsyncLoadFuture: - """Controllable future for worker async-load completion tests.""" - - def __init__(self, *, done: bool, error: BaseException | None = None) -> None: - self._done = done - self._error = error - self.result_calls = 0 - - def done(self) -> bool: - """Return whether this fake load has completed.""" - return self._done - - def result(self, timeout: float | None = None) -> None: - """Record result collection and optionally raise the seeded error.""" - del timeout - self.result_calls += 1 - 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(WorkerRuntime): - """Worker probe with pending async load futures.""" - - def __init__(self) -> None: - self._role = KVConnectorRole.WORKER - self._load_pipeline = _LoadPipelineProbe() - self._store_pipeline = _StorePipelineProbe() - self.released: list[object] = [] - self.load_loop_stopped = False - self.store_loop_stopped = False - self.load_thread_joined = False - self.store_thread_joined = False - - class _Loop: - def __init__(self, callback) -> None: - self._callback = callback - - def call_soon_threadsafe(self, callback) -> None: - del callback - self._callback() - - def stop(self) -> None: - return - - class _Thread: - def __init__(self, callback) -> None: - self._callback = callback - - def join(self, timeout: float | None = None) -> None: - del timeout - self._callback() - - del _Loop, _Thread - - def _reap_save_futures(self, block: bool) -> None: - del block - - def get_finished(self, finished_req_ids: set[str]): - """Expose worker completion polling without shutdown test doubles.""" - return WorkerRuntime.get_finished(self, finished_req_ids) - - def _submit_load_coroutine(self, coro): - coro.close() - - class _Done: - def result(self, timeout: float | None = None) -> None: - del timeout - - return _Done() - - def _submit_store_coroutine(self, coro): - coro.close() - - class _Done: - def result(self, timeout: float | None = None) -> None: - del timeout - - return _Done() - - def seed_load( - self, - req_id: str, - future: _AsyncLoadFuture, - *, - block_ids: list[int], - lease: object | None = None, - ) -> None: - """Seed one pending async load.""" - self._load_pipeline.pending[req_id] = _PendingLoad( - future=future, - block_ids=block_ids, - lease=lease, - ) - - -class _LoopProbe(WorkerRuntime): - """Minimal worker probe exposing background loop selection.""" - - def __init__(self) -> None: - self._load_pipeline = _LoadPipelineProbe() - self._store_pipeline = _StorePipelineProbe() - - @property - def loop_pair(self) -> tuple[object, object]: - """Return the load and store loops configured on this probe.""" - return self._load_pipeline, self._store_pipeline - - def submit_load(self) -> None: - """Submit a load coroutine through the worker helper.""" - self._submit_load_coroutine(self._noop()) - - def submit_store(self) -> None: - """Submit a store coroutine through the worker helper.""" - self._submit_store_coroutine(self._noop()) - - async def _noop(self) -> None: - """Return immediately for loop-selection tests.""" - return - - def test_dataclasses_instantiate(): """DaserConnectorMeta, ReqLoadSpec, ReqStoreSpec all instantiate cleanly.""" spec_load = ReqLoadSpec("k", 0, 1, [0], 0, 16, 0, 0) @@ -714,55 +285,6 @@ def test_dataclasses_instantiate(): assert spec_load.pos_offset == 0 -def test_connector_allows_runtime_config_from_ipc(monkeypatch, tmp_path): - """Worker startup can begin with socket_path only and fill config by IPC.""" - - class DummyIPCClient: - def __init__(self, socket_path): - self.socket_path = socket_path - - def get_runtime_config(self): - return { - "store_path": str(tmp_path / "daser.store"), - "slot_size": 1024, - "block_tokens": 4, - "model_id": "served-model", - "cache_reuse_mode": "prefix", - } - - class DummyBase: - def __init__(self, vllm_config, role, kv_cache_config=None): - self._role = role - - class DummyConfig: - kv_connector_extra_config = {"socket_path": "/tmp/daser.sock"} - - class DummyVLLMConfig: - kv_transfer_config = DummyConfig() - model_config = None - - monkeypatch.setattr( - "daser.connector.daser_connector.IPCClientSync", - DummyIPCClient, - ) - monkeypatch.setattr( - "daser.connector.daser_connector.KVConnectorBase_V1.__init__", - DummyBase.__init__, - ) - - connector = _RuntimeConfigProbe( - DummyVLLMConfig(), - role=KVConnectorRole.SCHEDULER, - ) - - assert connector.runtime_state == ( - str(tmp_path / "daser.store"), - 1024, - 4, - "served-model", - ) - - def test_connector_requests_cross_layer_nhd_layout() -> None: """DaseR asks vLLM for block-major cross-layer KV cache layout.""" assert DaserConnector.get_required_kvcache_layout(object()) == "NHD" @@ -810,408 +332,12 @@ class DummyVLLMConfig: role=KVConnectorRole.SCHEDULER, ) - assert isinstance(connector._cache_reuse_strategy, PrefixReuseStrategy) # noqa: SLF001 assert isinstance( # noqa: SLF001 connector._request_lifecycle._reuse_strategy(), # noqa: SLF001 PrefixReuseStrategy, ) -def test_start_load_kv_initializes_gds_after_server_creates_store( - monkeypatch, tmp_path -): - """Worker load path marks server transfer ready after deferred startup.""" - store_path = tmp_path / "daser.store" - store_path.write_bytes(b"\0" * 4096) - - connector = _WorkerProbe(str(store_path)) - - connector.start_load_kv(forward_context=object()) - - assert connector.transfer_ready is True - - -def test_start_load_kv_does_not_emit_info_timing( - monkeypatch, - caplog: pytest.LogCaptureFixture, -) -> None: - """Worker load timing stays out of the hot INFO logging path.""" - connector = _WorkerProbe("") - connector._transfer_ready = True # noqa: SLF001 - connector._transfer_mode = "iouring" # noqa: SLF001 - connector._skip_l2 = True # noqa: SLF001 - connector._slot_size = 4 # noqa: SLF001 - connector._store_staging_bytes = 64 # noqa: SLF001 - connector._kv_caches = { # noqa: SLF001 - "layer.0": torch.empty(2, 1, 1, 4, dtype=torch.uint8) - } - connector._layer_names = ["layer.0"] # noqa: SLF001 - connector._load_key_scale = 1.0 # noqa: SLF001 - connector._load_value_scale = 1.0 # noqa: SLF001 - connector._rope_delta_scale = 1.0 # noqa: SLF001 - connector._rope_base = 10000.0 # noqa: SLF001 - connector._rope_rotary_dim = 0 # noqa: SLF001 - connector._rope_is_neox_style = True # noqa: SLF001 - connector._meta = DaserConnectorMeta( # noqa: SLF001 - reqs_to_load={ - "req": ReqLoadSpec( - chunk_key="hit", - start_slot=0, - num_slots=1, - block_ids=[0], - file_offset=0, - token_count=4, - ) - } - ) - - class _Lease: - view = torch.empty(4, dtype=torch.uint8) - - def release(self) -> None: - return - - class _Future: - def done(self) -> bool: - return True - - def result(self, timeout: float): - del timeout - return { - "transfer_open_ms": 0.0, - "transfer_load_ms": 0.0, - "transfer_sync_ms": 0.0, - "transfer_stats_delta": {"l1_hits": 1, "l1_misses": 0, "l2_reads": 0}, - } - - def fake_submit_load_coroutine(coro): - coro.close() - return _Future() - - monkeypatch.setattr(connector, "_acquire_staging", lambda *args: _Lease()) - monkeypatch.setattr(connector, "_submit_load_coroutine", fake_submit_load_coroutine) - monkeypatch.setattr( - "daser.connector.worker_runtime.cupy.asarray", lambda tensor: tensor - ) - monkeypatch.setattr( - "daser.connector.worker_runtime.export_cuda_ipc_handle", lambda _: b"0" - ) - monkeypatch.setattr( - "daser.connector.worker_runtime.cuda_array_device_id", lambda _: 0 - ) - monkeypatch.setattr( - "daser.connector.worker_runtime.cuda_array_pointer", lambda _: 1 - ) - monkeypatch.setattr( - "daser.connector.worker_runtime._cuda_allocation_base_and_offset", - lambda _: (1, 0), - ) - monkeypatch.setattr( - "daser.connector.worker_runtime._copy_staging_to_kv_cache", - lambda **kwargs: 1, - ) - monkeypatch.setattr( - "daser.connector.worker_runtime._synchronize_cuda_tensor", - lambda tensor: None, - ) - - with caplog.at_level("INFO", logger="daser.connector.worker_runtime"): - connector.start_load_kv(forward_context=object()) - - assert "start_load_kv timing" not in caplog.text - - -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} - - load0 = _Client("load0") - load1 = _Client("load1") - - class Probe(WorkerRuntime): - def __init__(self) -> None: - self._load_pipeline = _LoadPipelineProbe(clients=[load0, 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 load0.calls and load1.calls - - -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._load_pipeline.pending) == {"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("") - connector._transfer_ready = True # noqa: SLF001 - connector._transfer_mode = "iouring" # noqa: SLF001 - connector._skip_l2 = True # noqa: SLF001 - connector._layer_names = [] # noqa: SLF001 - connector._meta = DaserConnectorMeta( # noqa: SLF001 - reqs_to_load={ - "req": ReqLoadSpec( - chunk_key="hit", - start_slot=0, - num_slots=1, - block_ids=[7], - file_offset=0, - token_count=4, - ) - } - ) - - connector.start_load_kv(forward_context=object()) - - assert connector.get_finished(set()) == (None, {"req"}) - assert connector.get_block_ids_with_load_errors() == {7} - - -def test_wait_for_save_defers_store_until_request_finished() -> None: - """Cold stores should be snapshotted after vLLM reports request finish.""" - connector = _FinishedSaveProbe() - connector.set_pending_meta( - { - "req": ReqStoreSpec( - chunk_key="stored", - start_slot=0, - num_slots=2, - block_ids=[4, 5], - file_offset=0, - token_count=8, - ) - }, - {"stored"}, - ) - - connector.wait_for_save() - - assert connector.staged_batches == [] - assert connector.submitted == 0 - assert connector.pending_finished_save_ids() == {"req"} - assert connector.pending_commit_keys() == set() - - finished_sending, finished_recving = connector.get_finished({"req"}) - - assert finished_recving is None - assert finished_sending == {"req"} - assert [block_ids for block_ids, _spans in connector.staged_batches] == [[4, 5]] - assert connector.submitted == 2 - assert connector.committed_after == [(1, ["stored"])] - - -def test_wait_for_save_groups_prefix_slot_stores_by_base_request() -> None: - """Synthetic prefix slot stores should finish with their base request.""" - connector = _FinishedSaveProbe() - connector.set_pending_meta( - { - "req:store:0": ReqStoreSpec( - chunk_key="stored-0", - start_slot=0, - num_slots=1, - block_ids=[4], - file_offset=0, - token_count=4, - ), - "req:store:1": ReqStoreSpec( - chunk_key="stored-1", - start_slot=1, - num_slots=1, - block_ids=[5], - file_offset=32, - token_count=4, - ), - }, - {"stored-0", "stored-1"}, - ) - - connector.wait_for_save() - - assert connector.pending_finished_save_ids() == {"req"} - - finished_sending, finished_recving = connector.get_finished({"req"}) - - assert finished_recving is None - assert finished_sending == {"req"} - assert [block_ids for block_ids, _spans in connector.staged_batches] == [[4, 5]] - assert connector.committed_after == [(1, ["stored-0", "stored-1"])] - - -def test_get_finished_holds_blocks_until_deferred_store_completes() -> None: - """Worker should not release finished request blocks before store is done.""" - connector = _FinishedSaveProbe() - connector.seed_finished_save( - "req", - { - "req": ReqStoreSpec( - chunk_key="stored", - start_slot=0, - num_slots=1, - block_ids=[4], - file_offset=0, - token_count=4, - ) - }, - ["stored"], - ) - - class _PendingFuture: - def done(self) -> bool: - return False - - def submit_pending(coro): - coro.close() - return _PendingFuture() - - connector.set_submit_store_coroutine(submit_pending) - - finished_sending, finished_recving = connector.get_finished({"req"}) - - assert finished_recving is None - assert finished_sending is None - assert connector.staged_batches - assert "req" in connector.pending_finished_save_ids() - - -def test_get_finished_reports_completed_deferred_store_on_later_step() -> None: - """Completed saves should be reported even after the original finish step.""" - connector = _FinishedSaveProbe() - pending_future = None - - class _PendingFuture: - def __init__(self) -> None: - self.complete = False - - def done(self) -> bool: - return self.complete - - def result(self, timeout: float): - del timeout - return None - - def submit_pending(coro): - nonlocal pending_future - coro.close() - pending_future = _PendingFuture() - return pending_future - - connector.seed_finished_save( - "req", - { - "req": ReqStoreSpec( - chunk_key="stored", - start_slot=0, - num_slots=1, - block_ids=[4], - file_offset=0, - token_count=4, - ) - }, - ["stored"], - ) - connector.set_submit_store_coroutine(submit_pending) - - assert connector.get_finished({"req"}) == (None, None) - assert pending_future is not None - pending_future.complete = True - - assert connector.get_finished(set()) == ({"req"}, None) - - -def test_get_finished_releases_request_when_no_store_batch_can_be_staged() -> None: - """A skipped staging batch should not keep scheduler blocks forever.""" - connector = _FinishedSaveProbe() - connector.seed_finished_save( - "req", - { - "req": ReqStoreSpec( - chunk_key="stored", - start_slot=0, - num_slots=1, - block_ids=[4], - file_offset=0, - token_count=4, - ) - }, - ["stored"], - ) - connector.disable_store_staging() - - assert connector.get_finished({"req"}) == ({"req"}, None) - assert connector.pending_finished_save_ids() == set() - assert connector.submitted == 0 - - -def test_get_finished_does_not_block_on_pending_async_load() -> None: - """Worker polling should not block on incomplete async load futures.""" - connector = _AsyncLoadProbe() - future = _AsyncLoadFuture(done=False) - connector.seed_load("req", future, block_ids=[4, 5]) - - finished_sending, finished_recving = connector.get_finished(set()) - - assert finished_sending is None - assert finished_recving is None - assert future.result_calls == 0 - assert "req" in connector._load_pipeline.pending # noqa: SLF001 - - -def test_get_finished_reports_completed_async_load() -> None: - """Completed async loads should release the waiting request.""" - connector = _AsyncLoadProbe() - future = _AsyncLoadFuture(done=True) - connector.seed_load("req", future, block_ids=[4, 5]) - - finished_sending, finished_recving = connector.get_finished(set()) - - assert finished_sending is None - assert finished_recving == {"req"} - assert future.result_calls == 1 - assert connector._load_pipeline.pending == {} # noqa: SLF001 - assert connector._load_pipeline.invalid_block_ids == set() # noqa: SLF001 - - -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) - - assert dispatcher.effective_inflight == 3 - - def test_load_staging_pool_depth_respects_available_cuda_memory(monkeypatch) -> None: """Load staging depth should not preallocate past available CUDA headroom.""" monkeypatch.setattr( @@ -1220,127 +346,94 @@ def test_load_staging_pool_depth_respects_available_cuda_memory(monkeypatch) -> lambda device=None: ((4 << 30), 80 << 30), ) - depth = _load_staging_pool_depth( + depth = load_staging_pool_depth( buffer_bytes=1536 << 20, pending_limit_bytes=12 << 30, device=torch.device("cuda"), + max_inflight=8, + reserve_bytes=1 << 30, ) assert depth == 2 -def test_load_request_dispatcher_submits_without_waiting_for_batch() -> None: - """Queued requests should be submitted immediately while slots are free.""" +def test_worker_transfer_ready_allows_skip_l2_without_store_path() -> None: + """L1-only mode has no store path but still has a valid transfer config.""" - class Probe(_AsyncLoadProbe): + class Pipeline: def __init__(self) -> None: - super().__init__() - 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=[], - ) - - def _consume_dispatcher_load(self, state, sample_tensor): - del sample_tensor - 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() - 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, - ), - 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) + self.initialized = False - assert probe.submitted == ["req-a", "req-b"] - assert [item.req_id for item in queued] == ["req-c"] - assert len(active) == 2 + def configure_rank_geometry(self, *args: int) -> None: + del args - 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) + def initialize_transfer(self) -> None: + self.initialized = True - assert probe.completed == ["req-a"] - assert probe.submitted == ["req-a", "req-b", "req-c"] - assert first_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() - future = _AsyncLoadFuture(done=True, error=RuntimeError("load failed")) - connector.seed_load("req", future, block_ids=[4, 5]) - - finished_sending, finished_recving = connector.get_finished(set()) + connector = WorkerRuntime.__new__(WorkerRuntime) + connector._transfer_ready = False # noqa: SLF001 + connector._pipelines_initialized = False # noqa: SLF001 + connector._store_path = "" # noqa: SLF001 + connector._slot_size = 1024 # noqa: SLF001 + connector._local_slot_size = 1024 # noqa: SLF001 + connector._tp_size = 1 # noqa: SLF001 + connector._server_tp_size = 1 # noqa: SLF001 + connector._tp_rank = 0 # noqa: SLF001 + connector._rank_stride_bytes = 0 # noqa: SLF001 + connector._transfer_mode = "iouring" # noqa: SLF001 + connector._skip_l2 = True # noqa: SLF001 + connector._refresh_runtime_config = lambda: None # noqa: SLF001 + connector._load_pipeline = Pipeline() # noqa: SLF001 + connector._store_pipeline = Pipeline() # noqa: SLF001 - assert finished_sending is None - assert finished_recving == {"req"} - assert connector._load_pipeline.pending == {} # noqa: SLF001 - assert connector._load_pipeline.invalid_block_ids == {4, 5} # noqa: SLF001 + assert connector._ensure_transfer_ready() is True # noqa: SLF001 + assert connector._transfer_ready is True # noqa: SLF001 + assert connector._pipelines_initialized is True # noqa: SLF001 + assert connector._load_pipeline.initialized is True # noqa: SLF001 + assert connector._store_pipeline.initialized is True # noqa: SLF001 -def test_shutdown_collects_failed_async_load_without_raising() -> None: - """Shutdown should still clean worker resources after load failures.""" - connector = _AsyncLoadProbe() - future = _AsyncLoadFuture(done=True, error=RuntimeError("load failed")) - connector.seed_load("req", future, block_ids=[4, 5]) +def test_worker_transfer_ready_propagates_refreshed_tp_geometry() -> None: + """Delayed TP geometry must reach pipelines before transfer initialization.""" - connector.shutdown() + class Pipeline: + def __init__(self) -> None: + self.geometry: tuple[int, ...] | None = None + self.initialized = False - assert connector._load_pipeline.pending == {} # noqa: SLF001 - assert connector._load_pipeline.invalid_block_ids == {4, 5} # noqa: SLF001 - assert connector._load_pipeline.shutdown_called # noqa: SLF001 - assert connector._store_pipeline.shutdown_called # noqa: SLF001 + def configure_rank_geometry(self, *args: int) -> None: + self.geometry = args + def initialize_transfer(self) -> None: + self.initialized = True -def test_worker_transfer_ready_allows_skip_l2_without_store_path() -> None: - """L1-only mode has no store path but still has a valid transfer config.""" - connector = _WorkerProbe("") + connector = WorkerRuntime.__new__(WorkerRuntime) + connector._transfer_ready = False # noqa: SLF001 + connector._pipelines_initialized = False # noqa: SLF001 + connector._store_path = "" # noqa: SLF001 + connector._slot_size = 0 # noqa: SLF001 + connector._local_slot_size = 512 # noqa: SLF001 + connector._tp_size = 2 # noqa: SLF001 + connector._server_tp_size = 1 # noqa: SLF001 + connector._tp_rank = 1 # noqa: SLF001 + connector._rank_stride_bytes = 0 # noqa: SLF001 connector._transfer_mode = "iouring" # noqa: SLF001 connector._skip_l2 = True # noqa: SLF001 + def refresh() -> None: + connector._slot_size = 1024 # noqa: SLF001 + connector._server_tp_size = 2 # noqa: SLF001 + connector._rank_stride_bytes = 8192 # noqa: SLF001 + + connector._refresh_runtime_config = refresh # noqa: SLF001 + connector._load_pipeline = Pipeline() # noqa: SLF001 + connector._store_pipeline = Pipeline() # noqa: SLF001 + assert connector._ensure_transfer_ready() is True # noqa: SLF001 - assert connector.transfer_ready is True + assert connector._load_pipeline.geometry == (8192, 1) # noqa: SLF001 + assert connector._store_pipeline.geometry == (8192, 1, 2) # noqa: SLF001 + assert connector._load_pipeline.initialized is True # noqa: SLF001 + assert connector._store_pipeline.initialized is True # noqa: SLF001 def test_worker_runtime_refreshes_l1_only_transfer_config(monkeypatch) -> None: @@ -1420,8 +513,8 @@ def get_runtime_config(self) -> dict[str, object]: assert [intent.chunk_key for intent in plan.intents] == rolling_keys(tokens, 128) -def test_runtime_config_ready_allows_skip_l2_without_store_path(monkeypatch): - """Scheduler runtime config should be ready when skip_l2 omits store_path.""" +def test_scheduler_runtime_config_is_owned_by_request_lifecycle(monkeypatch): + """Scheduler readiness is refreshed directly on its lifecycle owner.""" class DummyIPCClient: def __init__(self, socket_path): @@ -1463,8 +556,9 @@ class DummyVLLMConfig: role=KVConnectorRole.SCHEDULER, ) - assert connector._runtime_config_ready is True # noqa: SLF001 - assert connector._skip_l2 is True # noqa: SLF001 + lifecycle = connector._request_lifecycle # noqa: SLF001 + assert lifecycle._runtime_config_ready is True # noqa: SLF001 + assert lifecycle._model_id == "served-model" # noqa: SLF001 def test_scheduler_refreshes_runtime_config_before_lookup(monkeypatch): @@ -2328,17 +1422,21 @@ def test_register_kv_caches_warms_dynamic_rope_apply_once(monkeypatch): class Probe(WorkerRuntime): def __init__(self) -> None: self._slot_size = 0 - self._load_pipeline = _LoadPipelineProbe() - self._store_pipeline = _StorePipelineProbe() + self._tp_size = 1 + self._server_tp_size = 1 + self._tp_rank = 0 + self._rank_stride_bytes = 0 self._rope_rotary_dim = 8 self._rope_base = 10000.0 self._rope_is_neox_style = True - self._ipc_async = None - self._bg_loop = None def _ensure_transfer_ready(self) -> bool: return False + def _configure_pipelines(self, sample: torch.Tensor) -> int: + del sample + return 1 + calls = [] monkeypatch.setattr( worker, @@ -2382,8 +1480,10 @@ class Probe(WorkerRuntime): def __init__(self) -> None: self._kv_cache_config = Config() self._slot_size = 0 - self._load_pipeline = _LoadPipelineProbe() - self._store_pipeline = _StorePipelineProbe() + self._tp_size = 1 + self._server_tp_size = 1 + self._tp_rank = 0 + self._rank_stride_bytes = 0 self._rope_rotary_dim = 8 self._rope_base = 10000.0 self._rope_is_neox_style = True @@ -2395,6 +1495,10 @@ def _ensure_transfer_ready(self) -> bool: def _init_server_transfer(self) -> None: return + def _configure_pipelines(self, sample: torch.Tensor) -> int: + del sample + return 1 + @property def registration_state(self): return ( @@ -2422,165 +1526,6 @@ def registration_state(self): assert slot_size == kv_cache[0].nbytes -def test_stage_store_batch_does_not_warm_dynamic_rope_again(monkeypatch): - from daser.connector import worker_runtime as worker - - class Probe(WorkerRuntime): - def __init__(self) -> None: - self.kv_cache = torch.zeros((2, 8, 4, 2, 8), dtype=torch.float32) - self._layer_names = ["layer.0"] - self._layer_idx_map = {"layer.0": 0} - self._kv_caches = {"layer.0": self.kv_cache} - self.slot_size = self.kv_cache[:, 0].nbytes - self._slot_size = self.slot_size - self._store_pipeline = _StorePipelineProbe() - self._store_pipeline.staging_bytes = 4096 - self._store_pipeline.pending_staging_limit_bytes = 8192 - self._rope_rotary_dim = 8 - self._rope_base = 10000.0 - self._rope_is_neox_style = True - - def stage_store_batch(self, block_ids: list[int], spans: list[StoreWriteSpan]): - """Expose store staging through a public test helper.""" - return self._stage_store_batch(block_ids, spans) - - calls = [] - monkeypatch.setattr( - worker, - "_warm_rope_apply_backends", - lambda **kwargs: calls.append(kwargs), - ) - - probe = Probe() - staged = probe.stage_store_batch( - block_ids=[0, 1, 2], - spans=[StoreWriteSpan(0, probe.slot_size * 3, 0, "k0", 0, 3)], - ) - - assert staged is not None - assert calls == [] - - -def test_store_staging_wait_skips_future_without_lease() -> None: - """Store backpressure waits until a future returns a staging lease.""" - - class _Future: - def __init__(self, name: str) -> None: - self.name = name - - def done(self) -> bool: - return False - - def result(self, timeout: float) -> None: - assert timeout > 0 - completed.append(self.name) - - class Probe(WorkerRuntime): - def __init__(self) -> None: - self._store_pipeline = _StorePipelineProbe() - self._store_pipeline.staging_pool = FixedCudaStagingPool( - device=torch.device("cpu"), - buffer_bytes=16, - depth=1, - ) - lease = self._store_pipeline.staging_pool.acquire(8) - self._store_pipeline.pending_staging_limit_bytes = 64 - self._store_pipeline.pending_staging_bytes = 8 - self._store_pipeline.save_futures = [ - _SaveFuture(_Future("commit"), 0, None), - _SaveFuture(_Future("store"), 8, lease), - ] - - def wait_for_staging(self) -> None: - self._wait_for_store_staging_release(8) - - completed: list[str] = [] - probe = Probe() - probe.wait_for_staging() - - assert completed == ["commit", "store"] - assert probe._store_pipeline.staging_pool.available == 1 # noqa: SLF001 - - -def test_stage_store_batch_records_ready_event_without_synchronizing(monkeypatch): - from daser.connector import worker_runtime as worker - - class Probe(WorkerRuntime): - def __init__(self) -> None: - self.kv_cache = torch.zeros((2, 8, 4, 2, 8), dtype=torch.float32) - self._layer_names = ["layer.0"] - self._layer_idx_map = {"layer.0": 0} - self._kv_caches = {"layer.0": self.kv_cache} - self.slot_size = self.kv_cache[:, 0].nbytes - self._slot_size = self.slot_size - self._store_pipeline = _StorePipelineProbe() - self._store_pipeline.staging_bytes = 4096 - self._store_pipeline.pending_staging_limit_bytes = 8192 - self._rope_rotary_dim = 8 - self._rope_base = 10000.0 - self._rope_is_neox_style = True - - def stage_store_batch(self, block_ids: list[int], spans: list[StoreWriteSpan]): - """Expose store staging through a public test helper.""" - return self._stage_store_batch(block_ids, spans) - - synced = [] - recorded = object() - monkeypatch.setattr( - worker, - "_synchronize_cuda_tensor", - lambda tensor: synced.append(tensor), - ) - monkeypatch.setattr(worker, "_record_cuda_event", lambda tensor: recorded) - - probe = Probe() - staged = probe.stage_store_batch( - block_ids=[0, 1, 2], - spans=[StoreWriteSpan(0, probe.slot_size * 3, 0, "k0", 0, 3)], - ) - - assert staged is not None - assert staged.ready_event is recorded - assert synced == [] - - -def test_stage_store_batch_keeps_dynamic_rope_warmup_out_of_store_path(monkeypatch): - from daser.connector import worker_runtime as worker - - class Probe(WorkerRuntime): - def __init__(self) -> None: - self.kv_cache = torch.zeros((2, 8, 4, 2, 8), dtype=torch.float32) - self._layer_names = ["layer.0"] - self._layer_idx_map = {"layer.0": 0} - self._kv_caches = {"layer.0": self.kv_cache} - self.slot_size = self.kv_cache[:, 0].nbytes - self._slot_size = self.slot_size - self._store_pipeline = _StorePipelineProbe() - self._store_pipeline.staging_bytes = 4096 - self._store_pipeline.pending_staging_limit_bytes = 8192 - self._rope_rotary_dim = 8 - self._rope_base = 10000.0 - self._rope_is_neox_style = True - - def stage_store_batch(self, block_ids: list[int], spans: list[StoreWriteSpan]): - """Expose store staging through a public test helper.""" - return self._stage_store_batch(block_ids, spans) - - calls = [] - monkeypatch.setattr( - worker, - "_warm_rope_apply_backends", - lambda **kwargs: calls.append(kwargs), - ) - - probe = Probe() - spans = [StoreWriteSpan(0, probe.slot_size * 3, 0, "k0", 0, 3)] - assert probe.stage_store_batch([0, 1, 2], spans) is not None - assert probe.stage_store_batch([3, 4, 5], spans) is not None - - assert calls == [] - - def test_update_state_after_alloc_skips_chunks_beyond_external_prefix(): class MockConnector(RequestLifecycle): def __init__(self) -> None: @@ -3935,54 +2880,55 @@ def test_build_staging_store_batches_uses_spec_file_offset(): @pytest.mark.asyncio -async def test_commit_empty_stored_keys_does_not_publish_requested_chunks(): - """Skipped stale stores must not commit requested chunks.""" - connector = _CommitProbe() +async def test_store_cuda_export_selects_staged_buffer_device(monkeypatch) -> None: + """Background CUDA IPC export must select the TP rank's staged device.""" + from daser.connector import store_pipeline as store_module + + selected_devices: list[torch.device] = [] + transferred: list[dict] = [] - await connector.commit_stored_keys([], ["stale-key"]) + class Client: + async def transfer_store_cuda(self, **kwargs): + transferred.append(kwargs) + return [] - assert connector.committed == [[]] - assert connector.commit_metadata == [(0, 1)] + pipeline = StorePipeline.__new__(StorePipeline) + pipeline._client = Client() # noqa: SLF001 + buffer = SimpleNamespace(device=torch.device("cuda:1"), nbytes=32) + staged = StagedStoreBatch(buffer=buffer, spans=[], lease=object()) + cupy_buffer = object() + + monkeypatch.setattr(torch.cuda, "set_device", selected_devices.append) + monkeypatch.setattr(store_module.cupy, "asarray", lambda tensor: cupy_buffer) + monkeypatch.setattr(store_module, "cuda_array_pointer", lambda array: 4096) + monkeypatch.setattr( + store_module, + "cuda_allocation_base_and_offset", + lambda pointer: (pointer, 0), + ) + monkeypatch.setattr(store_module, "export_cuda_ipc_handle", lambda array: b"ipc") + monkeypatch.setattr(store_module, "cuda_array_device_id", lambda array: 1) + + await pipeline._write_cuda_buffer(staged) # noqa: SLF001 + + assert selected_devices == [torch.device("cuda:1")] + assert transferred[0]["device_id"] == 1 def test_tensor_parallel_rank_lanes_are_contiguous_and_disjoint() -> None: """Each TP rank maps a logical slot run into its own contiguous lane.""" local_slot_size = 32 rank_stride = 10 * local_slot_size + start_slot = 3 - rank_0 = _rank_lane_offset(3, local_slot_size, rank_stride, tp_rank=0) - rank_1 = _rank_lane_offset(3, local_slot_size, rank_stride, tp_rank=1) + rank_0 = start_slot * local_slot_size + rank_1 = rank_stride + start_slot * local_slot_size assert rank_0 == 3 * local_slot_size assert rank_1 == rank_stride + 3 * local_slot_size assert rank_0 + 2 * local_slot_size <= rank_1 -@pytest.mark.asyncio -async def test_worker_load_and_store_use_separate_ipc_clients(): - """Worker load RPCs should not queue behind store RPCs on one IPC client.""" - connector = _QueueProbe() - - await connector.load_via_public_helper() - await connector.store_via_public_helper() - - assert connector.load_calls == ["load"] - assert connector.store_calls == ["store"] - - -def test_worker_load_and_store_use_separate_background_loops(): - """Foreground loads should not queue behind background store loop work.""" - connector = _LoopProbe() - - connector.submit_load() - connector.submit_store() - - load_pipeline, store_pipeline = connector.loop_pair - assert load_pipeline.submit_calls == 1 - assert store_pipeline.submit_calls == 1 - assert load_pipeline is not store_pipeline - - def test_derive_staging_limits_scale_with_vram(monkeypatch): """GPU staging caps consider device size and currently free VRAM.""" diff --git a/tests/unit/test_block_aligned_bug.py b/tests/unit/test_block_aligned_bug.py index f4320de..b38bc53 100644 --- a/tests/unit/test_block_aligned_bug.py +++ b/tests/unit/test_block_aligned_bug.py @@ -428,7 +428,9 @@ def __init__(self): ) def test_trim_external_window_keeps_full_block_for_lmcache_style_minus_one(self): - from daser.connector.scheduler import _trim_chunk_to_external_window + from daser.connector.scheduler_planning import ( + _trim_chunk_to_external_window, + ) chunk = { "chunk_key": "k", From 95a071e08c925709d4356ef04737ae7e46540869 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Sun, 12 Jul 2026 20:38:06 +0800 Subject: [PATCH 3/4] refactor(connector): restore pipeline behavior and boundaries - separate scheduler and worker implementations into role-owned packages - restore concurrent streaming load and store pipeline behavior - cap independent fixed staging pools under one 6 GiB worker budget - restore load diagnostics and strengthen drain and tier benchmark gates - add behavior-contract coverage and update connector design documentation --- benchmarks/bench_staging_restore.py | 2 +- benchmarks/utils/vllm_bench.py | 65 ++++- daser/connector/daser_connector.py | 14 +- daser/connector/ipc_client.py | 2 +- daser/connector/scheduler/__init__.py | 6 + .../{scheduler.py => scheduler/adapter.py} | 0 .../lifecycle.py} | 4 +- .../planning.py} | 0 daser/connector/{ => scheduler}/reuse.py | 0 daser/connector/worker/__init__.py | 13 + .../{worker.py => worker/adapter.py} | 0 .../{load_pipeline.py => worker/load.py} | 272 ++++++++++++------ .../{worker_memory.py => worker/memory.py} | 138 ++++----- .../{worker_runtime.py => worker/runtime.py} | 41 ++- daser/connector/{ => worker}/staging.py | 0 .../{store_pipeline.py => worker/store.py} | 68 ++--- daser/server/ipc/server.py | 14 +- docs/design/architecture.md | 4 +- docs/design/components.md | 6 +- .../3_server_managed_transfer.md | 2 +- docs/optimizations/4_rope_apply_compile.md | 2 +- tests/connector/test_connector_layout.py | 4 +- tests/connector/test_daser_connector.py | 145 +++++----- tests/connector/test_gds_transfer.py | 6 +- tests/connector/test_worker_pipelines.py | 214 ++++++++++++++ tests/server/test_ipc_server.py | 27 +- tests/unit/test_benchmark_unified_utils.py | 86 ++++++ tests/unit/test_block_aligned_bug.py | 4 +- tests/unit/test_chunk_reuse_scheduler.py | 2 +- tests/unit/test_skip_save_flag.py | 2 +- 30 files changed, 792 insertions(+), 351 deletions(-) create mode 100644 daser/connector/scheduler/__init__.py rename daser/connector/{scheduler.py => scheduler/adapter.py} (100%) rename daser/connector/{request_lifecycle.py => scheduler/lifecycle.py} (99%) rename daser/connector/{scheduler_planning.py => scheduler/planning.py} (100%) rename daser/connector/{ => scheduler}/reuse.py (100%) create mode 100644 daser/connector/worker/__init__.py rename daser/connector/{worker.py => worker/adapter.py} (100%) rename daser/connector/{load_pipeline.py => worker/load.py} (74%) rename daser/connector/{worker_memory.py => worker/memory.py} (70%) rename daser/connector/{worker_runtime.py => worker/runtime.py} (95%) rename daser/connector/{ => worker}/staging.py (100%) rename daser/connector/{store_pipeline.py => worker/store.py} (91%) create mode 100644 tests/connector/test_worker_pipelines.py diff --git a/benchmarks/bench_staging_restore.py b/benchmarks/bench_staging_restore.py index 990e118..3b46e6b 100644 --- a/benchmarks/bench_staging_restore.py +++ b/benchmarks/bench_staging_restore.py @@ -20,7 +20,7 @@ # First Party from benchmarks.utils.constants import BLOCK_TOKENS # noqa: E402 -from daser.connector.staging import copy_staging_to_kv_cache # noqa: E402 +from daser.connector.worker.staging import copy_staging_to_kv_cache # noqa: E402 from daser.ops.rope_apply import clear_rope_apply_cache # noqa: E402 diff --git a/benchmarks/utils/vllm_bench.py b/benchmarks/utils/vllm_bench.py index 73f4005..c292ad9 100644 --- a/benchmarks/utils/vllm_bench.py +++ b/benchmarks/utils/vllm_bench.py @@ -11,7 +11,7 @@ import asyncio from collections.abc import Callable -from dataclasses import asdict +from dataclasses import asdict, replace import json import math from pathlib import Path @@ -297,9 +297,19 @@ def run_load( asyncio.run(_wait_lmcache_quiescent(manifest, settle_seconds=0.0)) elif backend_run.backend == "daser": _drain_daser(manifest, print_kv=print_kv) - warm_metrics, warm_hit_rate = _run_phase( - args, manifest, warm_raw, run_command=run_command - ) + if backend_run.backend == "daser" and args.evict: + warm_metrics, warm_hit_rate = _run_daser_evict_warm_phase( + args, + manifest, + warm_raw, + run_command=run_command, + ) + else: + warm_metrics, warm_hit_rate = _run_phase( + args, manifest, warm_raw, run_command=run_command + ) + if backend_run.backend == "daser" and args.evict: + _require_daser_evict_tier_activity(warm_metrics) cold_summary = _normalise_result(cold_raw) warm_summary = _normalise_result(warm_raw) _apply_phase_metrics(cold_summary, cold_hit_rate) @@ -373,6 +383,33 @@ def _run_phase( return _collect_phase_metrics(manifest, before_metrics) +def _run_daser_evict_warm_phase( + args: RunBenchArgs, + manifest: BenchmarkManifest, + raw_path: Path, + *, + run_command: Callable[[list[str]], Any], +) -> tuple[dict[str, Any], float | None]: + """Perturb LRU order before measuring a complete evict warm phase. + + A same-order scan of a working set larger than L1 can produce 100% L2 + reads even when most entries were resident. Replaying the first 20% of the + deterministic workload first makes the subsequent complete phase exercise + both resident L1 entries and evicted L2 entries. Both commands are included + in the returned warm metric delta; correctness and latency use only the + complete phase result. + """ + before_metrics = asyncio.run(collect_phase_metrics(manifest)) + prime_args = replace( + args, + bench_num_prompts=max(1, math.ceil(args.bench_num_prompts * 0.2)), + ) + prime_path = raw_path.with_name("vllm_bench_warm_prime.json") + run_command(_bench_command(prime_args, manifest.endpoints["vllm"], prime_path)) + run_command(_bench_command(args, manifest.endpoints["vllm"], raw_path)) + return _collect_phase_metrics(manifest, before_metrics) + + def _collect_phase_metrics( manifest: BenchmarkManifest, before_metrics: dict[str, Any] | None, @@ -414,11 +451,21 @@ def _drain_daser( endpoint = manifest.endpoints.get("daser") if endpoint is None: return - try: - response = httpx.post(f"{endpoint.url}/drain", timeout=30.0) - response.raise_for_status() - except Exception as exc: # noqa: BLE001 - print_kv("daser_drain_status", f"unavailable ({exc})") + response = httpx.post(f"{endpoint.url}/drain", timeout=30.0) + response.raise_for_status() + print_kv("daser_drain_status", "ok") + + +def _require_daser_evict_tier_activity(metrics: dict[str, Any]) -> None: + """Require one evict warm phase to exercise both DaseR cache tiers.""" + counters = metrics.get("backend_prometheus", {}) + l1_hits = float(counters.get("daser_l1_hits_total", 0.0)) + l2_reads = float(counters.get("daser_l2_reads_total", 0.0)) + if l1_hits <= 0 or l2_reads <= 0: + raise RuntimeError( + "DaseR evict warm phase must exercise both tiers: " + f"l1_hits={l1_hits:g} l2_reads={l2_reads:g}" + ) def _bench_command( diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index d15e252..44ef88d 100644 --- a/daser/connector/daser_connector.py +++ b/daser/connector/daser_connector.py @@ -17,13 +17,13 @@ # First Party from daser.connector.ipc_client import IPCClientSync from daser.connector.metadata import DaserConnectorMeta, ReqLoadSpec, ReqStoreSpec -from daser.connector.request_lifecycle import RequestLifecycle -from daser.connector.scheduler import SchedulerConnectorMixin -from daser.connector.staging import ( +from daser.connector.scheduler.adapter import SchedulerConnectorMixin +from daser.connector.scheduler.lifecycle import RequestLifecycle +from daser.connector.worker.adapter import WorkerConnectorMixin +from daser.connector.worker.runtime import WorkerRuntime +from daser.connector.worker.staging import ( DEFAULT_ROPE_DELTA_SCALE, ) -from daser.connector.worker import WorkerConnectorMixin -from daser.connector.worker_runtime import WorkerRuntime from daser.logging import init_logger logger = init_logger(__name__) @@ -70,8 +70,8 @@ class DaserConnector( The entrypoint remains in this module for vLLM's ``kv_connector_module_path``. Scheduler-role behavior lives in - ``daser.connector.scheduler`` and worker-role behavior lives in - ``daser.connector.worker``. + ``daser.connector.scheduler.adapter`` and worker-role behavior lives in + ``daser.connector.worker.adapter``. Args: vllm_config: full VllmConfig from vLLM. diff --git a/daser/connector/ipc_client.py b/daser/connector/ipc_client.py index f024c10..4fedc06 100644 --- a/daser/connector/ipc_client.py +++ b/daser/connector/ipc_client.py @@ -498,7 +498,7 @@ async def transfer_store_cuda( allocation_base_ptr: int, allocation_offset: int, producer_pid: int, - spans: list[dict[str, int]], + spans: list[dict[str, Any]], ) -> list[str]: """Store from a CUDA IPC buffer through the server transfer layer. diff --git a/daser/connector/scheduler/__init__.py b/daser/connector/scheduler/__init__.py new file mode 100644 index 0000000..87904f9 --- /dev/null +++ b/daser/connector/scheduler/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 + +from daser.connector.scheduler.adapter import SchedulerConnectorMixin +from daser.connector.scheduler.lifecycle import RequestLifecycle + +__all__ = ["RequestLifecycle", "SchedulerConnectorMixin"] diff --git a/daser/connector/scheduler.py b/daser/connector/scheduler/adapter.py similarity index 100% rename from daser/connector/scheduler.py rename to daser/connector/scheduler/adapter.py diff --git a/daser/connector/request_lifecycle.py b/daser/connector/scheduler/lifecycle.py similarity index 99% rename from daser/connector/request_lifecycle.py rename to daser/connector/scheduler/lifecycle.py index 71e94c0..ddcc69c 100644 --- a/daser/connector/request_lifecycle.py +++ b/daser/connector/scheduler/lifecycle.py @@ -13,8 +13,7 @@ from daser.connector.helpers import PendingStore from daser.connector.metadata import DaserConnectorMeta, ReqLoadSpec, ReqStoreSpec -from daser.connector.reuse import build_cache_reuse_strategy -from daser.connector.scheduler_planning import ( +from daser.connector.scheduler.planning import ( _base_req_id, _computed_tokens_after_step, _contiguous_prefix_tokens, @@ -25,6 +24,7 @@ _store_slot_index, _trim_chunk_to_external_window, ) +from daser.connector.scheduler.reuse import build_cache_reuse_strategy from daser.logging import init_logger logger = init_logger(__name__) diff --git a/daser/connector/scheduler_planning.py b/daser/connector/scheduler/planning.py similarity index 100% rename from daser/connector/scheduler_planning.py rename to daser/connector/scheduler/planning.py diff --git a/daser/connector/reuse.py b/daser/connector/scheduler/reuse.py similarity index 100% rename from daser/connector/reuse.py rename to daser/connector/scheduler/reuse.py diff --git a/daser/connector/worker/__init__.py b/daser/connector/worker/__init__.py new file mode 100644 index 0000000..c81b912 --- /dev/null +++ b/daser/connector/worker/__init__.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 + +from daser.connector.worker.adapter import WorkerConnectorMixin +from daser.connector.worker.load import LoadPipeline +from daser.connector.worker.runtime import WorkerRuntime +from daser.connector.worker.store import StorePipeline + +__all__ = [ + "LoadPipeline", + "StorePipeline", + "WorkerConnectorMixin", + "WorkerRuntime", +] diff --git a/daser/connector/worker.py b/daser/connector/worker/adapter.py similarity index 100% rename from daser/connector/worker.py rename to daser/connector/worker/adapter.py diff --git a/daser/connector/load_pipeline.py b/daser/connector/worker/load.py similarity index 74% rename from daser/connector/load_pipeline.py rename to daser/connector/worker/load.py index 7c03b8d..2820286 100644 --- a/daser/connector/load_pipeline.py +++ b/daser/connector/worker/load.py @@ -5,9 +5,11 @@ import asyncio from collections import deque from concurrent.futures import Future +from contextlib import nullcontext from dataclasses import dataclass, replace import os import threading +import time from typing import Any # Third Party @@ -17,11 +19,11 @@ from daser.connector.helpers import base_req_id from daser.connector.ipc_client import IPCClientAsync from daser.connector.metadata import ReqLoadSpec -from daser.connector.staging import copy_staging_to_kv_cache -from daser.connector.worker_memory import ( +from daser.connector.worker.memory import ( CudaStagingLease, FixedCudaStagingPool, ) +from daser.connector.worker.staging import copy_staging_to_kv_cache from daser.logging import init_logger from daser.transfer.cuda_ipc import ( cuda_allocation_base_and_offset, @@ -37,42 +39,58 @@ @dataclass -class _PendingLoad: - """Track one request load until vLLM can resume it.""" +class _LoadRequest: + """Own one base request's specs and completion future.""" + req_id: str + specs: dict[str, ReqLoadSpec] future: Future[None] - block_ids: list[int] + + @property + def block_ids(self) -> list[int]: + """Return all vLLM blocks affected by this request.""" + return [block_id for spec in self.specs.values() for block_id in spec.block_ids] @dataclass class _InflightLoadBatch: """Hold one submitted load batch and its fixed staging lease.""" - buffer_index: int + total_bytes: int per_req_ranges: list[Any] staging_lease: CudaStagingLease future: Any + submitted_at: float -@dataclass -class _QueuedLoadRequest: - """Represent one request waiting for a load staging buffer.""" - - req_id: str - spec_id: str - spec: ReqLoadSpec - future: Future[None] +@dataclass(frozen=True) +class _LoadBatchTiming: + """Record transfer and restore accounting for one load batch.""" + + bytes: int + copies: int + copy_runs: int + ipc_ms: float + wait_ms: float + copy_ms: float + worker_sync_ms: float + transfer_open_ms: float + transfer_load_ms: float + transfer_sync_ms: float + l1_hits: int + l1_misses: int + l2_reads: int @dataclass class _InflightRequestLoad: """Track active and remaining load batches for one request.""" - item: _QueuedLoadRequest + request: _LoadRequest buffer_index: int - batches: list[_LoadBatch] - next_batch: int + batches: deque[_LoadBatch] active: _InflightLoadBatch | None + completed: list[_LoadBatchTiming] class LoadPipeline: @@ -98,7 +116,7 @@ def __init__(self, socket_path: str, client_count: int) -> None: self._queue: asyncio.Queue[Any] | None = None self._queue_lock = threading.Lock() self._dispatcher_future: Any | None = None - self._pending: dict[str, Any] = {} + self._pending: dict[str, _LoadRequest] = {} self._invalid_block_ids: set[int] = set() self._staging_pool: FixedCudaStagingPool | None = None self._staging_registered = False @@ -206,16 +224,15 @@ def start(self, reqs_to_load: dict[str, ReqLoadSpec]) -> None: self.mark_failed(reqs_to_load, "no registered KV cache layout") return queue = self._ensure_queue() + grouped: dict[str, dict[str, ReqLoadSpec]] = {} for spec_id, spec in reqs_to_load.items(): - req_id = base_req_id(spec_id) - request_future: Future[None] = Future() - self._pending[req_id] = _PendingLoad( - future=request_future, - block_ids=list(spec.block_ids), - ) + grouped.setdefault(base_req_id(spec_id), {})[spec_id] = spec + for req_id, specs in grouped.items(): + request = _LoadRequest(req_id, specs, Future()) + self._pending[req_id] = request self._loop.call_soon_threadsafe( queue.put_nowait, - _QueuedLoadRequest(req_id, spec_id, spec, request_future), + request, ) self._ensure_dispatcher() @@ -233,13 +250,13 @@ def mark_failed( Async/thread-safety: Called on the worker thread before background submission. """ - block_ids = [ - block_id for spec in reqs_to_load.values() for block_id in spec.block_ids - ] - failed_future: Future[None] = Future() - failed_future.set_exception(RuntimeError(reason)) - for req_id in {base_req_id(req_id) for req_id in reqs_to_load}: - self._pending[req_id] = _PendingLoad(failed_future, list(block_ids)) + grouped: dict[str, dict[str, ReqLoadSpec]] = {} + for spec_id, spec in reqs_to_load.items(): + grouped.setdefault(base_req_id(spec_id), {})[spec_id] = spec + for req_id, specs in grouped.items(): + future: Future[None] = Future() + future.set_exception(RuntimeError(reason)) + self._pending[req_id] = _LoadRequest(req_id, specs, future) def collect_finished(self) -> set[str]: """Collect completed loads without blocking the worker thread. @@ -347,7 +364,7 @@ async def _run_dispatcher(self) -> None: free_buffers = deque( range(max(1, min(len(self._clients), self._staging_pool.depth))) ) - queued: deque[_QueuedLoadRequest] = deque() + queued: deque[_LoadRequest] = deque() active: list[_InflightRequestLoad] = [] try: while True: @@ -360,9 +377,15 @@ async def _run_dispatcher(self) -> None: return queued.append(item) while queued and free_buffers: - state = self._submit_request( - queued.popleft(), free_buffers.popleft() - ) + request = queued.popleft() + buffer_index = free_buffers.popleft() + try: + state = self._submit_request(request, buffer_index) + except BaseException as exc: + if not request.future.done(): + request.future.set_exception(exc) + free_buffers.append(buffer_index) + continue if state.active is None: free_buffers.append(state.buffer_index) else: @@ -372,7 +395,15 @@ async def _run_dispatcher(self) -> None: active_batch = state.active if active_batch is not None and not active_batch.future.done(): continue - reusable_buffer, request_done = self._consume_request(state) + try: + reusable_buffer, request_done = self._consume_request(state) + except BaseException as exc: + if not state.request.future.done(): + state.request.future.set_exception(exc) + active.remove(state) + free_buffers.append(state.buffer_index) + consumed = True + continue if request_done: active.remove(state) free_buffers.append(reusable_buffer) @@ -383,8 +414,8 @@ async def _run_dispatcher(self) -> None: await self._wait_for_completion(active) except BaseException as exc: for state in active: - if not state.item.future.done(): - state.item.future.set_exception(exc) + if not state.request.future.done(): + state.request.future.set_exception(exc) if state.active is not None: state.active.staging_lease.release() for item in queued: @@ -395,7 +426,7 @@ async def _run_dispatcher(self) -> None: async def _drain_queue( self, queue: asyncio.Queue[Any], - queued: deque[_QueuedLoadRequest], + queued: deque[_LoadRequest], ) -> None: while True: try: @@ -424,38 +455,45 @@ async def _wait_for_completion( timeout=_LOAD_DISPATCH_WAIT_TIMEOUT_S, return_when=asyncio.FIRST_COMPLETED, ) + for future in done: + future.exception() if not done: await asyncio.sleep(0) def _submit_request( self, - item: _QueuedLoadRequest, + request: _LoadRequest, buffer_index: int, ) -> _InflightRequestLoad: if self._staging_pool is None: raise RuntimeError("load staging pool is not configured") - spec = replace( - item.spec, - file_offset=( - self._tp_rank * self._rank_stride_bytes - + item.spec.start_slot * self._local_slot_size - ), - ) - batches = build_load_read_batches( - {item.spec_id: spec}, - self._local_slot_size, - max_batch_bytes=self._staging_pool.buffer_bytes, - include_req_ids=True, + specs = { + spec_id: replace( + spec, + file_offset=( + self._tp_rank * self._rank_stride_bytes + + spec.start_slot * self._local_slot_size + ), + ) + for spec_id, spec in request.specs.items() + } + batches = deque( + build_load_read_batches( + specs, + self._local_slot_size, + max_batch_bytes=self._staging_pool.buffer_bytes, + include_req_ids=True, + ) ) if not batches: - item.future.set_result(None) - return _InflightRequestLoad(item, buffer_index, [], 0, 0, None) + request.future.set_result(None) + return _InflightRequestLoad(request, buffer_index, batches, None, []) return _InflightRequestLoad( - item=item, + request=request, buffer_index=buffer_index, batches=batches, - next_batch=1, - active=self._submit_batch(batches[0], buffer_index), + active=self._submit_batch(batches.popleft(), buffer_index), + completed=[], ) def _submit_batch( @@ -491,11 +529,13 @@ def _submit_batch( producer_pid=os.getpid(), spans=spans, ) + submitted_at = time.perf_counter() return _InflightLoadBatch( - buffer_index=buffer_index, + total_bytes=total_bytes, per_req_ranges=per_req_ranges, staging_lease=lease, future=self._submit(transfer), + submitted_at=submitted_at, ) def _consume_request( @@ -505,48 +545,106 @@ def _consume_request( active = state.active if active is None: return state.buffer_index, True - self._consume_batch(active) - state.buffer_index = active.buffer_index - if state.next_batch >= len(state.batches): - if not state.item.future.done(): - state.item.future.set_result(None) + state.completed.append(self._consume_batch(active)) + if not state.batches: + self._log_request_timing(state) + if not state.request.future.done(): + state.request.future.set_result(None) return state.buffer_index, True state.active = self._submit_batch( - state.batches[state.next_batch], + state.batches.popleft(), state.buffer_index, ) - state.next_batch += 1 return state.buffer_index, False - def _consume_batch(self, state: _InflightLoadBatch) -> None: + def _consume_batch(self, state: _InflightLoadBatch) -> _LoadBatchTiming: try: - state.future.result(timeout=120.0) + wait_start = time.perf_counter() + response = state.future.result(timeout=120.0) + wait_ms = (time.perf_counter() - wait_start) * 1000 + ipc_ms = (time.perf_counter() - state.submitted_at) * 1000 + copy_start = time.perf_counter() + copies, copy_runs = self._restore_batch(state) + copy_ms = (time.perf_counter() - copy_start) * 1000 + sync_ms = 0.0 if self._cuda_stream is None: - self._restore_batch(state) + pass else: - with torch.cuda.stream(self._cuda_stream): - self._restore_batch(state) + sync_start = time.perf_counter() self._cuda_stream.synchronize() + sync_ms = (time.perf_counter() - sync_start) * 1000 + payload = response if isinstance(response, dict) else {} + stats = payload.get("transfer_stats_delta", {}) + stats = stats if isinstance(stats, dict) else {} + return _LoadBatchTiming( + bytes=state.total_bytes, + copies=copies, + copy_runs=copy_runs, + ipc_ms=ipc_ms, + wait_ms=wait_ms, + copy_ms=copy_ms, + worker_sync_ms=sync_ms, + transfer_open_ms=float(payload.get("transfer_open_ms", 0.0)), + transfer_load_ms=float(payload.get("transfer_load_ms", 0.0)), + transfer_sync_ms=float(payload.get("transfer_sync_ms", 0.0)), + l1_hits=int(stats.get("l1_hits", 0)), + l1_misses=int(stats.get("l1_misses", 0)), + l2_reads=int(stats.get("l2_reads", 0)), + ) finally: state.staging_lease.release() - def _restore_batch(self, state: _InflightLoadBatch) -> None: + def _restore_batch(self, state: _InflightLoadBatch) -> tuple[int, int]: staging = state.staging_lease.view - for run in build_load_copy_runs(state.per_req_ranges): - 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._local_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, - ) + runs = build_load_copy_runs(state.per_req_ranges) + copies = 0 + context = ( + torch.cuda.stream(self._cuda_stream) + if self._cuda_stream is not None + else nullcontext() + ) + with context: + for run in 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._local_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, + ) + return copies, len(runs) + + def _log_request_timing(self, state: _InflightRequestLoad) -> None: + rows = state.completed + logger.debug( + "[CONNECTOR] load timing req=%s batches=%d bytes=%d copy_runs=%d " + "gpu_copies=%d ipc_ms=%.3f dispatcher_wait_ms=%.3f copy_ms=%.3f " + "worker_sync_ms=%.3f transfer_open_ms=%.3f " + "transfer_load_ms=%.3f transfer_sync_ms=%.3f l1_hits=%d " + "l1_misses=%d l2_reads=%d", + state.request.req_id, + len(rows), + sum(row.bytes for row in rows), + sum(row.copy_runs for row in rows), + sum(row.copies for row in rows), + sum(row.ipc_ms for row in rows), + sum(row.wait_ms for row in rows), + sum(row.copy_ms for row in rows), + sum(row.worker_sync_ms for row in rows), + sum(row.transfer_open_ms for row in rows), + sum(row.transfer_load_ms for row in rows), + sum(row.transfer_sync_ms for row in rows), + sum(row.l1_hits for row in rows), + sum(row.l1_misses for row in rows), + sum(row.l2_reads for row in rows), + ) def _register_staging_buffers(self) -> None: pool = self._staging_pool diff --git a/daser/connector/worker_memory.py b/daser/connector/worker/memory.py similarity index 70% rename from daser/connector/worker_memory.py rename to daser/connector/worker/memory.py index 5b9f64b..7c84f4a 100644 --- a/daser/connector/worker_memory.py +++ b/daser/connector/worker/memory.py @@ -12,9 +12,8 @@ import torch DEFAULT_STORE_STAGING_BYTES = 1536 << 20 -DEFAULT_PENDING_STORE_STAGING_BYTES = 3072 << 20 +DEFAULT_STAGING_BUDGET_BYTES = 6144 << 20 MIN_STORE_STAGING_BYTES = 64 << 20 -MIN_STORE_STAGING_POOL_DEPTH = 1 class _CudaStagingLeaseOwner(Protocol): @@ -203,98 +202,67 @@ def release(self, lease: CudaStagingLease) -> None: raise ValueError("lease does not belong to this fixed staging pool") -def derive_staging_limits(device: torch.device) -> tuple[int, int]: - """Return bounded staging batch and pending byte limits for a device. +def derive_staging_layout( + device: torch.device, + local_slot_size: int, + max_load_inflight: int, + reserve_bytes: int, +) -> tuple[int, int, int, int]: + """Partition one CUDA staging budget between load and store pools. Args: device: Device that owns worker-side staging tensors. + local_slot_size: Minimum buffer size required for one KV slot. + max_load_inflight: Maximum useful load pool depth. + reserve_bytes: Free CUDA memory kept outside staging pools. Returns: - Tuple of single-buffer bytes and total pending bytes. + Buffer bytes, load depth, store depth, and combined allocation bytes. + + Raises: + ValueError: If one buffer per direction cannot fit the budget. Async/thread-safety: Reads CUDA device properties during worker initialization before request traffic starts. """ + if local_slot_size <= 0: + raise ValueError("local_slot_size must be positive") + if max_load_inflight <= 0: + raise ValueError("max_load_inflight must be positive") if device.type != "cuda": - return DEFAULT_STORE_STAGING_BYTES, DEFAULT_PENDING_STORE_STAGING_BYTES - props = torch.cuda.get_device_properties(device) - total = int(props.total_memory) - try: - free, _ = torch.cuda.mem_get_info(device) - free = int(free) - except (RuntimeError, TypeError, ValueError): - free = total - batch = min( - DEFAULT_STORE_STAGING_BYTES, - max(MIN_STORE_STAGING_BYTES, min(total // 50, free // 10)), - ) - pending = min( - DEFAULT_PENDING_STORE_STAGING_BYTES, - max(batch, min(total // 25, free // 5)), - ) - return batch, pending - - -def store_staging_pool_depth( - buffer_bytes: int, - pending_limit_bytes: int, -) -> int: - """Return fixed store staging pool depth for a 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 startup calculation. - """ - 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, - ) - - -def load_staging_pool_depth( - buffer_bytes: int, - pending_limit_bytes: int, - device: torch.device, - max_inflight: int, - reserve_bytes: int, -) -> int: - """Return fixed load staging depth under memory and inflight limits. - - Args: - buffer_bytes: Capacity of one fixed staging buffer. - pending_limit_bytes: Existing worker staging byte budget. - device: Device used for staging allocation. - max_inflight: Maximum request loads allowed in flight. - reserve_bytes: Free CUDA memory to leave unallocated. - - Returns: - Number of fixed load staging buffers to preallocate. + buffer_bytes = max(DEFAULT_STORE_STAGING_BYTES, local_slot_size) + budget_bytes = DEFAULT_STAGING_BUDGET_BYTES + else: + props = torch.cuda.get_device_properties(device) + total = int(props.total_memory) + try: + free, _ = torch.cuda.mem_get_info(device) + free = int(free) + except (RuntimeError, TypeError, ValueError): + free = total + usable = max(0, free - max(0, reserve_bytes)) + buffer_bytes = max( + local_slot_size, + min( + DEFAULT_STORE_STAGING_BYTES, + max(MIN_STORE_STAGING_BYTES, min(total // 50, free // 10)), + ), + ) + budget_bytes = min( + DEFAULT_STAGING_BUDGET_BYTES, + (2 * total) // 25, + usable, + ) - Async/thread-safety: - Called during worker initialization before request traffic. It may query - CUDA free memory but does not allocate. - """ - depth = min( - 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) - reserve_bytes) - memory_depth = max(1, usable_bytes // buffer_bytes) - return max(1, min(depth, memory_depth)) + minimum = 2 * buffer_bytes + if budget_bytes < minimum: + raise ValueError( + "CUDA staging budget cannot fit one load and one store buffer: " + f"required={minimum} available={budget_bytes}" + ) + total_depth = budget_bytes // buffer_bytes + store_depth = min(2, total_depth // 2) + load_depth = min(max_load_inflight, total_depth - store_depth) + allocated_bytes = buffer_bytes * (load_depth + store_depth) + return buffer_bytes, load_depth, store_depth, allocated_bytes diff --git a/daser/connector/worker_runtime.py b/daser/connector/worker/runtime.py similarity index 95% rename from daser/connector/worker_runtime.py rename to daser/connector/worker/runtime.py index d49ec7d..f4d75cb 100644 --- a/daser/connector/worker_runtime.py +++ b/daser/connector/worker/runtime.py @@ -16,21 +16,19 @@ # First Party from daser.connector.ipc_client import IPCClientSync -from daser.connector.load_pipeline import LoadPipeline from daser.connector.metadata import ( DaserConnectorMeta, ) -from daser.connector.staging import CROSS_LAYER_KV_CACHE_KEY, FUSED_RESTORE_MIN_SLOTS -from daser.connector.store_pipeline import StorePipeline -from daser.connector.worker_memory import ( - DEFAULT_STORE_STAGING_BYTES, +from daser.connector.worker.load import LoadPipeline +from daser.connector.worker.memory import ( FixedCudaStagingPool, - load_staging_pool_depth, - store_staging_pool_depth, + derive_staging_layout, ) -from daser.connector.worker_memory import ( - derive_staging_limits as _derive_staging_limits, +from daser.connector.worker.staging import ( + CROSS_LAYER_KV_CACHE_KEY, + FUSED_RESTORE_MIN_SLOTS, ) +from daser.connector.worker.store import StorePipeline from daser.logging import init_logger from daser.ops.rope_apply import ( apply_rope_delta_to_key_block as _apply_rope_delta_to_key_block, @@ -600,26 +598,21 @@ def _configure_pipelines(self, sample: torch.Tensor) -> int: Async/thread-safety: Called once on the worker thread during KV-cache registration. """ - staging_bytes, pending_limit_bytes = _derive_staging_limits(sample.device) - staging_bytes = max( - staging_bytes or DEFAULT_STORE_STAGING_BYTES, + staging_bytes, load_depth, store_depth, allocated_bytes = derive_staging_layout( + sample.device, self._local_slot_size, + _LOAD_REQUEST_MAX_INFLIGHT, + _LOAD_STAGING_RESERVE_BYTES, ) store_pool = FixedCudaStagingPool( device=sample.device, buffer_bytes=staging_bytes, - depth=store_staging_pool_depth(staging_bytes, pending_limit_bytes), + depth=store_depth, ) load_pool = FixedCudaStagingPool( device=sample.device, buffer_bytes=staging_bytes, - depth=load_staging_pool_depth( - staging_bytes, - pending_limit_bytes, - sample.device, - _LOAD_REQUEST_MAX_INFLIGHT, - _LOAD_STAGING_RESERVE_BYTES, - ), + depth=load_depth, ) self._store_pipeline.configure( kv_caches=self._kv_caches, @@ -647,11 +640,11 @@ def _configure_pipelines(self, sample: torch.Tensor) -> int: rope_is_neox_style=self._rope_is_neox_style, ) logger.info( - "[CONNECTOR] preallocated staging bytes=%d pending=%d " - "store_depth=%d load_depth=%d", + "[CONNECTOR] preallocated staging buffer_bytes=%d total_bytes=%d " + "load_depth=%d store_depth=%d", staging_bytes, - pending_limit_bytes, - store_pool.depth, + allocated_bytes, load_pool.depth, + store_pool.depth, ) return load_pool.depth diff --git a/daser/connector/staging.py b/daser/connector/worker/staging.py similarity index 100% rename from daser/connector/staging.py rename to daser/connector/worker/staging.py diff --git a/daser/connector/store_pipeline.py b/daser/connector/worker/store.py similarity index 91% rename from daser/connector/store_pipeline.py rename to daser/connector/worker/store.py index 13d9782..40e2ec7 100644 --- a/daser/connector/store_pipeline.py +++ b/daser/connector/worker/store.py @@ -14,17 +14,17 @@ from daser.connector.helpers import base_req_id from daser.connector.ipc_client import IPCClientAsync from daser.connector.metadata import ReqStoreSpec, StoreWriteSpan -from daser.connector.staging import ( +from daser.connector.worker.memory import ( + DEFAULT_STORE_STAGING_BYTES, + CudaStagingLease, + FixedCudaStagingPool, +) +from daser.connector.worker.staging import ( CROSS_LAYER_KV_CACHE_KEY, copy_cross_layer_kv_cache_to_staging, copy_kv_cache_to_staging, record_cuda_event, ) -from daser.connector.worker_memory import ( - DEFAULT_STORE_STAGING_BYTES, - CudaStagingLease, - FixedCudaStagingPool, -) from daser.logging import init_logger from daser.transfer.cuda_ipc import ( cuda_allocation_base_and_offset, @@ -193,23 +193,22 @@ def collect_finished(self, finished_req_ids: set[str]) -> set[str]: if not save.finished and save.future is None: continue if save.future is not None and save.future.done(): - save.future.result(timeout=120.0) - finished.add(req_id) - del self._pending_finished_saves[req_id] - has_inflight = any( + try: + save.future.result(timeout=120.0) + finished.add(req_id) + finally: + del self._pending_finished_saves[req_id] + + capacity = self._staging_pool.depth if self._staging_pool is not None else 1 + inflight = sum( save.future is not None for save in self._pending_finished_saves.values() ) - if not has_inflight: - next_save = next( - ( - save - for save in self._pending_finished_saves.values() - if save.finished and save.future is None - ), - None, - ) - if next_save is not None: - self._submit_save(next_save) + for save in self._pending_finished_saves.values(): + if inflight >= capacity: + break + if save.finished and save.future is None: + self._submit_save(save) + inflight += 1 return finished def shutdown(self) -> None: @@ -239,6 +238,7 @@ def shutdown(self) -> None: save = self._pending_finished_saves[req_id] try: self._submit_save(save) + assert save.future is not None save.future.result(timeout=120.0) except BaseException as exc: if first_error is None: @@ -269,11 +269,10 @@ def _submit_save(self, save: _DeferredFinishedSave) -> None: producer_event = record_cuda_event(sample) if sample is not None else None save.future = self._submit(self._store_finished_save(save, producer_event)) - def _stage_finished_save( + def _plan_finished_save( self, save: _DeferredFinishedSave, - producer_event: torch.cuda.Event | None, - ) -> list["StagedStoreBatch"] | None: + ) -> list[tuple[list[int], list[StoreWriteSpan]]]: if not self._kv_caches or self._staging_pool is None: return [] reqs_to_store = { @@ -286,35 +285,20 @@ def _stage_finished_save( ) for req_id, spec in save.reqs_to_store.items() } - batches = build_staging_store_batches( + return build_staging_store_batches( reqs_to_store, self._local_slot_size, max_batch_bytes=self._staging_bytes, ) - if len(batches) > self._staging_pool.available: - return None - staged_batches: list[StagedStoreBatch] = [] - try: - for block_ids, spans in batches: - staged_batches.append( - self._stage_batch(block_ids, spans, producer_event) - ) - except BaseException: - for staged in staged_batches: - staged.lease.release() - raise - return staged_batches async def _store_finished_save( self, save: _DeferredFinishedSave, producer_event: torch.cuda.Event | None, ) -> None: - staged_batches = self._stage_finished_save(save, producer_event) - if staged_batches is None: - raise RuntimeError("store staging capacity cannot satisfy one request") stored_keys: list[str] = [] - for staged in staged_batches: + for block_ids, spans in self._plan_finished_save(save): + staged = self._stage_batch(block_ids, spans, producer_event) try: stored_keys.extend(await self._write_cuda_buffer(staged)) finally: diff --git a/daser/server/ipc/server.py b/daser/server/ipc/server.py index 3571e62..07a6f0a 100644 --- a/daser/server/ipc/server.py +++ b/daser/server/ipc/server.py @@ -622,7 +622,7 @@ def _record_transfer_metrics( "Transfer size per operation in bytes.", buckets=(65536, 262144, 1048576, 4194304, 16777216, 67108864, 268435456), ).observe(nbytes, labels=labels) - self._record_l1_metrics() + self._record_tier_metrics() throughput_gbps = (nbytes / elapsed_s / 1_000_000_000) if elapsed_s > 0 else 0.0 logger.debug( "[IPC] transfer_%s summary backend=%s status=%s bytes=%d " @@ -635,18 +635,21 @@ def _record_transfer_metrics( throughput_gbps, ) - def _record_l1_metrics(self) -> None: - """Publish L1 cache hit/miss counters and usage gauges.""" + def _record_tier_metrics(self) -> None: + """Publish L1 cache metrics and the cumulative L2 read counter.""" transfer = self._transfer if transfer is None: return stats = transfer.stats current_hits = stats.l1_hits current_misses = stats.l1_misses + current_l2_reads = stats.l2_reads prev_hits = getattr(self, "_prev_l1_hits", 0) prev_misses = getattr(self, "_prev_l1_misses", 0) + prev_l2_reads = getattr(self, "_prev_l2_reads", 0) delta_hits = current_hits - prev_hits delta_misses = current_misses - prev_misses + delta_l2_reads = current_l2_reads - prev_l2_reads if delta_hits > 0: self._metrics.counter("daser_l1_hits_total", "L1 memory cache hits.").inc( delta_hits @@ -655,8 +658,13 @@ def _record_l1_metrics(self) -> None: self._metrics.counter( "daser_l1_misses_total", "L1 memory cache misses." ).inc(delta_misses) + if delta_l2_reads > 0: + self._metrics.counter( + "daser_l2_reads_total", "Reads served from the L2 storage tier." + ).inc(delta_l2_reads) self._prev_l1_hits = current_hits self._prev_l1_misses = current_misses + self._prev_l2_reads = current_l2_reads l1_used = transfer.l1_bytes_used l1_capacity = int(self._runtime_config.get("l1_size_bytes", 0)) self._metrics.gauge("daser_l1_bytes_used", "L1 memory cache bytes in use.").set( diff --git a/docs/design/architecture.md b/docs/design/architecture.md index 81408a1..a1f10ac 100644 --- a/docs/design/architecture.md +++ b/docs/design/architecture.md @@ -61,8 +61,8 @@ graph TB subgraph vllm["vLLM 进程"] VAPI["OpenAI-compatible HTTP API"] DC["DaserConnector
KVConnectorBase_V1"] - SCHED["scheduler.py
SCHEDULER role"] - WORKER["worker.py
WORKER role"] + SCHED["scheduler/adapter.py
SCHEDULER role"] + WORKER["worker/adapter.py
WORKER role"] LIFE["RequestLifecycle
pending state + IPC orchestration"] RUNTIME["WorkerRuntime
KV/meta/completion ownership"] LOAD["LoadPipeline
queue + load loop"] diff --git a/docs/design/components.md b/docs/design/components.md index 44edce4..73ae4e1 100644 --- a/docs/design/components.md +++ b/docs/design/components.md @@ -5,12 +5,12 @@ | 组件 | 进程 | 职责 | |------|------|------| | `DaserConnector` | vLLM | vLLM `KVConnectorBase_V1` 入口；保留在 `daser/connector/daser_connector.py` 供 `kv_connector_module_path` 加载 | -| `SchedulerConnectorMixin` | vLLM scheduler | `daser/connector/scheduler.py`；适配 vLLM scheduler hooks，不拥有请求 lifecycle 状态 | +| `SchedulerConnectorMixin` | vLLM scheduler | `daser/connector/scheduler/adapter.py`；适配 vLLM scheduler hooks，不拥有请求 lifecycle 状态 | | `RequestLifecycle` | vLLM scheduler | 集中 lookup、pending load/store/alloc/async-save、slot 分配、preemption、completion 和 connector metadata 构造 | -| `WorkerConnectorMixin` | vLLM worker | `daser/connector/worker.py`；适配 vLLM worker hooks 和启动期 kernel warmup，不拥有 pipeline 状态 | +| `WorkerConnectorMixin` | vLLM worker | `daser/connector/worker/adapter.py`；适配 vLLM worker hooks 和启动期 kernel warmup，不拥有 pipeline 状态 | | `WorkerRuntime` | vLLM worker | 集中 KV layout、step metadata、load/store completion 和 shutdown，组合两个独立 pipeline | | `LoadPipeline` / `StorePipeline` | vLLM worker | 分别拥有 load/store event loop、IPC client、staging lease、future、backpressure 和 transfer plan | -| worker memory module | vLLM worker | 负责 load/store 共用的 bounded CUDA staging pool、indexed lease 和 capacity derivation | +| worker memory module | vLLM worker | 负责 load/store 独立 fixed pool、indexed lease，以及约 6 GiB combined budget 的单点推导 | | staging tensor module | vLLM worker | 负责 slot-major KV copy、CUDA producer synchronization 和 RoPE restore，不拥有 request/IPC 状态 | | `IPCClientSync` | vLLM scheduler | 阻塞式 Unix socket 客户端，用于 `get_runtime_config`、`lookup`、`alloc_chunk` | | `IPCClientAsync` | vLLM worker | asyncio Unix socket 客户端，用于 `transfer_store`、`transfer_load`、`commit_chunks` | diff --git a/docs/optimizations/3_server_managed_transfer.md b/docs/optimizations/3_server_managed_transfer.md index 3f5acac..8d52a38 100644 --- a/docs/optimizations/3_server_managed_transfer.md +++ b/docs/optimizations/3_server_managed_transfer.md @@ -103,7 +103,7 @@ A same-process pointer path is used only for local unit-test and benchmark cases where producer and consumer PIDs match. Worker-side staging is shared by both GDS and iouring modes through -`daser.connector.staging`. `register_kv_caches` creates a bounded +`daser.connector.worker.staging`. `register_kv_caches` creates a bounded `CudaStagingPool` and preallocates one reusable buffer so the hot path does not pay a fresh CUDA allocation for the common batch size. Store batches keep their lease alive until the background transfer future completes; load batches release diff --git a/docs/optimizations/4_rope_apply_compile.md b/docs/optimizations/4_rope_apply_compile.md index 86bbe53..f2174e1 100644 --- a/docs/optimizations/4_rope_apply_compile.md +++ b/docs/optimizations/4_rope_apply_compile.md @@ -30,7 +30,7 @@ delta implementation there: TileLang operator entrypoint. - Naive eager RoPE is no longer part of production code; it is kept only inside the temporary micro benchmark and tests as a correctness oracle. -- `daser.connector.staging.apply_rope_delta_to_key_block()` remains the +- `daser.connector.worker.staging.apply_rope_delta_to_key_block()` remains the connector-facing helper and delegates into `daser.ops`. The production path is TileLang-only. TileLang import/compile/runtime failures diff --git a/tests/connector/test_connector_layout.py b/tests/connector/test_connector_layout.py index ef18510..d60dd69 100644 --- a/tests/connector/test_connector_layout.py +++ b/tests/connector/test_connector_layout.py @@ -23,11 +23,11 @@ def test_connector_entrypoint_delegates_scheduler_and_worker_methods() -> None: ) assert ( inspect.getmodule(DaserConnector.get_num_new_matched_tokens).__name__ - == "daser.connector.scheduler" + == "daser.connector.scheduler.adapter" ) assert ( inspect.getmodule(DaserConnector.start_load_kv).__name__ - == "daser.connector.worker" + == "daser.connector.worker.adapter" ) diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index 955f605..d33ea9d 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -18,61 +18,58 @@ rolling_prefix_key, rolling_prefix_keys, ) -from daser.connector.load_pipeline import ( - build_load_copy_runs as _build_load_copy_runs, -) -from daser.connector.load_pipeline import ( - build_load_read_batches as _build_load_read_batches, -) -from daser.connector.load_pipeline import ( - build_load_read_plan as _build_load_read_plan, -) from daser.connector.metadata import ( DaserConnectorMeta, ReqLoadSpec, ReqStoreSpec, StoreWriteSpan, ) -from daser.connector.request_lifecycle import RequestLifecycle -from daser.connector.reuse import PrefixReuseStrategy -from daser.connector.scheduler_planning import ( +from daser.connector.scheduler.lifecycle import RequestLifecycle +from daser.connector.scheduler.planning import ( _block_ids_for_chunk, _contiguous_prefix_tokens, _trim_chunk_to_external_window, ) -from daser.connector.staging import ( +from daser.connector.scheduler.reuse import PrefixReuseStrategy +from daser.connector.worker.load import ( + build_load_copy_runs as _build_load_copy_runs, +) +from daser.connector.worker.load import ( + build_load_read_batches as _build_load_read_batches, +) +from daser.connector.worker.load import ( + build_load_read_plan as _build_load_read_plan, +) +from daser.connector.worker.memory import ( + DEFAULT_STAGING_BUDGET_BYTES, + DEFAULT_STORE_STAGING_BYTES, + MIN_STORE_STAGING_BYTES, + FixedCudaStagingPool, + derive_staging_layout, +) +from daser.connector.worker.runtime import WorkerRuntime +from daser.connector.worker.staging import ( DEFAULT_ROPE_DELTA_SCALE, ) -from daser.connector.staging import ( +from daser.connector.worker.staging import ( apply_rope_delta_to_key_block as _apply_rope_delta_to_key_block, ) -from daser.connector.staging import ( +from daser.connector.worker.staging import ( copy_staging_to_kv_cache as _copy_staging_to_kv_cache, ) -from daser.connector.staging import ( +from daser.connector.worker.staging import ( record_cuda_event as _record_cuda_event, ) -from daser.connector.staging import ( +from daser.connector.worker.staging import ( synchronize_cuda_tensor as _synchronize_cuda_tensor, ) -from daser.connector.store_pipeline import ( +from daser.connector.worker.store import ( StagedStoreBatch, StorePipeline, ) -from daser.connector.store_pipeline import ( +from daser.connector.worker.store import ( build_staging_store_batches as _build_staging_store_batches, ) -from daser.connector.worker_memory import ( - DEFAULT_PENDING_STORE_STAGING_BYTES, - DEFAULT_STORE_STAGING_BYTES, - MIN_STORE_STAGING_BYTES, - FixedCudaStagingPool, - load_staging_pool_depth, -) -from daser.connector.worker_memory import ( - derive_staging_limits as _derive_staging_limits, -) -from daser.connector.worker_runtime import WorkerRuntime BLOCK_TOKENS = 4 NUM_LAYERS = 2 @@ -338,23 +335,29 @@ class DummyVLLMConfig: ) -def test_load_staging_pool_depth_respects_available_cuda_memory(monkeypatch) -> None: - """Load staging depth should not preallocate past available CUDA headroom.""" +def test_staging_layout_respects_available_cuda_headroom(monkeypatch) -> None: + """Combined staging pools stay within available CUDA headroom.""" + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda device: SimpleNamespace(total_memory=80 << 30), + ) 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"), - max_inflight=8, + buffer_bytes, load_depth, store_depth, allocated = derive_staging_layout( + torch.device("cuda"), + local_slot_size=64 << 20, + max_load_inflight=8, reserve_bytes=1 << 30, ) - assert depth == 2 + assert (buffer_bytes, load_depth, store_depth) == ((4 << 30) // 10, 5, 2) + assert allocated == buffer_bytes * 7 + assert allocated <= (4 << 30) - (1 << 30) def test_worker_transfer_ready_allows_skip_l2_without_store_path() -> None: @@ -457,7 +460,7 @@ def close(self) -> None: return monkeypatch.setattr( - "daser.connector.worker_runtime.IPCClientSync", + "daser.connector.worker.runtime.IPCClientSync", DummyIPCClient, ) connector = WorkerRuntime.__new__(WorkerRuntime) @@ -1417,7 +1420,7 @@ def test_restore_cross_layer_kv_cache_table_tilelang_matches_reference_cuda(): def test_register_kv_caches_warms_dynamic_rope_apply_once(monkeypatch): - from daser.connector import worker_runtime as worker + from daser.connector.worker import runtime as worker class Probe(WorkerRuntime): def __init__(self) -> None: @@ -1438,11 +1441,6 @@ def _configure_pipelines(self, sample: torch.Tensor) -> int: return 1 calls = [] - monkeypatch.setattr( - worker, - "_derive_staging_limits", - lambda device: (4096, 8192), - ) monkeypatch.setattr( worker, "_warm_rope_apply_backends", @@ -1468,7 +1466,7 @@ def _configure_pipelines(self, sample: torch.Tensor) -> int: def test_register_cross_layers_kv_cache_preserves_layer_order(monkeypatch): """Worker registration keeps vLLM layer names for slot-major staging.""" - from daser.connector import worker_runtime as worker + from daser.connector.worker import runtime as worker class Group: layer_names = ["layer.0", "layer.1"] @@ -1508,11 +1506,6 @@ def registration_state(self): self._slot_size, ) - monkeypatch.setattr( - worker, - "_derive_staging_limits", - lambda device: (4096, 8192), - ) monkeypatch.setattr(worker, "_warm_rope_apply_backends", lambda **kwargs: None) kv_cache = torch.zeros((8, 2, 2, 4, 2, 8), dtype=torch.float32) @@ -1767,7 +1760,7 @@ def __init__(self) -> None: calls: list[tuple[list[int], int, str | None, int]] = [] monkeypatch.setattr( - "daser.connector.reuse.rolling_prefix_keys", + "daser.connector.scheduler.reuse.rolling_prefix_keys", lambda tokens, block_tokens, initial_key=None, start_slot=0, **kwargs: ( calls.append((list(tokens), block_tokens, initial_key, start_slot)) or rolling_prefix_keys( @@ -2520,7 +2513,7 @@ def fake_rope( kv_block[:, :, 0, ..., :rotary_dim].add_(10.0) monkeypatch.setattr( - "daser.connector.staging.apply_rope_delta_to_kv_key_block", + "daser.connector.worker.staging.apply_rope_delta_to_kv_key_block", fake_rope, ) @@ -2572,7 +2565,7 @@ def fail_rope(*args, **kwargs): raise AssertionError("RoPE should not run for pos_offset=0") monkeypatch.setattr( - "daser.connector.staging.apply_rope_delta_to_key_block", + "daser.connector.worker.staging.apply_rope_delta_to_key_block", fail_rope, ) @@ -2617,7 +2610,7 @@ def fake_rope(kv_block, delta, rope_base, rotary_dim, is_neox_style): kv_block[:, :, 0, ..., :rotary_dim].add_(10) monkeypatch.setattr( - "daser.connector.staging._apply_rope_delta_with_tables", + "daser.connector.worker.staging._apply_rope_delta_with_tables", fake_rope, ) @@ -2665,7 +2658,7 @@ def fake_rope( kv_block[:, :, 0, ..., :rotary_dim].add_(3.0) monkeypatch.setattr( - "daser.connector.staging._apply_rope_delta_with_tables", + "daser.connector.worker.staging._apply_rope_delta_with_tables", fake_rope, ) @@ -2725,14 +2718,14 @@ def fail_rope(*args, **kwargs): raise AssertionError("legacy RoPE path should not run") monkeypatch.setattr( - "daser.connector.staging.apply_rope_delta_to_kv_key_block_table", + "daser.connector.worker.staging.apply_rope_delta_to_kv_key_block_table", fake_table_rope, ) monkeypatch.setattr( - "daser.connector.staging._apply_rope_delta_to_kv_key_block", + "daser.connector.worker.staging._apply_rope_delta_to_kv_key_block", fail_rope, ) - monkeypatch.setattr("daser.connector.staging._rope_table_cache", {}) + monkeypatch.setattr("daser.connector.worker.staging._rope_table_cache", {}) copies = _copy_staging_to_kv_cache( staging=staging, @@ -2773,7 +2766,7 @@ def fake_table_restore( return True monkeypatch.setattr( - "daser.connector.staging._restore_cross_layer_with_tables", + "daser.connector.worker.staging._restore_cross_layer_with_tables", fake_table_restore, ) @@ -2812,7 +2805,7 @@ def failing_restore(*args, **kwargs): raise RuntimeError("backend unavailable") monkeypatch.setattr( - "daser.connector.staging._restore_cross_layer_with_tables", + "daser.connector.worker.staging._restore_cross_layer_with_tables", failing_restore, ) @@ -2882,7 +2875,7 @@ def test_build_staging_store_batches_uses_spec_file_offset(): @pytest.mark.asyncio async def test_store_cuda_export_selects_staged_buffer_device(monkeypatch) -> None: """Background CUDA IPC export must select the TP rank's staged device.""" - from daser.connector import store_pipeline as store_module + from daser.connector.worker import store as store_module selected_devices: list[torch.device] = [] transferred: list[dict] = [] @@ -2929,8 +2922,8 @@ def test_tensor_parallel_rank_lanes_are_contiguous_and_disjoint() -> None: assert rank_0 + 2 * local_slot_size <= rank_1 -def test_derive_staging_limits_scale_with_vram(monkeypatch): - """GPU staging caps consider device size and currently free VRAM.""" +def test_derive_staging_layout_scales_with_vram(monkeypatch): + """One budget preserves balanced pools and the explicit 6 GiB ceiling.""" class Props: def __init__(self, total_memory: int) -> None: @@ -2946,15 +2939,15 @@ def __init__(self, total_memory: int) -> None: "mem_get_info", lambda device=None: (12 << 30, 24 << 30), ) - small_batch, small_pending = _derive_staging_limits(torch.device("cuda")) + small_batch, small_load, small_store, small_total = derive_staging_layout( + torch.device("cuda"), 64 << 20, 8, 1 << 30 + ) assert small_batch == max( MIN_STORE_STAGING_BYTES, min((24 << 30) // 50, (12 << 30) // 10), ) - assert small_pending == max( - small_batch, - min((24 << 30) // 25, (12 << 30) // 5), - ) + assert (small_load, small_store) == (2, 2) + assert small_total == small_batch * 4 monkeypatch.setattr( torch.cuda, @@ -2966,18 +2959,24 @@ def __init__(self, total_memory: int) -> None: "mem_get_info", lambda device=None: (64 << 30, 80 << 30), ) - large_batch, large_pending = _derive_staging_limits(torch.device("cuda")) + large_batch, large_load, large_store, large_total = derive_staging_layout( + torch.device("cuda"), 64 << 20, 8, 1 << 30 + ) assert large_batch == DEFAULT_STORE_STAGING_BYTES - assert large_pending == DEFAULT_PENDING_STORE_STAGING_BYTES + assert (large_load, large_store) == (2, 2) + assert large_total == DEFAULT_STAGING_BUDGET_BYTES monkeypatch.setattr( torch.cuda, "mem_get_info", lambda device=None: (8 << 30, 80 << 30), ) - tight_batch, tight_pending = _derive_staging_limits(torch.device("cuda")) + tight_batch, tight_load, tight_store, tight_total = derive_staging_layout( + torch.device("cuda"), 64 << 20, 8, 1 << 30 + ) assert tight_batch == (8 << 30) // 10 - assert tight_pending == (8 << 30) // 5 + assert (tight_load, tight_store) == (5, 2) + assert tight_total == tight_batch * 7 def test_store_cuda_staging_pool_reuses_preallocated_buffer(): diff --git a/tests/connector/test_gds_transfer.py b/tests/connector/test_gds_transfer.py index 1378028..3394350 100644 --- a/tests/connector/test_gds_transfer.py +++ b/tests/connector/test_gds_transfer.py @@ -89,7 +89,7 @@ def test_missing_file_raises(tmp_path): def test_fixed_staging_pool_reuses_two_preallocated_buffers() -> None: """Fixed staging pool reuses bounded preallocated buffers.""" # First Party - from daser.connector.worker_memory import FixedCudaStagingPool + from daser.connector.worker.memory import FixedCudaStagingPool pool = FixedCudaStagingPool( device=torch.device("cpu"), @@ -114,7 +114,7 @@ def test_fixed_staging_pool_reuses_two_preallocated_buffers() -> None: 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.worker_memory import FixedCudaStagingPool + from daser.connector.worker.memory import FixedCudaStagingPool pool = FixedCudaStagingPool( device=torch.device("cpu"), @@ -139,7 +139,7 @@ def release_first() -> None: def test_fixed_staging_pool_rejects_oversized_request() -> None: """Fixed staging pool rejects requests larger than one buffer.""" # First Party - from daser.connector.worker_memory import FixedCudaStagingPool + from daser.connector.worker.memory import FixedCudaStagingPool pool = FixedCudaStagingPool( device=torch.device("cpu"), diff --git a/tests/connector/test_worker_pipelines.py b/tests/connector/test_worker_pipelines.py new file mode 100644 index 0000000..bfc541c --- /dev/null +++ b/tests/connector/test_worker_pipelines.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from concurrent.futures import Future +import time +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from daser.connector.metadata import ReqLoadSpec, ReqStoreSpec +from daser.connector.worker.load import LoadPipeline +from daser.connector.worker.memory import FixedCudaStagingPool +from daser.connector.worker.store import StagedStoreBatch, StorePipeline + + +class _ManualFuture: + def __init__(self) -> None: + self.complete = False + + def done(self) -> bool: + return self.complete + + def result(self, timeout: float) -> None: + del timeout + + +def _store_spec(key: str, blocks: list[int]) -> ReqStoreSpec: + return ReqStoreSpec(key, 0, len(blocks), blocks, 0, len(blocks)) + + +def test_store_pipeline_defers_and_dispatches_up_to_pool_depth() -> None: + pipeline = StorePipeline.__new__(StorePipeline) + pipeline._pending_finished_saves = {} # noqa: SLF001 + pipeline._staging_pool = SimpleNamespace(depth=2) # noqa: SLF001 + submitted: list[_ManualFuture] = [] + + def submit(save: Any) -> None: + future = _ManualFuture() + save.future = future + submitted.append(future) + + pipeline._submit_save = submit # type: ignore[method-assign] # noqa: SLF001 + pipeline.queue_finished( + {req: _store_spec(req, [index]) for index, req in enumerate(("a", "b", "c"))}, + {"a", "b", "c"}, + ) + + assert pipeline.collect_finished(set()) == set() + assert submitted == [] + assert pipeline.collect_finished({"a", "b", "c"}) == set() + assert len(submitted) == 2 + + submitted[0].complete = True + assert pipeline.collect_finished(set()) == {"a"} + assert len(submitted) == 3 + submitted[1].complete = True + submitted[2].complete = True + assert pipeline.collect_finished(set()) == {"b", "c"} + + +def test_store_pipeline_streams_request_larger_than_pool_depth() -> None: + pipeline = StorePipeline.__new__(StorePipeline) + pipeline._pending_finished_saves = {} # noqa: SLF001 + pipeline._staging_pool = SimpleNamespace(depth=1) # noqa: SLF001 + pipeline._kv_caches = {"layer": torch.empty(1)} # noqa: SLF001 + pipeline._local_slot_size = 16 # noqa: SLF001 + pipeline._rank_stride_bytes = 0 # noqa: SLF001 + pipeline._tp_rank = 0 # noqa: SLF001 + pipeline._tp_size = 1 # noqa: SLF001 + pipeline._staging_bytes = 32 # noqa: SLF001 + released: list[int] = [] + writes: list[int] = [] + + class Lease: + view = torch.empty(1) + + def release(self) -> None: + released.append(1) + + class Client: + async def commit_chunks(self, keys: list[str], **kwargs: int) -> None: + assert keys == ["large"] + assert kwargs == {"tp_rank": 0, "tp_size": 1} + + def stage(block_ids: list[int], spans: list[Any], event: Any) -> StagedStoreBatch: + del event + return StagedStoreBatch(torch.empty(1), spans, Lease()) + + async def write(staged: StagedStoreBatch) -> list[str]: + writes.append(sum(span.nbytes for span in staged.spans) // 16) + return ["large"] + + def submit(coro: Any) -> Future[None]: + future: Future[None] = Future() + try: + asyncio.run(coro) + future.set_result(None) + except BaseException as exc: + future.set_exception(exc) + return future + + pipeline._client = Client() # type: ignore[assignment] # noqa: SLF001 + pipeline._stage_batch = stage # type: ignore[method-assign] # noqa: SLF001 + pipeline._write_cuda_buffer = write # type: ignore[method-assign] # noqa: SLF001 + pipeline._submit = submit # type: ignore[method-assign] # noqa: SLF001 + pipeline._submit_save = StorePipeline._submit_save.__get__(pipeline) # noqa: SLF001 + pipeline._kv_caches = {} # noqa: SLF001 + pipeline.queue_finished({"large": _store_spec("large", [0, 1, 2, 3, 4])}, {"large"}) + pipeline._kv_caches = {"layer": torch.empty(1)} # noqa: SLF001 + + assert pipeline.collect_finished({"large"}) == set() + assert pipeline.collect_finished(set()) == {"large"} + assert writes == [2, 2, 1] + assert len(released) == 3 + + +class _LoadClient: + def __init__(self, fail_offset: int | None = None) -> None: + self.calls: list[int] = [] + self.fail_offset = fail_offset + + async def transfer_load_registered_cuda(self, **kwargs: Any) -> dict[str, Any]: + offset = int(kwargs["spans"][0]["file_offset"]) + self.calls.append(offset) + if offset == self.fail_offset: + raise RuntimeError("load failed") + return { + "transfer_open_ms": 1.0, + "transfer_load_ms": 2.0, + "transfer_sync_ms": 3.0, + "transfer_stats_delta": {"l1_hits": 4, "l1_misses": 5, "l2_reads": 6}, + } + + async def close(self) -> None: + return None + + +def _load_spec(key: str, blocks: list[int], offset: int = 0) -> ReqLoadSpec: + return ReqLoadSpec(key, offset // 16, len(blocks), blocks, offset, len(blocks)) + + +def _load_pipeline( + monkeypatch: pytest.MonkeyPatch, client: _LoadClient +) -> LoadPipeline: + monkeypatch.setattr( + "daser.connector.worker.load.copy_staging_to_kv_cache", + lambda **kwargs: 1, + ) + pipeline = LoadPipeline("unused.sock", client_count=2) + pipeline._clients = [client, client] # type: ignore[assignment] # noqa: SLF001 + pipeline.configure( + kv_caches={"layer": torch.empty(1)}, + layer_names=["layer"], + local_slot_size=16, + rank_stride_bytes=0, + tp_rank=0, + staging_pool=FixedCudaStagingPool(torch.device("cpu"), 32, 2), + load_key_scale=1.0, + load_value_scale=1.0, + rope_delta_scale=1.0, + rope_base=10000.0, + rope_rotary_dim=0, + rope_is_neox_style=True, + ) + pipeline._staging_registered = True # noqa: SLF001 + return pipeline + + +def _wait_finished(pipeline: LoadPipeline, expected: set[str]) -> set[str]: + deadline = time.monotonic() + 2.0 + finished: set[str] = set() + while time.monotonic() < deadline and finished != expected: + finished.update(pipeline.collect_finished()) + time.sleep(0.005) + return finished + + +def test_load_pipeline_handles_empty_and_multibatch_requests( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = _LoadClient() + pipeline = _load_pipeline(monkeypatch, client) + try: + pipeline.start( + { + "empty": _load_spec("empty", []), + "large:load:0": _load_spec("large", [0, 1, 2, 3, 4]), + } + ) + assert _wait_finished(pipeline, {"empty", "large"}) == {"empty", "large"} + assert client.calls == [0, 32, 64] + finally: + pipeline.shutdown() + + +def test_load_failure_invalidates_only_failed_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = _LoadClient(fail_offset=0) + pipeline = _load_pipeline(monkeypatch, client) + try: + pipeline.start( + { + "bad": _load_spec("bad", [7], 0), + "good": _load_spec("good", [8], 16), + } + ) + assert _wait_finished(pipeline, {"bad", "good"}) == {"bad", "good"} + assert pipeline.take_invalid_block_ids() == {7} + assert sorted(client.calls) == [0, 16] + finally: + pipeline.shutdown() diff --git a/tests/server/test_ipc_server.py b/tests/server/test_ipc_server.py index f217b37..1994f79 100644 --- a/tests/server/test_ipc_server.py +++ b/tests/server/test_ipc_server.py @@ -3,6 +3,7 @@ # Standard import asyncio import os +from types import SimpleNamespace from typing import Any # Third Party @@ -19,7 +20,7 @@ from daser.server.doc_registry import DocRegistry from daser.server.ipc import IPCServer from daser.server.metadata_store import MetadataStore -from daser.transfer.base import TransferLayer +from daser.transfer.base import TransferLayer, TransferStats SLOT_SIZE = 4096 BLOCK_TOKENS = 4 @@ -227,6 +228,30 @@ async def test_ipc_server_records_operation_metrics(tmp_path) -> None: ) +def test_ipc_server_records_tier_counter_deltas(tmp_path) -> None: + """Tier metrics publish monotonic L1 and L2 deltas exactly once.""" + registry = MetricsRegistry() + server = IPCServer( + str(tmp_path / "test.sock"), + make_core(), + make_runtime_config(tmp_path), + metrics_registry=registry, + ) + transfer = SimpleNamespace( + stats=TransferStats(l1_hits=2, l1_misses=3, l2_reads=4), + l1_bytes_used=1024, + ) + server._transfer = transfer # type: ignore[assignment] # noqa: SLF001 + + server._record_tier_metrics() # noqa: SLF001 + server._record_tier_metrics() # noqa: SLF001 + + rendered = registry.render_prometheus() + assert "daser_l1_hits_total 2.0" in rendered + assert "daser_l1_misses_total 3.0" in rendered + assert "daser_l2_reads_total 4.0" in rendered + + @pytest.mark.asyncio async def test_ipc_server_records_external_prefix_cache_metrics(tmp_path) -> None: """IPC can publish vLLM-equivalent external prefix cache counters.""" diff --git a/tests/unit/test_benchmark_unified_utils.py b/tests/unit/test_benchmark_unified_utils.py index 4218bfc..564c955 100644 --- a/tests/unit/test_benchmark_unified_utils.py +++ b/tests/unit/test_benchmark_unified_utils.py @@ -8,6 +8,7 @@ from pathlib import Path import subprocess import sys +from types import SimpleNamespace from typing import Any import pytest @@ -2466,6 +2467,91 @@ async def fake_collect( assert hit_rate == 0.75 +def test_daser_evict_gate_requires_l1_and_l2_activity() -> None: + """One warm evict phase must include positive deltas from both tiers.""" + vllm_bench._require_daser_evict_tier_activity( # noqa: SLF001 + { + "backend_prometheus": { + "daser_l1_hits_total": 3.0, + "daser_l2_reads_total": 2.0, + } + } + ) + for counters in ( + {}, + {"daser_l1_hits_total": 1.0}, + {"daser_l1_hits_total": 1.0, "daser_l2_reads_total": 0.0}, + ): + with pytest.raises(RuntimeError, match="must exercise both tiers"): + vllm_bench._require_daser_evict_tier_activity( # noqa: SLF001 + {"backend_prometheus": counters} + ) + + +def test_daser_evict_warm_phase_primes_same_seed_prefix( + tmp_path: Path, + monkeypatch, +) -> None: + """Evict warm metrics include a short LRU perturbation and full phase.""" + commands: list[list[str]] = [] + + async def fake_collect( + manifest: Any, + before_metrics: dict[str, Any] | None = None, + ) -> dict[str, Any]: + del manifest + if before_metrics is None: + return {"backend_prometheus": {}} + return { + "backend_prometheus": { + "daser_l1_hits_total": 1.0, + "daser_l2_reads_total": 1.0, + }, + "hit_ratios": {}, + } + + monkeypatch.setattr(vllm_bench, "collect_phase_metrics", fake_collect) + args = RunBenchArgs(model="model", bench_num_prompts=10) + manifest = SimpleNamespace( + endpoints={"vllm": SimpleNamespace(url="http://127.0.0.1:8001")} + ) + + metrics, _hit_rate = vllm_bench._run_daser_evict_warm_phase( # noqa: SLF001 + args, + manifest, + tmp_path / "warm.json", + run_command=commands.append, + ) + + assert [command[command.index("--num-prompts") + 1] for command in commands] == [ + "2", + "10", + ] + assert [command[command.index("--seed") + 1] for command in commands] == [ + "42", + "42", + ] + assert metrics["backend_prometheus"]["daser_l1_hits_total"] == 1.0 + + +def test_daser_drain_failure_aborts_benchmark(monkeypatch) -> None: + """A failed cold-to-warm barrier is a benchmark failure.""" + monkeypatch.setattr( + vllm_bench.httpx, + "post", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("timeout")), + ) + manifest = SimpleNamespace( + endpoints={"daser": SimpleNamespace(url="http://127.0.0.1:9999")} + ) + + with pytest.raises(RuntimeError, match="timeout"): + vllm_bench._drain_daser( # noqa: SLF001 + manifest, + print_kv=lambda key, value: None, + ) + + def test_run_bench_vllm_bench_entrypoint_runs_openai_rows( tmp_path: Path, monkeypatch, diff --git a/tests/unit/test_block_aligned_bug.py b/tests/unit/test_block_aligned_bug.py index b38bc53..b812918 100644 --- a/tests/unit/test_block_aligned_bug.py +++ b/tests/unit/test_block_aligned_bug.py @@ -11,7 +11,7 @@ from daser.connector.helpers import PendingStore, hash_tokens # First Party -from daser.connector.request_lifecycle import RequestLifecycle +from daser.connector.scheduler.lifecycle import RequestLifecycle class TestTokeniseAndTruncateBug: @@ -428,7 +428,7 @@ def __init__(self): ) def test_trim_external_window_keeps_full_block_for_lmcache_style_minus_one(self): - from daser.connector.scheduler_planning import ( + from daser.connector.scheduler.planning import ( _trim_chunk_to_external_window, ) diff --git a/tests/unit/test_chunk_reuse_scheduler.py b/tests/unit/test_chunk_reuse_scheduler.py index 79ad68e..bc6434d 100644 --- a/tests/unit/test_chunk_reuse_scheduler.py +++ b/tests/unit/test_chunk_reuse_scheduler.py @@ -4,7 +4,7 @@ from typing import Any # First Party -from daser.connector.request_lifecycle import RequestLifecycle +from daser.connector.scheduler.lifecycle import RequestLifecycle BLOCK_TOKENS = 4 diff --git a/tests/unit/test_skip_save_flag.py b/tests/unit/test_skip_save_flag.py index 6751153..3468b1c 100644 --- a/tests/unit/test_skip_save_flag.py +++ b/tests/unit/test_skip_save_flag.py @@ -17,7 +17,7 @@ from daser.connector.helpers import PendingStore, hash_tokens # First Party -from daser.connector.request_lifecycle import RequestLifecycle +from daser.connector.scheduler.lifecycle import RequestLifecycle BLOCK_TOKENS = 16 From 92353965778886fd4d98c38e5a5dfde2b75d3923 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Sun, 12 Jul 2026 20:43:21 +0800 Subject: [PATCH 4/4] fix(tests): skip worker pipelines without optional deps - declare torch, vLLM, and CuPy requirements before connector imports - keep worker pipeline behavior tests active in full development environments --- tests/connector/test_worker_pipelines.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/connector/test_worker_pipelines.py b/tests/connector/test_worker_pipelines.py index bfc541c..686cbe7 100644 --- a/tests/connector/test_worker_pipelines.py +++ b/tests/connector/test_worker_pipelines.py @@ -7,6 +7,11 @@ from typing import Any import pytest + +pytest.importorskip("torch") +pytest.importorskip("vllm") +pytest.importorskip("cupy") + import torch from daser.connector.metadata import ReqLoadSpec, ReqStoreSpec