[Bugfix][Helios] Guard cross-attn KV cache against address reuse by yancaocn · Pull Request #8064 · vllm-project/vllm-omni

yancaocn · 2026-09-23T07:26:16Z

The cross-attention KV cache in HeliosTransformer3DModel reuses the projected text-encoder K/V across a request's denoise steps. It was keyed solely by a storage fingerprint of the source tensor (data_ptr / shape / stride / dtype / device / _version), which is not unique across requests: once the source is freed, PyTorch's caching allocator may reuse its address for an unrelated tensor of the same shape, producing a false cache hit -- request B silently reused request A's cross-attention K/V (cross-prompt contamination), intermittent and hard to reproduce.

Extract the cache into SourceTensorLRUCache and store a weakref to the source tensor alongside each cached value. On lookup, a dead weakref (source freed) turns an address-reuse collision into a miss instead of a false hit. Intra-request reuse is preserved (the source is held alive across all denoise steps); stale entries are evicted lazily on lookup and proactively on insert.

clear_cross_attention_cache() is unchanged (still clears both caches); it remains available for explicit per-request eviction.

Tests: tests/diffusion/models/helios/test_cross_attn_cache.py
test_cross_attn_cache.py::TestIntraRequestReuse (3) PASSED
test_cross_attn_cache.py::TestAddressReuseAfterFree (4) PASSED
test_cross_attn_cache.py::TestLRUEviction (2) PASSED
test_cross_attn_cache.py::TestPutOverwrite (1) PASSED
test_cross_attn_cache.py::TestClear (1) PASSED
test_cross_attn_cache.py::TestProactiveStaleEviction (1) PASSED
======================== 12 passed in 1.08s ========================

The cross-attention KV cache in HeliosTransformer3DModel reuses the projected text-encoder K/V across a request's denoise steps. It was keyed solely by a storage fingerprint of the source tensor (data_ptr / shape / stride / dtype / device / _version), which is not unique across requests: once the source is freed, PyTorch's caching allocator may reuse its address for an unrelated tensor of the same shape, producing a false cache hit -- request B silently reused request A's cross-attention K/V (cross-prompt contamination), intermittent and hard to reproduce. Extract the cache into SourceTensorLRUCache and store a weakref to the source tensor alongside each cached value. On lookup, a dead weakref (source freed) turns an address-reuse collision into a miss instead of a false hit. Intra-request reuse is preserved (the source is held alive across all denoise steps); stale entries are evicted lazily on lookup and proactively on insert. clear_cross_attention_cache() is unchanged (still clears both caches); it remains available for explicit per-request eviction. Tests: tests/diffusion/models/helios/test_cross_attn_cache.py Signed-off-by: yancaocn <yancaochn@163.com>

vllm-omni-review-bot · 2026-09-23T07:30:22Z

This PR appears to belong to: docs/design/module/diffusion/offloader.md, docs/design/module/diffusion/diffusion_model_integration.md, docs/design/module/diffusion/index.md.

Module owners: @wtomin @RuixiangMa @david6666666

Routing: @wtomin via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @RuixiangMa via module of the changed files, module named in the PR description, semantic router; @david6666666 via module of the changed files, module named in the PR description

@yancaocn, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

vllm-omni-review-bot · 2026-09-23T11:01:34Z

Omni ReviewBot: superseded

The CI failure noted on 97fe096dc477 refers to an earlier head; the pull request now points at 7627003fe652.

vllm-omni-review-bot · 2026-09-23T20:42:14Z

Omni ReviewBot routing record

Assigned Direct under experiment vllm-omni-strict-5050-20260829.

vllm-omni-review-bot

Omni ReviewBot review

PR description

Helios reused projected cross-attention K/V across a request's denoise steps by keying only a source-tensor storage fingerprint (data_ptr, shape, stride, dtype, device, _version). After that source was freed, PyTorch's caching allocator could recycle the address for a later request with the same layout, so the next lookup silently returned the previous request's K/V. This change extracts the cache into SourceTensorLRUCache, stores a weakref beside each value, and treats a dead weakref as a miss, with lazy eviction on lookup and proactive eviction on insert. Intra-request reuse still hits while the encoder tensor stays alive; clear_cross_attention_cache() is unchanged.

Change flow

