[generate] Fix CUDAGraphs tensor overwrite for encoder-decoder static cache by ydshieh · Pull Request #48199 · huggingface/transformers

ydshieh · 2026-08-21T20:02:49Z

What

Fixes T5ModelIntegrationTests::test_compile_static_cache which fails with:

RuntimeError: Error: accessing tensor output of CUDAGraphs that has been
overwritten by a subsequent run.

Root cause

_prepare_static_cache must call early_initialization on cache buffers
before any CUDAGraph capture so that mark_static_address is applied
while is_torchdynamo_compiling() is False. Without it, lazy_initialization
runs inside the capture context, mark_static_address is skipped, and
CUDAGraphs pools the tensor memory — overwriting it on the next replay.

For encoder-decoder models both the self-attention and cross-attention
caches need early init. The self-attention cache is straightforward, but the
cross-attention cache is the key issue: its K/V tensors are computed from
encoder hidden states on the first decode step, which happens inside the
compiled graph. So lazy_initialization runs inside graph capture, and
mark_static_address is never called.

Regression introduced by #47731 ("Stop setting the static cache as an
attribute to save memory"): before that PR, _prepare_static_cache reused
self._cache across generate calls, so lazy_initialization had already
run outside any compiled context on the first call. After #47731, a fresh
StaticCache is always created, making eager initialization necessary.

Decoder-only models (e.g. afmoe) — also affected, not yet fixed here

The same regression from #47731 also affects decoder-only models with
test_compile_static_cache. Before #47731, self._cache was reused across
generate calls, so cumulative_length (and other buffers) always had the
same tensor address. After #47731, a fresh StaticCache is created on every
call, so cumulative_length is at a different memory address each time —
CUDAGraph replay references the old address and triggers the overwrite error.

Confirmed: afmoe::test_compile_static_cache fails with the same
RuntimeError: Error: accessing tensor output of CUDAGraphs on the Aug 05 CI
commit (9f66415aec). The existing elif prefill_chunk_size block does not
cover this case.

This PR only fixes the encoder-decoder path (T5). The decoder-only fix
needs a separate look — would like @Cyrilvallez to have a first look before
going further.

Changes

_prepare_static_cache: add an if self.config.is_encoder_decoder branch
that calls early_initialization on both self_attention_cache and
cross_attention_cache before graph capture.
The existing elif prefill_chunk_size decoder-only block is unchanged.

Test plan

T5ModelIntegrationTests::test_compile_static_cache — verified PASS on A10G runner (torch 2.13, sm_86)
All 14 T5ModelIntegrationTests pass

ydshieh · 2026-08-21T20:06:14Z

run-slow: t5

github-actions · 2026-08-21T20:07:42Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/t5"]
quantizations: []

HuggingFaceDocBuilderDev · 2026-08-21T20:15:26Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

github-actions · 2026-08-21T20:21:32Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	7601fa82	workflow commit (merge commit)
PR	6a63cb10	branch commit (from PR)
main	e453228e	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

… cache `_prepare_static_cache` now calls `early_initialization` on BOTH the self-attention and cross-attention caches for encoder-decoder models. Previously only the self-attention cache was initialized eagerly. For encoder-decoder models (e.g. T5), cross-attention K/V tensors are computed from encoder hidden states on the first decode step, which happens inside the CUDAGraph capture context. When `lazy_initialization` runs there, `is_torchdynamo_compiling()` is True so `mark_static_address` is skipped. CUDAGraphs then pools the tensor memory addresses, causing: RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run. Regression introduced by #47731 ("Stop setting the static cache as an attribute to save memory"), which changed `_prepare_static_cache` to always create a fresh `StaticCache` instead of reusing `self._cache`. Before that PR, the same cache object was reused across `generate` calls and was already initialized from the first call (outside any compiled context), so `lazy_initialization` never ran inside capture. Also folds the decoder-only early-init (previously only for chunked prefill, #46421) into the same unconditional block, since the same CUDAGraph issue applies to any compiled generate call. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

ydshieh · 2026-08-21T21:24:21Z

run-slow: auto, bert, gpt2, t5, modernbert, vit, clip, detr, table_transformer, got_ocr2, whisper, wav2vec2, qwen2_audio, speecht5, csm, llama, gemma3, qwen2, mistral3, qwen2_5_vl, llava, smolvlm, internvl, gemma3n, gpt_oss, qwen2_5_omni, pi0

github-actions · 2026-08-21T21:25:36Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/auto", "models/bert", "models/clip", "models/csm", "models/detr", "models/gemma3", "models/gemma3n", "models/got_ocr2", "models/gpt2", "models/gpt_oss", "models/internvl", "models/llama", "models/llava", "models/mistral3", "models/modernbert", "models/pi0", "models/qwen2", "models/qwen2_5_omni", "models/qwen2_5_vl", "models/qwen2_audio", "models/smolvlm", "models/speecht5", "models/t5", "models/table_transformer", "models/vit", "models/wav2vec2", "models/whisper"]
quantizations: []

github-actions · 2026-08-21T22:14:31Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	ca036c1a	workflow commit (merge commit)
PR	7d2a05a8	branch commit (from PR)
main	d56c55bf	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

