From 7d2a05a875957f4fa2b673402acab22f09977e8a Mon Sep 17 00:00:00 2001 From: ydshieh Date: Fri, 21 Aug 2026 22:01:28 +0200 Subject: [PATCH 1/7] [generate] Fix CUDAGraphs tensor overwrite for encoder-decoder static 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 --- src/transformers/generation/utils.py | 39 ++++++++++++++++++---------- 1 file changed, 26 insertions(+), 13 deletions(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index edfe829b82f3..d98160550b0a 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1885,19 +1885,32 @@ def _prepare_static_cache( "offloading": offload_cache, } cache = EncoderDecoderCache(cache, StaticCache(**cross_attention_cache_kwargs)) - elif prefill_chunk_size is not None: - # Chunked prefill compiles the prefill, so eagerly init the fresh cache to avoid a recompile next call - # (#46421). Skipped (-> lazy init) when it can't 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 - cache.early_initialization( - batch_size=batch_size, - num_heads=num_heads, - head_dim=head_dim, - dtype=self.dtype, - device=self.device, - ) + + # 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" + # For encoder-decoder models, 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. + # This also covers the chunked-prefill case (#46421), so the previous `elif prefill_chunk_size` block + # is folded in here. 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, + ) + if self.config.is_encoder_decoder: + cache.self_attention_cache.early_initialization(**early_init_kwargs) + cache.cross_attention_cache.early_initialization(**early_init_kwargs) + else: + cache.early_initialization(**early_init_kwargs) # Set the current length on the current model, to avoid recompilation later if we can self._previous_max_cache_length = effective_length From 0e9de38274eec88a0711260eb907a48564ae25ff Mon Sep 17 00:00:00 2001 From: ydshieh Date: Sat, 22 Aug 2026 11:21:07 +0200 Subject: [PATCH 2/7] Clarify comment: chunked-prefill block folded in, not specific to that case Co-Authored-By: Claude Sonnet 4.6 --- src/transformers/generation/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index d98160550b0a..b5bcb2fe88cb 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1894,8 +1894,9 @@ def _prepare_static_cache( # For encoder-decoder models, 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. - # This also covers the chunked-prefill case (#46421), so the previous `elif prefill_chunk_size` block - # is folded in here. Skipped when the cache cannot be initialized on a single device. + # Note: chunked prefill previously had its own `elif prefill_chunk_size` block for the same reason + # (#46421); it is folded here since the fix applies equally to all compiled generate calls. + # 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 From a611fca9e80188be9ff521bdb31bcc487ff32fb1 Mon Sep 17 00:00:00 2001 From: ydshieh Date: Sat, 22 Aug 2026 11:23:47 +0200 Subject: [PATCH 3/7] Fix issue reference: #46421 is the issue, #46446 is the PR that fixed it Co-Authored-By: Claude Sonnet 4.6 --- src/transformers/generation/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index b5bcb2fe88cb..82842f7b5b50 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1895,7 +1895,7 @@ def _prepare_static_cache( # 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. # Note: chunked prefill previously had its own `elif prefill_chunk_size` block for the same reason - # (#46421); it is folded here since the fix applies equally to all compiled generate calls. + # (#46446); it is folded here since the fix applies equally to all compiled generate calls. # 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: From 564e3fa65b04625d1ba9a62aaef023e6145d346e Mon Sep 17 00:00:00 2001 From: ydshieh Date: Sat, 22 Aug 2026 11:31:25 +0200 Subject: [PATCH 4/7] Clarify comment: replace 'folded in' with plain language Co-Authored-By: Claude Sonnet 4.6 --- src/transformers/generation/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 82842f7b5b50..95dd3263a712 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1895,7 +1895,8 @@ def _prepare_static_cache( # 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. # Note: chunked prefill previously had its own `elif prefill_chunk_size` block for the same reason - # (#46446); it is folded here since the fix applies equally to all compiled generate calls. + # (#46446); that block is removed and its logic absorbed here, since the fix applies equally to + # all compiled generate calls, not just chunked prefill. # 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: From f544b291dccbacb5f74cba3c878276b80593bd27 Mon Sep 17 00:00:00 2001 From: ydshieh Date: Sat, 22 Aug 2026 11:41:04 +0200 Subject: [PATCH 5/7] Scope fix to encoder-decoder only; restore elif prefill_chunk_size for decoder-only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- src/transformers/generation/utils.py | 59 ++++++++++++++++------------ 1 file changed, 34 insertions(+), 25 deletions(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 95dd3263a712..1a10aa058677 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1886,33 +1886,42 @@ def _prepare_static_cache( } cache = EncoderDecoderCache(cache, StaticCache(**cross_attention_cache_kwargs)) - # 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" - # For encoder-decoder models, 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. - # Note: chunked prefill previously had its own `elif prefill_chunk_size` block for the same reason - # (#46446); that block is removed and its logic absorbed here, since the fix applies equally to - # all compiled generate calls, not just chunked prefill. - # 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, - ) - if self.config.is_encoder_decoder: + 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) - else: - cache.early_initialization(**early_init_kwargs) + elif prefill_chunk_size is not None: + # Chunked prefill compiles the prefill, so eagerly init the fresh cache to avoid a recompile + # next call (#46446). Skipped (-> lazy init) when it can't 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 + cache.early_initialization( + batch_size=batch_size, + num_heads=num_heads, + head_dim=head_dim, + dtype=self.dtype, + device=self.device, + ) # Set the current length on the current model, to avoid recompilation later if we can self._previous_max_cache_length = effective_length From 892f3558949d6cf223d3a653a15905ddd32c36f0 Mon Sep 17 00:00:00 2001 From: ydshieh Date: Sun, 23 Aug 2026 22:13:09 +0200 Subject: [PATCH 6/7] fix --- src/transformers/generation/utils.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 1a10aa058677..4cb80517e7cb 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1900,13 +1900,13 @@ def _prepare_static_cache( 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, - ) + 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) elif prefill_chunk_size is not None: From 347b9fc91ca448598ee5aacd7d32014bdc1c8fb5 Mon Sep 17 00:00:00 2001 From: ydshieh Date: Mon, 24 Aug 2026 09:11:11 +0200 Subject: [PATCH 7/7] Fix Pyright type narrowing: assert EncoderDecoderCache in encoder-decoder 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 --- src/transformers/generation/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 4cb80517e7cb..75d47ad47249 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1887,6 +1887,7 @@ def _prepare_static_cache( 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 # 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