flowchart TD
  hs["[EXISTING] encoder_hidden_states across denoise steps"]:::existing
  xf["[CHANGED] HeliosTransformer3DModel projection and KV caches"]:::changed
  cache["[NEW] SourceTensorLRUCache fingerprint plus weakref"]:::new
  hit["[EXISTING] intra-request K/V reuse while source is live"]:::existing
  miss["[CHANGED] freed-source address reuse becomes a miss"]:::changed
  hs --> xf --> cache
  cache --> hit
  cache --> miss
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

No actionable findings.

vllm-omni-review-bot · 2026-09-24T14:55:33Z

Omni ReviewBot: superseded

The CI failure noted on 7627003fe652 refers to an earlier head; the pull request now points at 6201665da0b2.

vllm-omni-review-bot

Omni ReviewBot review

PR description

Helios was reusing projected cross-attention K/V across a request's denoise steps by fingerprinting the source tensor's storage (pointer, shape, stride, dtype, device, version). That fingerprint is not unique after the source is freed, so the caching allocator could give the same address to a later request and silently serve the previous prompt's K/V. The change extracts that map into SourceTensorLRUCache, stores a weakref to the source beside each value, and treats a dead weakref as a miss (with lazy eviction on lookup and proactive eviction on insert). Intra-request reuse still hits while the source stays alive; clear_cross_attention_cache() still wipes both caches.

Change flow

flowchart TD
  enc["[EXISTING] Per-request encoder hidden states"]:::existing
  model["[CHANGED] HeliosTransformer3DModel projection and KV caches"]:::changed
  cache["[NEW] SourceTensorLRUCache fingerprint plus weakref"]:::new
  denoise["[EXISTING] Intra-request denoise-step K/V reuse"]:::existing
  tests["[NEW] Address-reuse regression tests"]:::new
  enc --> model
  model --> cache
  cache --> denoise
  cache --> tests
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

See inline comments below.

vllm-omni-review-bot · 2026-09-25T20:59:55Z

+
+from vllm_omni.diffusion.models.helios.cross_attn_cache import SourceTensorLRUCache
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]


[P1] PR body names the new test module but reports no Test Result