ydshieh · 2026-08-22T09:16:23Z

+                dtype=self.dtype,
+                device=self.device,
+            )
+            if self.config.is_encoder_decoder:


Why touch instead of overriding in ?

Overriding _prepare_static_cache in T5 is technically feasible, but the bug is not T5-specific: any encoder-decoder model using static_cache + torch.compile would hit the same issue (cross-attention K/V are allocated inside the graph capture, so mark_static_address is never called). T5 is just the only encoder-decoder with a test_compile_static_cache CI test today.

Putting the fix here means BART, Whisper, and any future encoder-decoder model get the correct behaviour automatically, instead of requiring the same override to be copy-pasted into each. The elif prefill_chunk_size block that was already here is the exact same pattern for decoder-only models (#46421), so this just extends it symmetrically to the encoder-decoder case.

…t case Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…r decoder-only The previous commit incorrectly removed the `elif prefill_chunk_size` guard for decoder-only models. For decoder-only models without chunked prefill, lazy_initialization runs during the (non-compiled) prefill, so mark_static_address is correctly applied outside graph capture — early init is not needed there. Restore the original elif prefill_chunk_size block for decoder-only models, and scope the new encoder-decoder early init as a separate if branch. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

ydshieh · 2026-08-22T09:43:25Z

run-slow: t5

github-actions · 2026-08-22T09:44:53Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/t5"]
quantizations: []

github-actions · 2026-08-22T09:57:46Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	63a7770d	workflow commit (merge commit)
PR	f544b291	branch commit (from PR)
main	d56c55bf	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

ydshieh · 2026-08-23T20:11:07Z

+
+        if self.config.is_encoder_decoder:
+            # Eagerly initialize the K/V buffers so that `mark_static_address` is applied before any
+            # `torch.compile` / CUDAGraph capture context. Without this, `lazy_initialization` runs inside
+            # graph capture (where `is_torchdynamo_compiling()` is True), causing `mark_static_address` to
+            # be skipped. CUDAGraphs then treats those tensors as pooled memory and can overwrite them
+            # between runs, raising:
+            #   "accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run"
+            # BOTH self-attention and cross-attention caches need early init: cross-attention K/V are
+            # computed from encoder hidden states on the first decode step, which happens inside the
+            # compiled graph, so they suffer the same lazy-init-inside-capture problem.
+            # Skipped when the cache cannot be initialized on a single device.
+            init_shape = self._get_static_cache_init_shape()
+            if init_shape is not None:
+                num_heads, head_dim = init_shape
+                early_init_kwargs = dict(
+                    batch_size=batch_size,
+                    num_heads=num_heads,
+                    head_dim=head_dim,
+                    dtype=self.dtype,
+                    device=self.device,
+                )
+                cache.self_attention_cache.early_initialization(**early_init_kwargs)
+                cache.cross_attention_cache.early_initialization(**early_init_kwargs)


@Cyrilvallez Let me know if you still want to have this fix only in T5 modeling code. I think the issue is there for all encoder-decoder models, but among them, only T5 has test_compile_static_cache.

ydshieh · 2026-08-23T20:13:34Z

run-slow: t5

github-actions · 2026-08-23T20:14:49Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/t5"]
quantizations: []

github-actions · 2026-08-23T20:24:10Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	57d7c0c8	workflow commit (merge commit)
PR	892f3558	branch commit (from PR)
main	3d6c7201	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

…oder branch Pyright infers `cache` as `StaticCache | EncoderDecoderCache` after the conditional reassignment, so accessing `.self_attention_cache` / `.cross_attention_cache` inside the guarded `if self.config.is_encoder_decoder` block raises unresolved-attribute errors. Adding an `assert isinstance` narrows the type without changing runtime behaviour. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

ydshieh · 2026-08-24T07:24:29Z

            cache = EncoderDecoderCache(cache, StaticCache(**cross_attention_cache_kwargs))
+
+        if self.config.is_encoder_decoder:
+            assert isinstance(cache, EncoderDecoderCache)  # always true: assigned above when is_encoder_decoder


this isinstance(cache, EncoderDecoderCache) is to avoid Pyright (check_types.py) complain (fail)

github-actions · 2026-08-24T07:25:34Z

CI recap

Dashboard: View test results in Grafana
Latest run: 32700407065:1
Result: failure | Jobs: 16 | Tests: 132,869 | Failures: 1 | Duration: 13h 59m

ydshieh · 2026-09-07T13:06:34Z

gently ping @Cyrilvallez again to help us to have all important models have ✅ tests 🙏

vasqu · 2026-09-08T16:05:14Z

+                    "device": self.device,
+                }
+                cache.self_attention_cache.early_initialization(**early_init_kwargs)
+                cache.cross_attention_cache.early_initialization(**early_init_kwargs)


I assume enc-dec always have the same shapes for head dim and num heads? Would investigate on real models

Wouldn't it make much more sense to adjust the encoder-decoder cache's implementation of early_initialization

For 1. I could also see some adjustments for _get_static_cache_init_shape which might be related to #48619 (meaning we could combine something using the new utils introduced there)

