004a5711dc
Signed-off-by: yancaocn <yancaochn@163.com> Co-authored-by: yancaocn <yancaochn@163.com> Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>
101 lines
3.9 KiB
Python
101 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Weakref-guarded LRU cache for cross-attention key/value reuse.
|
|
|
|
``HeliosTransformer3DModel`` projects the text-encoder output into
|
|
cross-attention K/V once per request and reuses the projection across the
|
|
~50 denoise steps of that request. The cache is keyed by a storage
|
|
fingerprint of the source tensor (``data_ptr`` / shape / stride / dtype /
|
|
device / ``_version``). That fingerprint is **not** unique across requests:
|
|
once the source tensor is freed, PyTorch's caching allocator may hand its
|
|
address to an unrelated tensor of the same shape, producing a false cache
|
|
hit and cross-prompt contamination (request B silently reuses request A's
|
|
projected K/V).
|
|
|
|
``SourceTensorLRUCache`` stores a :mod:`weakref` to the source tensor next
|
|
to the cached value and verifies the source is still alive on every lookup.
|
|
An address reused by a *different* tensor still matches the fingerprint,
|
|
but the original source's weakref is dead by then, so the lookup becomes a
|
|
miss instead of a silent wrong hit. Within a single request the source
|
|
tensor is held alive across all denoise steps, so legitimate intra-request
|
|
reuse still hits.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import weakref
|
|
from collections import OrderedDict
|
|
|
|
import torch
|
|
|
|
|
|
class SourceTensorLRUCache:
|
|
"""Bounded LRU cache keyed by a source-tensor fingerprint + liveness.
|
|
|
|
The key is derived from the source tensor's storage fingerprint
|
|
(``data_ptr`` / shape / stride / dtype / device / ``_version``); a
|
|
weakref to the source is stored alongside the value and checked on
|
|
lookup so that an allocator address reused after the source is freed
|
|
never yields a false hit. Stale entries (dead weakref) are evicted
|
|
lazily on lookup and proactively on insert.
|
|
"""
|
|
|
|
def __init__(self, max_size: int = 2) -> None:
|
|
self._cache: OrderedDict[tuple, tuple[weakref.ReferenceType, object]] = OrderedDict()
|
|
self._max_size = max_size
|
|
|
|
@staticmethod
|
|
def _key(tensor: torch.Tensor) -> tuple:
|
|
try:
|
|
version = tensor._version
|
|
except RuntimeError:
|
|
version = None
|
|
return (
|
|
tensor.data_ptr(),
|
|
tuple(tensor.shape),
|
|
tuple(tensor.stride()),
|
|
tensor.dtype,
|
|
tensor.device.type,
|
|
tensor.device.index,
|
|
version,
|
|
)
|
|
|
|
def get(self, source: torch.Tensor):
|
|
"""Return the cached value for ``source`` or ``None`` on miss.
|
|
|
|
A miss is also returned when the source tensor that originally
|
|
populated a matching entry has been freed (its weakref is dead),
|
|
even if the fingerprint now matches a new tensor reusing the
|
|
address; the stale entry is evicted in that case.
|
|
"""
|
|
key = self._key(source)
|
|
entry = self._cache.get(key)
|
|
if entry is None:
|
|
return None
|
|
src_ref, value = entry
|
|
if src_ref() is None:
|
|
# Source freed: the address may have been reused by an unrelated
|
|
# tensor, so the fingerprint match is a coincidence -> miss.
|
|
self._cache.pop(key, None)
|
|
return None
|
|
self._cache.move_to_end(key)
|
|
return value
|
|
|
|
def put(self, source: torch.Tensor, value) -> None:
|
|
key = self._key(source)
|
|
# Proactively drop entries whose source has been freed so cached
|
|
# values do not linger until the LRU bound evicts them.
|
|
if self._cache:
|
|
for stale in [k for k, v in self._cache.items() if v[0]() is None]:
|
|
del self._cache[stale]
|
|
self._cache[key] = (weakref.ref(source), value)
|
|
self._cache.move_to_end(key)
|
|
while len(self._cache) > self._max_size:
|
|
self._cache.popitem(last=False)
|
|
|
|
def clear(self) -> None:
|
|
self._cache.clear()
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._cache)
|