diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index edfe829b82f3..75d47ad47249 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1885,9 +1885,34 @@ def _prepare_static_cache( "offloading": offload_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 + # 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) 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. + # 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