vasqu · 2026-09-08T16:09:19Z

cc @CyrilBurner

CyrilBurner

Are you sure?? Prefill is never compiled so should not be an issue

CyrilBurner · 2026-09-09T05:47:51Z

+            # Eagerly initialize the K/V buffers so that `mark_static_address` is applied before any
+            # `torch.compile` / CUDAGraph capture context. Without this, `lazy_initialization` runs inside
+            # graph capture (where `is_torchdynamo_compiling()` is True), causing `mark_static_address` to
+            # be skipped. CUDAGraphs then treats those tensors as pooled memory and can overwrite them
+            # between runs, raising:
+            #   "accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run"
+            # BOTH self-attention and cross-attention caches need early init: cross-attention K/V are
+            # computed from encoder hidden states on the first decode step, which happens inside the
+            # compiled graph, so they suffer the same lazy-init-inside-capture problem.
+            # Skipped when the cache cannot be initialized on a single device.
+            init_shape = self._get_static_cache_init_shape()
+            if init_shape is not None:
+                num_heads, head_dim = init_shape
+                early_init_kwargs = {
+                    "batch_size": batch_size,
+                    "num_heads": num_heads,
+                    "head_dim": head_dim,
+                    "dtype": self.dtype,
+                    "device": self.device,
+                }
+                cache.self_attention_cache.early_initialization(**early_init_kwargs)
+                cache.cross_attention_cache.early_initialization(**early_init_kwargs)


I don't get it, I believe all this comment is wrong since prefill is never compiled - since it's an AI comment, I believe the AI probably hallucinated it

Thanks for the note! To clarify: "the compiled graph" in the comment refers to the compiled decoder step, not the encoder. The encoder runs eagerly, but the cross-attention K/V computation (and lazy_initialization) happens on the first decoder step, which is inside the compiled graph — that's where mark_static_address gets skipped.

vasqu · 2026-09-17T13:15:02Z

cc @Cyrilvallez for the redirect :D

but yes this is a valid issue because the decoder part of the enc-dec models are not properly accounted for (even if the AI comment is heavy lol)

Cyrilvallez · 2026-09-18T07:30:55Z

Superseded by #48923

ydshieh force-pushed the fix_t5_compile_static_cache branch from 6a63cb1 to 7d2a05a Compare August 21, 2026 21:17

ydshieh commented Aug 22, 2026

View reviewed changes

ydshieh and others added 4 commits August 22, 2026 11:21

Clarify comment: chunked-prefill block folded in, not specific to tha…

0e9de38

…t case Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Fix issue reference: #46421 is the issue, #46446 is the PR that fixed it

a611fca

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Clarify comment: replace 'folded in' with plain language

564e3fa

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

ydshieh commented Aug 23, 2026

View reviewed changes

fix

892f355

ydshieh commented Aug 24, 2026

View reviewed changes

ydshieh requested a review from Cyrilvallez August 26, 2026 10:59

ydshieh mentioned this pull request Aug 28, 2026

Implement VibeVoice #40546

Merged

sergereview Bot mentioned this pull request Sep 5, 2026

[serge] Fix 2 integration tests for model glm4_moe failing with OOM (other (2)) #48551

Merged

vasqu reviewed Sep 8, 2026

View reviewed changes

CyrilBurner reviewed Sep 9, 2026

View reviewed changes

Cyrilvallez closed this Sep 18, 2026

Conversation

ydshieh commented Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Root cause

Decoder-only models (e.g. afmoe) — also affected, not yet fixed here

Changes

Test plan

Uh oh!

ydshieh commented Aug 21, 2026

Uh oh!

github-actions Bot commented Aug 21, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Aug 21, 2026

Uh oh!

github-actions Bot commented Aug 21, 2026

CI Results

Commit Info

Uh oh!

ydshieh commented Aug 21, 2026

Uh oh!

github-actions Bot commented Aug 21, 2026

Uh oh!

github-actions Bot commented Aug 21, 2026

CI Results

Commit Info

Uh oh!

ydshieh Aug 22, 2026

Choose a reason for hiding this comment

Uh oh!

ydshieh commented Aug 22, 2026

Uh oh!

github-actions Bot commented Aug 22, 2026

Uh oh!

github-actions Bot commented Aug 22, 2026

CI Results

Commit Info

Uh oh!

ydshieh Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

ydshieh commented Aug 23, 2026

Uh oh!

github-actions Bot commented Aug 23, 2026

Uh oh!

github-actions Bot commented Aug 23, 2026

CI Results

Commit Info

Uh oh!

ydshieh Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Aug 24, 2026

CI recap

Uh oh!

ydshieh commented Sep 7, 2026

Uh oh!

vasqu Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu commented Sep 8, 2026

Uh oh!

CyrilBurner left a comment

Choose a reason for hiding this comment

Uh oh!

CyrilBurner Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

ydshieh Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu commented Sep 17, 2026

Uh oh!

Cyrilvallez commented Sep 18, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

ydshieh commented Aug 21, 2026 •

edited

Loading

ydshieh Aug 24, 2026 •

edited

Loading