From 004a5711dc1495e38fd12c4deae0438f680f68fa Mon Sep 17 00:00:00 2001 From: Yan Cao <31481315+yancaocn@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:44:51 +0800 Subject: [PATCH] [Bugfix][Helios] Guard cross-attn KV cache against address reuse (#8064) Signed-off-by: yancaocn Co-authored-by: yancaocn Co-authored-by: Zhou Taichang --- .../models/helios/test_cross_attn_cache.py | 197 ++++++++++++++++++ .../models/helios/cross_attn_cache.py | 100 +++++++++ .../models/helios/helios_transformer.py | 51 ++--- 3 files changed, 309 insertions(+), 39 deletions(-) create mode 100644 tests/diffusion/models/helios/test_cross_attn_cache.py create mode 100644 vllm_omni/diffusion/models/helios/cross_attn_cache.py diff --git a/tests/diffusion/models/helios/test_cross_attn_cache.py b/tests/diffusion/models/helios/test_cross_attn_cache.py new file mode 100644 index 000000000..bd5c12b14 --- /dev/null +++ b/tests/diffusion/models/helios/test_cross_attn_cache.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Regression tests for the cross-attention KV cache liveness guard. + +These tests pin the fix for cross-prompt contamination. 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 -- request B would silently reuse request A's +projected cross-attention 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. + +The address-reuse scenario is reproduced deterministically with two views +over the same storage: the second view's fingerprint is identical to the +first's (same ``data_ptr`` / shape / dtype), but the first view's weakref +is dead once the view object is freed, so the second lookup must miss. +""" + +import gc + +import pytest +import torch + +from vllm_omni.diffusion.models.helios.cross_attn_cache import SourceTensorLRUCache + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +# --------------------------------------------------------------------------- +# Intra-request reuse (the cache's intended purpose) +# --------------------------------------------------------------------------- + + +class TestIntraRequestReuse: + """Within one request the source tensor is alive across denoise steps.""" + + def test_same_source_hits(self): + cache = SourceTensorLRUCache(max_size=2) + a = torch.randn(4, 8) + cache.put(a, "A") + assert cache.get(a) == "A" + + def test_repeated_get_keeps_hitting(self): + cache = SourceTensorLRUCache(max_size=2) + a = torch.randn(4, 8) + cache.put(a, "A") + for _ in range(50): + assert cache.get(a) == "A" + + def test_distinct_live_sources_do_not_collide(self): + cache = SourceTensorLRUCache(max_size=2) + a = torch.randn(4, 8) + b = torch.randn(4, 8) + assert a.data_ptr() != b.data_ptr() + cache.put(a, "A") + # b has a different data_ptr -> different key -> miss + assert cache.get(b) is None + assert cache.get(a) == "A" + + +# --------------------------------------------------------------------------- +# The bug: an address reused after the source is freed must NOT false-hit +# --------------------------------------------------------------------------- + + +class TestAddressReuseAfterFree: + """The contamination bug: freed source's address reused by a new tensor.""" + + def test_freed_source_reused_address_does_not_false_hit(self): + cache = SourceTensorLRUCache(max_size=2) + storage = torch.empty(4, 8) + a = storage[:] # view; a.data_ptr() == storage.data_ptr() + cache.put(a, "A") + assert cache.get(a) == "A" + ptr_a = a.data_ptr() + del a + gc.collect() + # A new view over the SAME storage -> identical fingerprint + # (data_ptr / shape / stride / dtype / device / version=0). + b = storage[:] + assert b.data_ptr() == ptr_a + # Without the weakref guard this would falsely return "A" + # (cross-prompt contamination). With the guard the dead weakref + # turns the fingerprint match into a miss. + assert cache.get(b) is None + assert len(cache) == 0 # stale entry evicted on miss + + def test_recompute_after_miss_stores_new_source(self): + cache = SourceTensorLRUCache(max_size=2) + storage = torch.empty(4, 8) + a = storage[:] + cache.put(a, "A") + del a + gc.collect() + b = storage[:] # reuses address + assert cache.get(b) is None + cache.put(b, "B") # the recompute path in the transformer + assert cache.get(b) == "B" + assert len(cache) == 1 + + def test_two_consecutive_requests_no_contamination(self): + # End-to-end simulation of the serving sequence that triggered the + # bug: request A populates the cache, A is freed, request B reuses + # A's address. B must never observe A's cached value. + cache = SourceTensorLRUCache(max_size=2) + storage = torch.empty(2, 16) + + a = storage[:] + cache.put(a, "A_projected_kv") + del a + gc.collect() + + b = storage[:] + assert cache.get(b) is None # B does not inherit A's projection + cache.put(b, "B_projected_kv") + assert cache.get(b) == "B_projected_kv" + + def test_freed_then_distinct_address_misses(self): + # When the new tensor gets a DIFFERENT address, the key differs and + # it trivially misses. The weakref guard must not break this path. + cache = SourceTensorLRUCache(max_size=2) + a = torch.randn(4, 8) + cache.put(a, "A") + del a + gc.collect() + b = torch.randn(4, 8) + assert cache.get(b) is None + + +# --------------------------------------------------------------------------- +# LRU semantics (unchanged by the fix, but must still hold) +# --------------------------------------------------------------------------- + + +class TestLRUEviction: + def test_lru_evicts_oldest(self): + cache = SourceTensorLRUCache(max_size=2) + a, b, c = (torch.randn(1) for _ in range(3)) + cache.put(a, "A") + cache.put(b, "B") + cache.put(c, "C") + assert cache.get(a) is None # evicted + assert cache.get(b) == "B" + assert cache.get(c) == "C" + + def test_get_renews_recency(self): + cache = SourceTensorLRUCache(max_size=2) + a, b, c = (torch.randn(1) for _ in range(3)) + cache.put(a, "A") + cache.put(b, "B") + assert cache.get(a) == "A" # a is now most-recent + cache.put(c, "C") # evicts b (oldest) + assert cache.get(b) is None + assert cache.get(a) == "A" + assert cache.get(c) == "C" + + +class TestPutOverwrite: + def test_same_source_overwrites_value(self): + cache = SourceTensorLRUCache(max_size=2) + a = torch.randn(4, 8) + cache.put(a, "A1") + cache.put(a, "A2") + assert cache.get(a) == "A2" + assert len(cache) == 1 + + +class TestClear: + def test_clear_empties_cache(self): + cache = SourceTensorLRUCache(max_size=2) + a, b = torch.randn(1), torch.randn(1) + cache.put(a, "A") + cache.put(b, "B") + assert len(cache) == 2 + cache.clear() + assert len(cache) == 0 + assert cache.get(a) is None + assert cache.get(b) is None + + +class TestProactiveStaleEviction: + def test_put_evicts_dead_entries(self): + cache = SourceTensorLRUCache(max_size=4) + a = torch.randn(4, 8) + cache.put(a, "A") + del a + gc.collect() + b = torch.randn(4, 8) + cache.put(b, "B") # proactively drops the dead "a" entry + assert len(cache) == 1 + assert cache.get(b) == "B" diff --git a/vllm_omni/diffusion/models/helios/cross_attn_cache.py b/vllm_omni/diffusion/models/helios/cross_attn_cache.py new file mode 100644 index 000000000..82ef357cb --- /dev/null +++ b/vllm_omni/diffusion/models/helios/cross_attn_cache.py @@ -0,0 +1,100 @@ +# 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) diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py index 6c6ec223d..6f716256c 100644 --- a/vllm_omni/diffusion/models/helios/helios_transformer.py +++ b/vllm_omni/diffusion/models/helios/helios_transformer.py @@ -3,7 +3,6 @@ # Adapted from Helios (https://github.com/BestWishYsh/Helios) import math -from collections import OrderedDict from collections.abc import Iterable from functools import lru_cache from typing import TYPE_CHECKING, Any @@ -28,6 +27,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelOutput +from vllm_omni.diffusion.models.helios.cross_attn_cache import SourceTensorLRUCache if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -790,44 +790,19 @@ class HeliosTransformer3DModel(nn.Module): # 5. Output norm & projection self.norm_out = HeliosOutputNorm(inner_dim, eps) self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) - self._projected_encoder_cache: OrderedDict[tuple, torch.Tensor] = OrderedDict() - self._cross_attn_kv_cache: OrderedDict[tuple, list[tuple[torch.Tensor, torch.Tensor]]] = OrderedDict() self._cross_attn_cache_size = 2 + # Cross-attention projection reuse is intra-request: the source + # encoder hidden states are held alive across the denoise steps of a + # single request. SourceTensorLRUCache guards each entry with a + # weakref to the source so an allocator address reused after the + # source is freed cannot produce a cross-prompt false hit. + self._projected_encoder_cache = SourceTensorLRUCache(self._cross_attn_cache_size) + self._cross_attn_kv_cache = SourceTensorLRUCache(self._cross_attn_cache_size) @property def dtype(self) -> torch.dtype: return next(self.parameters()).dtype - @staticmethod - def _tensor_cache_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, - ) - - @staticmethod - def _get_from_lru(cache: OrderedDict, key: tuple): - value = cache.get(key) - if value is not None: - cache.move_to_end(key) - return value - - def _put_lru(self, cache: OrderedDict, key: tuple, value) -> None: - cache[key] = value - cache.move_to_end(key) - while len(cache) > self._cross_attn_cache_size: - cache.popitem(last=False) - def clear_cross_attention_cache(self) -> None: self._projected_encoder_cache.clear() self._cross_attn_kv_cache.clear() @@ -839,13 +814,12 @@ class HeliosTransformer3DModel(nn.Module): if not self._cache_enabled(): return self.condition_embedder.text_embedder(encoder_hidden_states) - cache_key = self._tensor_cache_key(encoder_hidden_states) - cached = self._get_from_lru(self._projected_encoder_cache, cache_key) + cached = self._projected_encoder_cache.get(encoder_hidden_states) if cached is not None: return cached projected = self.condition_embedder.text_embedder(encoder_hidden_states) - self._put_lru(self._projected_encoder_cache, cache_key, projected) + self._projected_encoder_cache.put(encoder_hidden_states, projected) return projected def _get_cross_attn_key_values( @@ -855,13 +829,12 @@ class HeliosTransformer3DModel(nn.Module): if not self._cache_enabled(): return None - cache_key = self._tensor_cache_key(encoder_hidden_states) - cached = self._get_from_lru(self._cross_attn_kv_cache, cache_key) + cached = self._cross_attn_kv_cache.get(encoder_hidden_states) if cached is not None: return cached key_values = [block.attn2.project_kv(encoder_hidden_states) for block in self.blocks] - self._put_lru(self._cross_attn_kv_cache, cache_key, key_values) + self._cross_attn_kv_cache.put(encoder_hidden_states, key_values) return key_values def forward(