The frozen PR description only lists Tests: tests/diffusion/models/helios/test_cross_attn_cache.py. It does not record a pytest command, marker expression, or result. This module is marked core_model, diffusion, and cpu (line 32), which is the L1/L2 every-PR selector in pyproject.toml (core_model: L1&L2 tests (run in each PR)). The command that would prove that guard collected this file is equivalent to pytest tests/diffusion/models/helios/test_cross_attn_cache.py -m "core_model and diffusion and cpu". The body shows neither that selector nor its result. .buildkite/** was not materialized here and git show HEAD:.buildkite/** could not be executed, so the exact uploaded step label could not be confirmed; the missing Test Result is still a contract miss against the L1/L2 core_model CPU diffusion collection.

vllm-omni-review-bot · 2026-09-25T21:01:57Z

Omni ReviewBot: finding feedback

[p1] PR body names the new test module but reports no Test Result — tests/diffusion/models/helios/test_cross_attn_cache.py:32

The frozen PR description only lists Tests: tests/diffusion/models/helios/test_cross_attn_cache.py. It does not record a pytest command, marker expression, or result. This module is marked core_model, diffusion, and cpu (line 32), which is the L1/L2 every-PR selector in pyproject.toml (core_model: L1&L2 tests (run in each PR)). The command that would prove that guard collected this file is equivalent to pytest tests/diffusion/models/helios/test_cross_attn_cache.py -m "core_model and diffusion and cpu". The body shows neither that selector nor its result. .buildkite/** was not materialized here and git show HEAD:.buildkite/** could not be executed, so the exact uploaded step label could not be confirmed; the missing Test Result is still a contract miss against the L1/L2 core_model CPU diffusion collection.

This finding appears in the bot's COMMENT review, but GitHub could not place it as an inline diff comment. If you are the PR author and disagree, react 👎 to this comment. The disagreement will be shown to the maintainer; it does not approve or merge the PR.

fhfuih

Any local test results on your machine? Preferrably you can attach a pair of before/after videos and their SSIM

yancaocn · 2026-09-29T14:54:49Z

@tzhouam pls review this PR,thanks.

tzhouam · 2026-09-30T02:15:10Z

@wtz2333 @tsinghua-code
We need to come up with an abstraction for the cross attn cache across models.

tzhouam

LGTM

yancaocn requested review from NickCao, wtomin and yenuo26 as code owners September 23, 2026 07:26

vllm-omni-review-bot reviewed Sep 23, 2026

View reviewed changes

hsliuustc0106 added the bug Something isn't working label Sep 24, 2026

Merge branch 'main' into helios-crossattn-cache-fix

7627003

vllm-omni-review-bot reviewed Sep 25, 2026

View reviewed changes

fhfuih reviewed Sep 26, 2026

View reviewed changes

Levius-Fubuki mentioned this pull request Sep 27, 2026

[Doc][Helios] Benchmark dense attention backends on H20 #8206

Open

Tao20060602 mentioned this pull request Sep 30, 2026

[RFC]: Deterministic Rollout Support for vLLM-Omni #4864

Open

Merge branch 'main' into helios-crossattn-cache-fix

6201665

tzhouam self-requested a review September 30, 2026 02:23

tzhouam approved these changes Sep 30, 2026

View reviewed changes

tzhouam added the ready label to trigger buildkite CI label Sep 30, 2026

Merge branch 'main' into helios-crossattn-cache-fix

d427ed1

tzhouam added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Sep 30, 2026

Jungle430 mentioned this pull request Sep 30, 2026

[Core] Add token_ids output type and enforce canonical names #7843

Merged

tzhouam merged commit 004a571 into vllm-project:main Sep 30, 2026
7 of 9 checks passed

tsinghua-code mentioned this pull request Oct 13, 2026

[Core] Introduce a generic cross-attention abstraction with LingBot integration #8859

Merged

Navigation Menu

Repository navigation

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

[Bugfix][Helios] Guard cross-attn KV cache against address reuse#8064
tzhouam merged 4 commits into
vllm-project:mainfrom
yancaocn:helios-crossattn-cache-fix

yancaocn commented Sep 23, 2026 •

edited

Loading

Uh oh!

vllm-omni-review-bot commented Sep 23, 2026

Uh oh!

vllm-omni-review-bot commented Sep 23, 2026 •

edited

Loading

Uh oh!

vllm-omni-review-bot commented Sep 23, 2026

Uh oh!

vllm-omni-review-bot left a comment

Uh oh!

vllm-omni-review-bot commented Sep 24, 2026 •

edited

Loading

Uh oh!

vllm-omni-review-bot left a comment

Uh oh!

vllm-omni-review-bot Sep 25, 2026

Uh oh!

vllm-omni-review-bot commented Sep 25, 2026

Uh oh!

fhfuih left a comment

Uh oh!

yancaocn commented Sep 29, 2026

Uh oh!

tzhouam commented Sep 30, 2026

Uh oh!

tzhouam left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants


		from vllm_omni.diffusion.models.helios.cross_attn_cache import SourceTensorLRUCache

		pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

Conversation

yancaocn commented Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vllm-omni-review-bot commented Sep 23, 2026

Uh oh!

vllm-omni-review-bot commented Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Omni ReviewBot: superseded

Uh oh!

vllm-omni-review-bot commented Sep 23, 2026

Uh oh!

vllm-omni-review-bot left a comment

Choose a reason for hiding this comment

Omni ReviewBot review

PR description

Change flow

Uh oh!

vllm-omni-review-bot commented Sep 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Omni ReviewBot: superseded

Uh oh!

vllm-omni-review-bot left a comment

Choose a reason for hiding this comment

Omni ReviewBot review

PR description

Change flow

Uh oh!

vllm-omni-review-bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

vllm-omni-review-bot commented Sep 25, 2026

Omni ReviewBot: finding feedback

Uh oh!

fhfuih left a comment

Choose a reason for hiding this comment

Uh oh!

yancaocn commented Sep 29, 2026

Uh oh!

tzhouam commented Sep 30, 2026

Uh oh!

tzhouam left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

yancaocn commented Sep 23, 2026 •

edited

Loading

vllm-omni-review-bot commented Sep 23, 2026 •

edited

Loading

vllm-omni-review-bot commented Sep 24, 2026 •

edited

Loading