[Bugfix][Helios] Guard cross-attn KV cache against address reuse (#8064)

Signed-off-by: yancaocn <yancaochn@163.com>
Co-authored-by: yancaocn <yancaochn@163.com>
Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>

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"

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)

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(