77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Concrete memory objects (RFC #4480).
|
|
|
|
``LatentBuffer`` backs its storage with a plain (monolithic) buffer. Attention
|
|
KV is deliberately absent: for DreamZero it is owned by the AR-Diffusion engine's
|
|
paged pool (PR #4534), and the RFC's ``PagedKV`` arrives in Phase 1 as a
|
|
handle wrapping that engine state rather than a dense buffer here. The RFC's
|
|
``FixedState`` and ``RetrievalStore`` land with their first consumers, but the
|
|
names are reserved.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
from collections.abc import Iterable
|
|
from typing import TypeVar
|
|
|
|
from vllm_omni.experimental.world_models.session_state.accounting import (
|
|
SeenStorages,
|
|
device_bytes,
|
|
)
|
|
from vllm_omni.experimental.world_models.session_state.base import StateObject
|
|
|
|
# The frame/chunk type a ``LatentBuffer`` holds (e.g. ``np.ndarray`` pixel
|
|
# frames, ``torch.Tensor`` latent chunks).
|
|
ItemT = TypeVar("ItemT")
|
|
|
|
|
|
class LatentBuffer(StateObject[ItemT, list[ItemT]]):
|
|
"""Append / ring buffer of latent or pixel frames.
|
|
|
|
A bounded ``deque`` (``maxlen`` set at ``allocate()`` time). The object
|
|
only stores and views the frames; model-specific stacking stays in the
|
|
caller.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._buf: deque[ItemT] | None = None
|
|
|
|
def allocate(self, *, maxlen: int | None = None) -> None:
|
|
self._buf = deque(maxlen=maxlen)
|
|
|
|
def append(self, payload: ItemT) -> None:
|
|
if self._buf is None:
|
|
raise RuntimeError("LatentBuffer is not allocated; call allocate() first.")
|
|
self._buf.append(payload)
|
|
|
|
def extend(self, payloads: Iterable[ItemT]) -> None:
|
|
if self._buf is None:
|
|
raise RuntimeError("LatentBuffer is not allocated; call allocate() first.")
|
|
self._buf.extend(payloads)
|
|
|
|
def commit(self, payload: ItemT | None = None) -> None:
|
|
if payload is not None:
|
|
self.append(payload)
|
|
|
|
def view(self, *, include_staged: bool = True) -> list[ItemT]:
|
|
if self._buf is None:
|
|
raise RuntimeError("LatentBuffer is not allocated; call allocate() first.")
|
|
return list(self._buf)
|
|
|
|
def __len__(self) -> int:
|
|
return 0 if self._buf is None else len(self._buf)
|
|
|
|
def reset(self) -> None:
|
|
self._buf = None
|
|
|
|
def nbytes_by_device(self, seen: SeenStorages) -> dict[str, int]:
|
|
if self._buf is None:
|
|
return {}
|
|
return device_bytes(list(self._buf), seen)
|
|
|
|
@property
|
|
def resident(self) -> bool:
|
|
return self._buf is not None
|