# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Runs WriteTasks and the step device→host staging pool. The manager owns request/slot identity and when to submit. This class owns the staging pool, the GPU-byte budget, the copy queues, and the single committer that writes into the CPU block pool. Two device→host paths — the step path does not write `chunk.host` in the committer: JOIN_NEXT_STEP save already launched a whole-step device→host into a staging slot and set `chunk.host` as views. Committer waits that `step_d2h_event`, then writes the pool. JOIN_ON_FINISH committer copies the device clone → owned host tensors, then writes the pool. Async: high-priority then low-priority queues, then pool write. Eager: submit() does wait+write inline (CPU tests / no CUDA). """ from __future__ import annotations import logging import threading import time from collections import deque from collections.abc import Callable from dataclasses import dataclass, field from enum import Enum, auto from typing import Literal, NamedTuple import torch from vllm_omni.core.prefix_cache.block_pool import PrefixBlockPool from vllm_omni.core.prefix_cache.interface import ( OmniPrefixCacheStagingTimeoutError, OmniPrefixCacheUnmatchError, PrefixCacheConfig, ReqId, TensorName, Tid, WriteSchedule, ) logger = logging.getLogger(__name__) class StagingBufferHolder(NamedTuple): """One owner of a staging-buffer slot. The slot is free when none remain. Not a buffer state — concurrent owners share the same slot: - for_step: claimed at save, released when materialize/discard consumes the ctx - for_task: bound before WriteTask submit, released when that task completes """ kind: Literal["step", "task"] owner_id: int @classmethod def for_step(cls, step_id: int) -> StagingBufferHolder: return cls("step", step_id) @classmethod def for_task(cls, tid: int) -> StagingBufferHolder: return cls("task", tid) @dataclass class _BudgetTicket: """GPU-byte-budget charge for one device clone shared by several writes. Charged once at ``reserve()``; freed when the last pinned tid releases. Immediate: the whole-step clone, pinned by that step's JOIN_NEXT_STEP tasks. Deferred: the shared deferred clone, pinned by every task a chunk of it was appended to. Does not store tensors. """ nbytes: int tids: set[Tid] = field(default_factory=set) _freed: bool = False def release(self, tid: Tid) -> int: """Drop `tid`; return the bytes to uncharge (nbytes once, when the last holder leaves; 0 otherwise). Idempotent per tid.""" self.tids.discard(tid) if self.tids or self._freed: return 0 self._freed = True return self.nbytes @dataclass class _WriteChunk: """One save's slots + tensors appended onto a WriteTask. Controller identity is the parent ``tid``, not a request. Deferred writes grow by appending one chunk per save. """ slots_cpu: torch.Tensor # int64 flat row ids, in token order tensors: dict[TensorName, torch.Tensor] # freeze (device or eager CPU) host: dict[TensorName, torch.Tensor] = field(default_factory=dict) # host view of rows [0:n) # GPU-byte charge of the clone these tensors view; shared across the # chunks cut from one save. None for chunks with no device clone. budget: _BudgetTicket | None = None class TaskState(Enum): """WriteTask lifecycle: a strict chain, plus FAILED from any non-terminal. PENDING -> QUEUED -> COPYING -> HOST_READY -> WRITTEN PENDING registered, not on a copy queue (deferred: waiting for finish / budget pressure) QUEUED on a copy queue COPYING one thread owns the copy stage (staging: waits the step device→host; deferred: copies the device clone) HOST_READY host rows ready (`host_ready` set); device clone dropped WRITTEN in the CPU pool (`done` set) FAILED committer could not finish (`host_ready` + `done` set; manager raises from every later entry) Moves only through `WriteTask.transition`; skipping a step is illegal. """ PENDING = auto() QUEUED = auto() COPYING = auto() HOST_READY = auto() WRITTEN = auto() FAILED = auto() @property def is_terminal(self) -> bool: return self in (TaskState.WRITTEN, TaskState.FAILED) # The one legal forward move from each non-terminal state. _CHAIN: dict[TaskState, TaskState] = { TaskState.PENDING: TaskState.QUEUED, TaskState.QUEUED: TaskState.COPYING, TaskState.COPYING: TaskState.HOST_READY, TaskState.HOST_READY: TaskState.WRITTEN, } _NOT_YET_COPYING = frozenset({TaskState.PENDING, TaskState.QUEUED}) def _is_legal_move(frm: TaskState, to: TaskState) -> bool: return (to is TaskState.FAILED and not frm.is_terminal) or _CHAIN.get(frm) is to @dataclass class WriteTask: """One write of (slot, key) rows for a single request. Identity: `tid` is the handle in the manager's (slot, key) tables. `req_id` + `write_n` mark whose write this is and the nth time that request opened a write. One write may cover several keys. `state` is the single source of truth for where the write is (see `TaskState`). `host_ready` / `done` are wait primitives set by the transitions into HOST_READY / WRITTEN / FAILED, never directly. How HOST_READY is reached: - JOIN_NEXT_STEP: `chunk.host` is a staging view set at save; the copy stage only waits `step_d2h_event`. - JOIN_ON_FINISH: the copy stage copies device→host into `chunk.host`. `JOIN_NEXT_STEP` is queued at submit. `JOIN_ON_FINISH` stays PENDING until finish or GPU-byte-budget pressure escalates it; budget flush takes the unfinished task with the oldest `enqueued_time`. Concurrent readers/writers: - Staging readers (materialize clone, committer pool write) all wait the same `step_d2h_event` before touching the view. - A later task taking the same (slot, key) records those rows in `reassigned`; the old pool write skips them. - `append_chunk` is refused once the copy stage is claimed; the caller opens a fresh task rather than mutating a closed one. - `lock` covers `state` / `reassigned` / `append_chunk` / host↔freeze / `slot_to_row`. `scatter_rows` snapshots `reassigned`. """ tid: Tid req_id: ReqId write_n: int # 1-based: nth write opened by this request schedule: WriteSchedule chunks: list[_WriteChunk] # On-device clone finished on the compute stream. Copy/read streams # must wait this before touching `chunks[].tensors`, or they can read # the next CUDA-graph static-buffer overwrite. freeze_event: torch.cuda.Event | None = None state: TaskState = TaskState.PENDING # Slots this write no longer owns (a newer write took them over). reassigned: dict[TensorName, torch.Tensor] = field(default_factory=dict) # Set on entering HOST_READY (or FAILED): this write's host rows are ready. host_ready: threading.Event = field(default_factory=threading.Event) # Set on entering WRITTEN (or FAILED): the CPU-pool write is over. done: threading.Event = field(default_factory=threading.Event) # Guards state / reassigned / append / host↔freeze. lock: threading.Lock = field(default_factory=threading.Lock) # time.monotonic() at submit; GPU-byte flush picks the oldest of these. enqueued_time: float = field(default_factory=time.monotonic) # Immediate path: this write's view into the shared step staging page. staging_slot: int | None = None # Immediate path: this step's CUDA device→host event (shared). None if deferred. step_d2h_event: torch.cuda.Event | None = None # slot -> (which chunk, row in that tensor). Built on demand # when a write has more than one `_WriteChunk`; pool write uses slot, the tensor uses row. _slot_to_row: dict[int, tuple[int, int]] | None = None def add_reassigned(self, key: str, slots: torch.Tensor) -> None: with self.lock: prev = self.reassigned.get(key) self.reassigned[key] = slots.clone() if prev is None else torch.cat([prev, slots]) def reassigned_intersects(self, chunk: _WriteChunk) -> bool: """True if this write already lost ownership of any (key, slot) the chunk would add; the caller opens a fresh WriteTask instead.""" with self.lock: if not self.reassigned: return False for key in chunk.tensors: taken = self.reassigned.get(key) if taken is not None and taken.numel() and bool(torch.isin(chunk.slots_cpu, taken).any()): return True return False # ------------------------------------------------------------ lifecycle def transition(self, to: TaskState) -> None: """Move to `to`; raises unless it is the next chain step or FAILED.""" with self.lock: self._transition_locked(to) def try_transition(self, to: TaskState) -> bool: """Move to `to` if legal from the current state. True if moved.""" with self.lock: if not _is_legal_move(self.state, to): return False self._transition_locked(to) return True def _transition_locked(self, to: TaskState, wake: bool = True) -> None: if not _is_legal_move(self.state, to): raise OmniPrefixCacheUnmatchError(f"task {self.tid}: illegal transition {self.state.name} -> {to.name}") self.state = to if not wake: return if to is TaskState.HOST_READY: self.host_ready.set() elif to is TaskState.WRITTEN: self.done.set() elif to is TaskState.FAILED: self.host_ready.set() self.done.set() def claim_copy(self) -> bool: """QUEUED -> COPYING. Only one thread may run the copy stage; True if this caller won.""" return self.try_transition(TaskState.COPYING) @property def is_terminal(self) -> bool: return self.state.is_terminal def append_chunk(self, chunk: _WriteChunk, freeze_event: torch.cuda.Event | None = None) -> TaskState | None: """Grow this write with one save's rows. Returns None when appended, else the state that closed the task (COPYING or later). `freeze_event` is stored in the same snapshot: events on one compute stream are ordered, so the newest also covers every earlier clone. """ with self.lock: if self.state not in _NOT_YET_COPYING: return self.state self.chunks.append(chunk) self._slot_to_row = None if freeze_event is not None: self.freeze_event = freeze_event return None def get_host_tensor(self, si: int, key: str) -> torch.Tensor | None: """`chunks[si]` host if written, else device freeze. One snapshot.""" with self.lock: chunk = self.chunks[si] src = chunk.host.get(key) if src is None: src = chunk.tensors.get(key) return src def set_host_tensor(self, rows: list[tuple[_WriteChunk, str, torch.Tensor]]) -> None: """COPYING -> HOST_READY: write these host tensors, drop the device freeze.""" with self.lock: for chunk, key, tensor in rows: chunk.host[key] = tensor self._clear_tensors() self._transition_locked(TaskState.HOST_READY) def mark_host_ready(self) -> None: """COPYING -> HOST_READY for staging views: wait the step device→host, drop freeze.""" if self.step_d2h_event is not None: self.step_d2h_event.synchronize() with self.lock: self._clear_tensors() self._transition_locked(TaskState.HOST_READY) def mark_failed(self, *, defer_wake: bool = False) -> bool: """-> FAILED from any non-terminal state; unblocks joiners. False if already terminal. ``defer_wake``: the caller publishes the failure record first, then calls ``wake_failed_waiters``.""" with self.lock: if self.state.is_terminal: return False self._clear_tensors() self._transition_locked(TaskState.FAILED, wake=not defer_wake) return True def wake_failed_waiters(self) -> None: """Set the events after the failure record is published (see ``mark_failed(defer_wake=True)``): a woken joiner must find it.""" self.host_ready.set() self.done.set() def mark_done(self) -> None: """HOST_READY -> WRITTEN.""" self.transition(TaskState.WRITTEN) def scatter_rows(self) -> list[tuple[TensorName, torch.Tensor, torch.Tensor]]: """`(key, slots, host)` to write, one entry per key. Omits slots in `reassigned`. A slot written twice by this task (preempt + resume onto the same block) keeps the later chunk's row. """ with self.lock: reassigned = {k: s.clone() for k, s in self.reassigned.items()} by_key: dict[TensorName, tuple[list[torch.Tensor], list[torch.Tensor]]] = {} for chunk in self.chunks: for k, host in chunk.host.items(): s, h = by_key.setdefault(k, ([], [])) s.append(chunk.slots_cpu) h.append(host) out: list[tuple[TensorName, torch.Tensor, torch.Tensor]] = [] for k, (s, h) in by_key.items(): slots = s[0] if len(s) == 1 else torch.cat(s) host = h[0] if len(h) == 1 else torch.cat(h, dim=0) if len(s) > 1 and torch.unique(slots).numel() != slots.numel(): last_pos = {int(slot): i for i, slot in enumerate(slots.tolist())} keep_idx = torch.tensor(sorted(last_pos.values()), dtype=torch.int64) slots, host = slots[keep_idx], host[keep_idx] taken = reassigned.get(k) if taken is not None and taken.numel(): keep = ~torch.isin(slots, taken) if not bool(keep.any()): continue slots, host = slots[keep], host[keep] out.append((k, slots, host)) return out def clear_tensors(self) -> None: """Drop the device freeze. Host is unchanged (staging wait / fail).""" with self.lock: self._clear_tensors() def _clear_tensors(self) -> None: for chunk in self.chunks: chunk.tensors = {} def budget_tickets(self) -> list[_BudgetTicket]: """Distinct tickets across this task's chunks (by identity; one per save).""" out: dict[int, _BudgetTicket] = {} for chunk in self.chunks: if chunk.budget is not None: out.setdefault(id(chunk.budget), chunk.budget) return list(out.values()) def slot_to_row(self) -> dict[int, tuple[int, int]]: with self.lock: if self._slot_to_row is None: m: dict[int, tuple[int, int]] = {} for si, chunk in enumerate(self.chunks): for ri, slot in enumerate(chunk.slots_cpu.tolist()): m[slot] = (si, ri) self._slot_to_row = m return self._slot_to_row class StagingBufferPool: """Reusable pinned landing zone for ONE whole-step device→host at save. Per-task `chunk.host` is a row-range view into a slot, so the committer skips a per-task device→host. Slots recycle; this is not the CPU block pool. A slot stays busy while anyone still holds it: the step (until materialize/discard) and each immediate write that views the page (until its pool write). Prefix hits do not hold a slot — they wait for the pool write and read the durable pool. Saves with only leftover mm still claim a slot (empty views) so every step id shares this bound. A full pool waits; timeout then errors. """ def __init__(self, depth: int, capacity: int): self.depth = depth self.capacity = capacity # rows per slot; a larger step raises self._bufs: dict[TensorName, torch.Tensor] = {} # [depth*capacity, width] self._busy: list[set[StagingBufferHolder]] = [set() for _ in range(depth)] self._slot_free_condition = threading.Condition() self._closed = False def _buf(self, key: str, width: int, dtype: torch.dtype, pin: bool) -> torch.Tensor: buf = self._bufs.get(key) if buf is None or buf.shape[-1] != width or buf.dtype != dtype: buf = torch.empty((self.depth * self.capacity, width), dtype=dtype, pin_memory=pin) self._bufs[key] = buf return buf def claim(self, holder: StagingBufferHolder, timeout: float) -> int: """Grab a free slot for `holder`. Waits until one is free, then times out. ``timeout<=0`` fails immediately if none are free. """ deadline = time.monotonic() + max(0.0, timeout) with self._slot_free_condition: while True: if self._closed: raise OmniPrefixCacheUnmatchError("staging pool shut down while waiting for a slot") for slot in range(self.depth): if not self._busy[slot]: self._busy[slot].add(holder) return slot remaining = deadline - time.monotonic() if remaining <= 0: raise OmniPrefixCacheStagingTimeoutError(f"timed out waiting for a staging slot after {timeout}s") self._slot_free_condition.wait(timeout=remaining) def bind(self, slot: int, holder: StagingBufferHolder) -> None: with self._slot_free_condition: self._busy[slot].add(holder) def release(self, slot: int, holder: StagingBufferHolder) -> None: with self._slot_free_condition: self._busy[slot].discard(holder) if not self._busy[slot]: self._slot_free_condition.notify() def close(self) -> None: with self._slot_free_condition: self._closed = True self._slot_free_condition.notify_all() def views(self, slot: int, key: str, n: int, width: int, dtype: torch.dtype, pin: bool) -> torch.Tensor: base = slot * self.capacity return self._buf(key, width, dtype, pin)[base : base + n] @dataclass class StepD2HClaim: """One whole-step landing in StagingBufferPool. Return of ``stage_step_host``. The manager stores this on ``_StepContext`` until materialize/discard releases the step holder. """ staging_slot: int # StagingBufferPool index views: dict[TensorName, torch.Tensor] # host rows [0:n) event: torch.cuda.Event | None = None # torch.cuda.Event; None on eager/CPU class OmniPrefixCacheController: """Staging pool + committer. Step device→host is launched at save; this thread waits that event (JOIN_NEXT_STEP) or copies deferred rows (JOIN_ON_FINISH), then writes into the CPU pool. """ def __init__(self, pool: PrefixBlockPool, config: PrefixCacheConfig, eager: bool | None = None): self._pool = pool self._config = config self._eager = (not torch.cuda.is_available()) if eager is None else eager self._tasks: dict[Tid, WriteTask] = {} self._completed: deque[Tid] = deque() # pool write done, awaiting manager drain self._failed: deque[Tid] = deque() # write failed; manager raises from every later entry (sticky) self._staged_bytes = 0 self._lock = threading.Lock() self._wake = threading.Condition(self._lock) self._queue_hi: deque[Tid] = deque() # JOIN_NEXT_STEP + forced JOIN_ON_FINISH self._queue_lo: deque[Tid] = deque() # JOIN_ON_FINISH waiting for finish/budget self._blocked: list[Tid] = [] # host copy done, awaiting pool write self._shutdown = False self._copy_stream: torch.cuda.Stream | None = None self._read_stream: torch.cuda.Stream | None = None self._worker: threading.Thread | None = None self._staging_pool = StagingBufferPool(config.staging_depth, config.staging_capacity_tokens) if not self._eager: self._copy_stream = torch.cuda.Stream() self._read_stream = torch.cuda.Stream() self._worker = threading.Thread(target=self._worker_loop, name="omni-prefix-cache-committer", daemon=True) self._worker.start() # --------------------------------------------------------- step device→host staging def _d2h_on_stream( self, stream: torch.cuda.Stream, freeze_event: torch.cuda.Event | None, copy: Callable[[], None], ) -> torch.cuda.Event: """Issue device→host on `stream` after freeze; return the done event.""" with torch.cuda.stream(stream): if freeze_event is not None: stream.wait_event(freeze_event) copy() ev = torch.cuda.Event() ev.record() return ev def stage_step_host( self, tensors: dict[str, torch.Tensor], n: int, freeze_event: torch.cuda.Event | None, step_holder: StagingBufferHolder, ) -> StepD2HClaim: """Claim a staging slot and, when `tensors` is non-empty, launch ONE whole-step device→host into it. Leftover-only saves pass empty `tensors` and still take a slot (empty views) so every step id shares this bound. A full pool waits for materialize/discard; timeout then errors. A step larger than the page overflows the next slot — that is a config break. The caller binds tasks after submit; `step_holder` is released by materialize/discard via staging_release. """ if tensors and n > self._staging_pool.capacity: raise OmniPrefixCacheUnmatchError( f"step has {n} tokens; staging capacity is {self._staging_pool.capacity} " "(size staging_capacity_tokens to max_num_batched_tokens)" ) slot = self._staging_pool.claim(step_holder, self._config.staging_claim_timeout_s) try: pin = not self._eager views: dict[str, torch.Tensor] = {} event: torch.cuda.Event | None = None if self._eager or all(t.device.type == "cpu" for t in tensors.values()): for key, src in tensors.items(): v = self._staging_pool.views(slot, key, n, int(src.shape[-1]), src.dtype, pin) v.copy_(src) views[key] = v else: def _copy_to_staging() -> None: for key, src in tensors.items(): v = self._staging_pool.views(slot, key, n, int(src.shape[-1]), src.dtype, pin) v.copy_(src, non_blocking=True) views[key] = v event = self._d2h_on_stream(self._copy_stream, freeze_event, _copy_to_staging) return StepD2HClaim(staging_slot=slot, views=views, event=event) except Exception: self._staging_pool.release(slot, step_holder) raise def staging_bind(self, slot: int, holder: StagingBufferHolder) -> None: self._staging_pool.bind(slot, holder) def staging_release(self, slot: int, holder: StagingBufferHolder) -> None: self._staging_pool.release(slot, holder) def _release_task_slot(self, task: WriteTask) -> None: """Drop the task's staging hold at a terminal state (pool write or failure).""" if task.staging_slot is not None: self._staging_pool.release(task.staging_slot, StagingBufferHolder.for_task(task.tid)) def in_flight_tasks(self) -> int: """Registered tasks not yet drained (diagnostics only).""" return len(self._tasks) # ------------------------------------------------------------------ submit def register(self, task: WriteTask, queued: bool = True) -> None: """Make a task visible (registry + QUEUED) without running anything; safe under the manager's state lock. queued=False (deferred tasks) stays PENDING on the GPU clone until finish/abort or the GPU-byte budget forces a copy. Queued tasks must then go through ``dispatch``. Caller must reserve() the task bytes first (budget flush can block; the manager does that outside the state lock) and pin the task on its budget ticket(s) before register. """ task.enqueued_time = time.monotonic() with self._lock: self._tasks[task.tid] = task if queued: task.transition(TaskState.QUEUED) def dispatch(self, tasks: list[WriteTask]) -> None: """Hand registered, queued tasks to the copy path. Threaded: enqueue. Eager: the copy + pool write run here, inline — never call this under the manager's state lock.""" if not tasks: return if self._eager: try: for task in tasks: self._run_eager(task) except Exception: # Later tasks were registered but will never be dispatched. # Fail them too so shared budget owners and waiters release. for task in tasks: self._fail_task(task.tid) raise return with self._wake: for task in tasks: (self._queue_hi if task.schedule is WriteSchedule.JOIN_NEXT_STEP else self._queue_lo).append(task.tid) self._wake.notify_all() def submit(self, task: WriteTask, queued: bool = True) -> None: """register + dispatch in one call (callers not holding the state lock).""" self.register(task, queued) if queued: self.dispatch([task]) def append_chunk( self, task: WriteTask, chunk: _WriteChunk, freeze_event: torch.cuda.Event | None = None ) -> TaskState | None: """Append to a pending task. None when appended, else the closing state.""" return task.append_chunk(chunk, freeze_event) def pin_budget(self, ticket: _BudgetTicket, tid: int) -> None: """Record that ``tid`` holds a view of the clone ``ticket`` charges.""" with self._lock: ticket.tids.add(tid) def reserve(self, nbytes: int) -> None: """Reserve GPU-clone bytes; blocking flush happens here, so callers must not hold the manager's state lock.""" self._reserve_bytes(nbytes) def _release_staged_bytes(self, task: WriteTask) -> None: """Drop this task's hold on its budget ticket(s). Idempotent; a ticket uncharges once, when its last holder leaves.""" tickets = task.budget_tickets() with self._wake: for ticket in tickets: self._staged_bytes -= ticket.release(task.tid) def _reserve_bytes(self, nbytes: int) -> None: # GPU-byte budget: force-copy oldest pending tasks until under # budget. Bounded wait: their device→host has usually long completed. while True: with self._lock: pending = [tid for tid, t in self._tasks.items() if not t.is_terminal] if self._staged_bytes + nbytes <= self._config.gpu_staging_bytes or not pending: # Under budget or no pending tasks; admit reservation. self._staged_bytes += nbytes return oldest = min(pending, key=lambda tid: self._tasks[tid].enqueued_time) logger.warning("omni prefix cache: staging cap hit, force-flushing task %d", oldest) self.escalate([oldest]) self.join([oldest]) # ------------------------------------------------------------- lifecycle def escalate(self, tids: list[int]) -> None: """Move pending tasks to the front of the high-priority queue. PENDING -> QUEUED at the head; QUEUED on the low-priority queue moves up; QUEUED already high-priority is a no-op. COPYING and later belong to the worker and are untouched: the worker claims under `_wake` too, so a task cannot be popped and re-queued behind its back. """ if self._eager: self.dispatch([task for tid in tids if (task := self._tasks.get(tid)) is not None and not task.is_terminal]) return with self._wake: for tid in tids: task = self._tasks.get(tid) if task is None: continue if task.try_transition(TaskState.QUEUED): self._queue_hi.appendleft(tid) elif task.state is TaskState.QUEUED and tid in self._queue_lo: self._queue_lo.remove(tid) self._queue_hi.appendleft(tid) self._wake.notify_all() def join(self, tids: list[int]) -> None: """Block until each task has finished the CPU-pool write (or failed).""" self._wait_tasks(tids, "done") def join_host_ready(self, tids: list[int]) -> None: """Block until each task's device→host is complete (`host_ready`). Staging: committer has waited `step_d2h_event`. Deferred: committer has written `chunk.host`. Does not wait for the CPU-pool write. """ self._wait_tasks(tids, "host_ready") def _wait_tasks(self, tids: list[int], event: Literal["done", "host_ready"]) -> None: timeout = self._config.staging_claim_timeout_s deadline = time.monotonic() + timeout for tid in tids: task = self._tasks.get(tid) if task is None: continue if not getattr(task, event).wait(max(0.0, deadline - time.monotonic())): raise OmniPrefixCacheUnmatchError( f"task {tid} ({task.schedule.value}) did not reach {event} within {timeout:g}s " f"(state={task.state.name}, in_flight_tasks={len(self._tasks)})" ) if task.state is TaskState.FAILED: # FAILED sets both events; success here would read rows the # write never produced. The one-shot failure record may have # been drained by another thread — do not rely on it. raise OmniPrefixCacheUnmatchError(f"task {tid} ({task.schedule.value}) write failed before {event}") def drain_completed(self) -> list[int]: """Pop pool-written tasks from `_completed` and drop them from `_tasks`. Staging holders were already released at the pool write.""" out: list[int] = [] with self._lock: while self._completed: out.append(self._completed.popleft()) for tid in out: self._tasks.pop(tid, None) return out def drain_failed(self) -> list[int]: """Pop failed task ids from `_failed`. Does not drop `_tasks`.""" out: list[int] = [] with self._lock: while self._failed: out.append(self._failed.popleft()) return out def get_task(self, tid: int) -> WriteTask | None: return self._tasks.get(tid) def shutdown(self) -> None: self._staging_pool.close() with self._wake: self._shutdown = True self._wake.notify_all() if self._worker is not None: self._worker.join(timeout=5.0) # ------------------------------------------------------------ fetch_host @torch.inference_mode() def fetch_host(self, task: WriteTask, slots: torch.Tensor, key: str) -> torch.Tensor: """Rows for `slots` of one not-yet-done JOIN_ON_FINISH task. `_slot_ref` puts JOIN_NEXT_STEP tids in `join_tids` (wait then pool). This path reads committer-written `chunk.host`, or the device clone if that device→host has not landed. """ if task.step_d2h_event is not None: task.step_d2h_event.synchronize() return self._rows_from(task, slots, key) def _rows_from(self, task: WriteTask, slots: torch.Tensor, key: str) -> torch.Tensor: """Map `slots` to rows across one or more `_WriteChunk`s; preserve caller order.""" s2r = task.slot_to_row() idx = [s2r[int(s)] for s in slots.tolist()] parts: list[torch.Tensor] = [] order: list[int] = [] pos = 0 rows_groups: dict[int, list[tuple[int, int]]] = {} for si, ri in idx: rows_groups.setdefault(si, []).append((pos, ri)) pos += 1 for si, items in rows_groups.items(): src = task.get_host_tensor(si, key) if src is None: continue rows_idx = torch.tensor([ri for _, ri in items], dtype=torch.long) picked = self._slice_rows(task, src, rows_idx, host=(src.device.type == "cpu")) parts.append(picked) order.extend(p for p, _ in items) if not parts: raise KeyError(f"key {key} not present in task {task.tid}") cat = torch.cat(parts, dim=0) out = torch.empty_like(cat) out[torch.tensor(order, dtype=torch.long)] = cat return out def _slice_rows(self, task: WriteTask, src: torch.Tensor, rows_idx: torch.Tensor, host: bool) -> torch.Tensor: n = int(rows_idx.numel()) # Ascending-run check without materializing an arange: endpoints plus # a monotonic diff are enough, and the common case is one long run. contiguous = ( n > 0 and int(rows_idx[-1]) - int(rows_idx[0]) == n - 1 and (n < 2 or bool((rows_idx.diff() == 1).all())) ) def _pick() -> torch.Tensor: return src[rows_idx[0] : rows_idx[0] + n] if contiguous else src.index_select(0, rows_idx) if host or src.device.type == "cpu": return _pick() if self._read_stream is None: return _pick().detach().cpu() copied: list[torch.Tensor] = [] def _copy_to_cpu() -> None: copied.append(_pick().to("cpu", non_blocking=True)) ev = self._d2h_on_stream(self._read_stream, task.freeze_event, _copy_to_cpu) ev.synchronize() return copied[0] # ------------------------------------------------------------ eager mode @torch.inference_mode() def _run_eager(self, task: WriteTask) -> None: """Eager stand-in for queue + worker: walk the whole chain inline.""" task.try_transition(TaskState.QUEUED) # no-op if submit already queued it if not task.claim_copy(): # Copy stage already ran (a held stub, or a prior escalate); finish the write. if task.state is TaskState.HOST_READY: self._scatter(task) return if task.schedule is WriteSchedule.JOIN_NEXT_STEP: # Host is already a staging view (copied at save). Drop freeze. # `step_d2h_event` is None on CPU; `mark_host_ready` skips wait. task.mark_host_ready() else: # JOIN_ON_FINISH: no copy stream; freeze → owned host inline. task.set_host_tensor( [ (chunk, k, t.detach().cpu() if t.device.type != "cpu" else t.clone()) for chunk in task.chunks for k, t in chunk.tensors.items() ] ) self._release_staged_bytes(task) self._scatter(task) # ---------------------------------------------------------- worker loop def _worker_loop(self) -> None: # A dying committer would strand every join() forever, so the loop # never propagates: it fails the offending task and keeps serving. while True: tid = None try: task: WriteTask | None = None with self._wake: while not self._shutdown and not self._queue_hi and not self._queue_lo: if self._blocked: break # submit / escalate / shutdown all notify. self._wake.wait() if self._shutdown and not self._queue_hi and not self._queue_lo and not self._blocked: return if self._queue_hi: tid = self._queue_hi.popleft() elif self._queue_lo: tid = self._queue_lo.popleft() if tid is not None: task = self._tasks.get(tid) if task is not None: # Claim in the same critical section as the pop: # escalate holds `_wake` and sees COPYING. task.transition(TaskState.COPYING) if task is not None: self._copy_task(task) with self._wake: assert task.tid not in self._blocked, f"task {task.tid} reached HOST_READY twice" self._blocked.append(task.tid) self._scatter_host_ready() except BaseException: logger.exception("omni prefix cache committer failed on task %s; releasing waiters", tid) self._fail_task(tid) @torch.inference_mode() def _copy_task(self, task: WriteTask) -> None: """COPYING -> HOST_READY. Staging: wait the save-time device→host event. Deferred: this is the device→host into owned `chunk.host` tensors. """ if task.schedule is WriteSchedule.JOIN_NEXT_STEP: # `chunk.host` is already a staging view; device→host ran at save. # `mark_host_ready` waits `step_d2h_event` if one was recorded. # No per-task copy. Pool write is `_scatter_host_ready`. task.mark_host_ready() self._release_staged_bytes(task) return # One device cat + one device→host per key, not one per chunk (a # long request has one chunk per step). chunk_bytes = self._config.copy_chunk_bytes by_key: dict[str, list[_WriteChunk]] = {} for chunk in task.chunks: for k in chunk.tensors: by_key.setdefault(k, []).append(chunk) pending: list[tuple[str, list[_WriteChunk], list[torch.Tensor]]] = [] if self._copy_stream is None: raise OmniPrefixCacheUnmatchError("deferred prefix cache copy requires a CUDA copy stream") with torch.cuda.stream(self._copy_stream): if task.freeze_event is not None: self._copy_stream.wait_event(task.freeze_event) for k, chunks in by_key.items(): srcs = [c.tensors[k] for c in chunks] src = srcs[0] if len(srcs) == 1 else torch.cat(srcs, dim=0) rows = int(src.shape[0]) row_bytes = max(1, src[:1].numel() * src.element_size()) step = max(1, rows if rows * row_bytes <= chunk_bytes else chunk_bytes // row_bytes) parts = [src[s : s + step].to("cpu", non_blocking=True) for s in range(0, max(rows, 1), step)] pending.append((k, chunks, parts)) ev = torch.cuda.Event() ev.record() ev.synchronize() rows_out: list[tuple[_WriteChunk, str, torch.Tensor]] = [] for k, chunks, parts in pending: host = parts[0] if len(parts) == 1 else torch.cat(parts, dim=0) if len(chunks) == 1: rows_out.append((chunks[0], k, host)) continue for chunk, view in zip(chunks, torch.split(host, [int(c.tensors[k].shape[0]) for c in chunks], dim=0)): rows_out.append((chunk, k, view)) task.set_host_tensor(rows_out) self._release_staged_bytes(task) def _fail_task(self, tid: int | None) -> None: """-> FAILED: release waiters for a task the committer could not complete. No-op once the task is terminal. Budget release is idempotent, so a raise after a successful copy stage does not uncharge twice. """ task = self._tasks.get(tid) if tid is not None else None if task is None or not task.mark_failed(defer_wake=True): return self._release_staged_bytes(task) with self._wake: if tid in self._blocked: self._blocked.remove(tid) self._release_task_slot(task) with self._lock: # Publish the failure: rows behind already-published block hashes # never landed, which the manager must raise on (hiding it would # crash on every future hit that touches these slots). self._failed.append(task.tid) # Events last: a joiner woken by them must find the record on drain, # not read never-written pool rows in a done-but-unpublished window. task.wake_failed_waiters() @torch.inference_mode() def _scatter_host_ready(self) -> None: """HOST_READY -> WRITTEN for every `_blocked` task that has reached HOST_READY.""" with self._wake: ready = [tid for tid in self._blocked if (t := self._tasks.get(tid)) and t.state is TaskState.HOST_READY] for tid in ready: self._blocked.remove(tid) for tid in ready: task = self._tasks[tid] try: self._scatter(task) except BaseException: # Attribute the failure to THIS task: letting it propagate # would fail whichever entry the worker loop happened to be # copying, double-release its bytes, and strand this one's # join() forever. logger.exception("omni prefix cache scatter failed on task %s; releasing waiters", tid) self._fail_task(tid) @torch.inference_mode() def _scatter(self, task: WriteTask) -> None: for key, slots, host in task.scatter_rows(): self._pool.write(key, slots, host) # `done` is set LAST: a join(done) waiter (the save barrier) must find # the completion record on its next drain and the staging slot free — # a done-but-undrained window would leave occupancy IN_TRANSIT and # make the COW preserve skip the reused rows. self._release_task_slot(task) with self._lock: self._completed.append(task.tid) task.mark_done()