Fix several bugs in cache_implementation=static

src/transformers/cache_utils.py

@@ -44,6 +44,7 @@ class CacheLayerMixin(ABC):
     """Base, abstract class for a single layer's cache."""
 
     is_compileable = False
+    supports_early_init = True
     # Subclasses can set ``layer_type`` to auto-register themselves in
     # ``LAYER_TYPE_CACHE_MAPPING`` at import time (used by ``DynamicCache`` to dispatch
     # per-layer cache classes from ``config.layer_types``).
@@ -386,12 +387,8 @@ def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tens
         devices, dtypes etc later on for each `update` (which could break the static dynamo addresses as well).
 
         If this is unwanted, one can call `early_initialization(...)` on the Cache directly, which will call this
-        function ahead-of-time (this is required for `torch.export` for example). Note that for `compile`, as we
-        internally don't compile the prefill, this is guaranteed to have been called already when compiling.
-        If compiling the prefill as well, e.g. calling `model.compile(...)` before `generate` with a static cache,
-        it is still supported in general, but without guarantees depending on the compilation options (e.g. cuda graphs,
-        i.e. `mode="reduce-overhead"` is known to fail). But it will in general work correctly, and prefill should
-        not be compiled anyway for performances!
+        function ahead-of-time (this is required for `torch.export` for example). It is also required whenever the
+        prefill itself ends up in a compiled region (with chunked prefill for instance).
         """
         self.dtype, self.device = key_states.dtype, key_states.device
         self.max_batch_size, self.num_heads = key_states.shape[:2]
@@ -857,6 +854,8 @@ class LinearAttentionCacheLayerMixin(ABC):
 
     # All shapes are static by essence in a LinearAttention layer, so it is compileable
     is_compileable = True
+    # Linear attention layers track their own conv/recurrent states; they don't use the key/value early-init path.
+    supports_early_init = False
 
     def __init__(self):
         self.conv_states: torch.Tensor | None = None
@@ -1239,6 +1238,8 @@ def early_initialization(
             )
 
         for layer, layer_num_heads, layer_head_dim in zip(self.layers, num_heads, head_dim):
+            if not layer.supports_early_init or layer.is_initialized:
+                continue
             # Note that the initialization needs all dimensions (except -2), as well as device and dtype, so we use
             # this fake tensor approach. It has size 0 on the -2 dimension, so it does not allocate any data (it only
             # creates an empty tensor with correct shape, dtype and device), which is very efficient and practical
@@ -1370,8 +1371,9 @@ def max_batch_size(self) -> int:
     @property
     def max_cache_len(self) -> int:
         """Return the maximum cache length of the cache"""
-        values = [layer.max_cache_len for layer in self.layers]
-        return max(values)
+        # Linear attention layers have no `max_cache_len`; skip them so a hybrid cache reports its attention layers'.
+        values = [layer.max_cache_len for layer in self.layers if hasattr(layer, "max_cache_len")]
+        return max(values) if values else 0
 
     @property
     def is_compileable(self) -> bool:
@@ -1384,7 +1386,8 @@ def is_compileable(self) -> bool:
     @property
     def is_initialized(self) -> bool:
         """Return whether the cache data is initialized"""
-        return len(self.layers) > 0 and all(layer.is_initialized for layer in self.layers)
+        layers = [layer for layer in self.layers if layer.supports_early_init]
+        return len(layers) > 0 and all(layer.is_initialized for layer in layers)
 
     @property
     def is_sliding(self) -> list[bool]:

src/transformers/generation/configuration_utils.py

@@ -173,6 +173,12 @@ class GenerationConfig(PushToHubMixin):
             our [cache documentation](https://huggingface.co/docs/transformers/en/kv_cache) for further information.
         cache_config (`dict`, *optional*, default to `None`):
             Arguments used in the key-value cache class can be passed in `cache_config`.
+        max_cache_len (`int`, *optional*):
+            Only used with static caches (`cache_implementation` set to `"static"` or `"offloaded_static"`).
+            Pre-sizes the cache to this length instead of the current call's `max_length`. Set it once to the
+            largest call you expect so that repeated `generate()` calls with a longer prompt or a larger
+            `max_new_tokens` (up to this ceiling) reuse the same cache instead of triggering a reallocation and a
+            `torch.compile` recompilation.
 
         > Parameters for manipulation of the model output logits
 
@@ -388,6 +394,7 @@ def __init__(self, **kwargs):
         self.use_cache = kwargs.pop("use_cache", None)
         self.cache_implementation = kwargs.pop("cache_implementation", None)
         self.cache_config = kwargs.pop("cache_config", None)
+        self.max_cache_len = kwargs.pop("max_cache_len", None)
 
         # Parameters for manipulation of the model output logits
         self.temperature = kwargs.pop("temperature", None)
@@ -654,6 +661,11 @@ def validate(self, strict=False, user_set_attributes: set[str] | None = None):
                 f"Invalid `cache_implementation` ({self.cache_implementation}). Choose one of: "
                 f"{valid_cache_implementations}"
             )
+        if self.max_cache_len is not None and self.cache_implementation not in ALL_STATIC_CACHE_IMPLEMENTATIONS:
+            logger.warning_once(
+                f"`max_cache_len` is only used with static caches ({STATIC_CACHE_IMPLEMENTATIONS}); it will be "
+                f"ignored with `cache_implementation={self.cache_implementation!r}`."
+            )
         # 1.3. Performance attributes
         if self.compile_config is not None and not isinstance(self.compile_config, CompileConfig):
             raise ValueError(

src/transformers/generation/utils.py

@@ -1784,8 +1784,33 @@ def _prepare_generation_config(
 
         return generation_config, model_kwargs
 
+    def _get_static_cache_init_shape(self: "GenerativePreTrainedModel") -> tuple[int, int] | None:
+        """
+        Returns the per-rank `(num_heads, head_dim)` to eagerly initialize a `StaticCache`, with the head count sharded
+        for tensor parallelism. Returns `None` when the cache cannot be early initialized.
+        """
+        if hasattr(self, "hf_device_map") and len(set(self.hf_device_map.values())) > 1:
+            # The model layers are on different devices
+            return None
+        text_config = self.config.get_text_config(decoder=True)
+        tp_size = getattr(self, "_tp_size", None) or 1
+        num_key_value_heads = getattr(text_config, "num_key_value_heads", None) or text_config.num_attention_heads
+        if num_key_value_heads % tp_size != 0:
+            # The model cannot be evenly sharded by head
+            return None
+        if getattr(text_config, "qk_head_dim", None) is not None:
+            # MLA models have distinct key (`qk_head_dim`) and value (`v_head_dim`) sizes.
+            return None
+        head_dim = getattr(text_config, "head_dim", None) or text_config.hidden_size // text_config.num_attention_heads
+        return num_key_value_heads // tp_size, head_dim
+
     def _prepare_static_cache(
-        self: "GenerativePreTrainedModel", cache_implementation: str, batch_size: int, max_cache_len: int, model_kwargs
+        self: "GenerativePreTrainedModel",
+        cache_implementation: str,
+        batch_size: int,
+        max_cache_len: int,
+        prefill_chunk_size: int | None,
+        model_kwargs,
     ) -> Cache:
         """
         Sets a cache for `generate`, that will persist across calls. A new cache will only be initialized a
@@ -1831,6 +1856,19 @@ def _prepare_static_cache(
                     "offloading": offload_cache,
                 }
                 self._cache = EncoderDecoderCache(self._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
+                    self._cache.early_initialization(
+                        batch_size=batch_size,
+                        num_heads=num_heads,
+                        head_dim=head_dim,
+                        dtype=self.dtype,
+                        device=self.device,
+                    )
         else:
             self._cache.reset()
         return self._cache
@@ -1929,10 +1967,18 @@ def _prepare_cache_for_generation(
                     f"and will be removed in v5.13. Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, "
                     "and the layer structure will be inferred automatically."
                 )
+            # `max_cache_len` lets the static cache be sized for the worst case across calls, so that later calls
+            # with a longer prompt or a larger `max_new_tokens` (up to that ceiling) reuse the same cache instead of
+            # triggering a reallocation (and a `torch.compile` recompilation). Without it, the cache is sized to the
+            # current call's `max_length` only. See #46424.
+            if generation_config.max_cache_len is not None:
+                max_cache_length = max(max_cache_length, generation_config.max_cache_len)
+            cache_batch_size = max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size
             model_kwargs[cache_name] = self._prepare_static_cache(
                 cache_implementation=generation_config.cache_implementation,
-                batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size,
+                batch_size=cache_batch_size,
                 max_cache_len=max_cache_length,
+                prefill_chunk_size=generation_config.prefill_chunk_size,
                 model_kwargs=model_kwargs,
             )
         elif generation_config.cache_implementation == "quantized":

src/transformers/masking_utils.py

@@ -1520,6 +1520,9 @@ def create_masks_for_generate(
     # case we return a single mask.
     if hasattr(effective_config, "layer_types"):
         layer_patterns = set(effective_config.layer_types)
+        # Without a registered attention-mask function, defer to the model by returning the raw attention mask
+        if any(layer_type not in LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING for layer_type in layer_patterns):
+            return attention_mask
         if len(layer_patterns) == 1:
             return LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING[next(iter(layer_patterns))](**mask_kwargs)
         causal_masks = {}

tests/utils/test_cache_utils.py

@@ -14,6 +14,7 @@
 
 import copy
 import unittest
+from unittest.mock import patch
 
 import pytest
 from packaging import version
@@ -55,6 +56,7 @@
         convert_and_export_with_cache,
         pipeline,
     )
+    from transformers.cache_utils import LinearAttentionLayer, StaticLayer
     from transformers.integrations.executorch import export_with_dynamic_cache
 
 
@@ -108,6 +110,48 @@ def _random_kvs(config):
         self.assertTrue(cached_keys.shape == (1, 1, 10, 128))
         self.assertTrue(cached_values.shape == (1, 1, 10, 128))
 
+    def test_early_initialization_does_not_corrupt_linear_attention_layers(self):
+        """
+        Regression test: `early_initialization` initialized *every* layer through the attention `(num_heads, head_dim)`
+        key/value layout, but linear attention layers track their own statically-shaped conv/recurrent states. So it
+        pre-allocated their `conv_states` with the wrong shape (a spurious 0-length axis) and flagged them
+        initialized; the first real `update_conv_state` then hit `conv_states.copy_(...)` with mismatched shapes and
+        raised a `RuntimeError` (the customer-facing failure, observed e.g. on Neuron). Linear attention layers must
+        be left untouched so they lazily take the correct shape on their first update.
+        """
+        config = LlamaConfig(num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, hidden_size=32)
+        config.layer_types = ["full_attention", "linear_attention"]
+        cache = StaticCache(config=config, max_cache_len=8)
+        self.assertIsInstance(cache.layers[0], StaticLayer)
+        linear_layer = cache.layers[1]
+        self.assertIsInstance(linear_layer, LinearAttentionLayer)
+        # Linear attention layers opt out of the key/value early-init path; attention layers opt in.
+        self.assertTrue(cache.layers[0].supports_early_init)
+        self.assertFalse(linear_layer.supports_early_init)
+
+        cache.early_initialization(batch_size=1, num_heads=2, head_dim=8, dtype=torch.float32, device=torch_device)
+        # The attention layer is initialized via the key/value layout, as expected.
+        self.assertTrue(cache.layers[0].is_initialized)
+
+        # Simulate the first forward updating the linear layer's conv state with its real (conv) shape. Before the
+        # fix, `early_initialization` had pre-allocated `conv_states` with the wrong key/value shape, so this raised
+        # `RuntimeError: ... must match ...`; with the fix the layer lazily takes the correct shape here instead.
+        conv_states = torch.zeros((1, 8, 4), dtype=torch.float32, device=torch_device)  # (batch, channels, kernel)
+        updated_conv_states = linear_layer.update_conv_state(conv_states)
+        self.assertEqual(updated_conv_states.shape, conv_states.shape)
+
+        # The cache-level `is_initialized` must also tolerate linear attention layers (which don't expose the flag);
+        # before the fix this raised `AttributeError`. It reflects the attention layer's state.
+        self.assertTrue(cache.is_initialized)
+
+    def test_max_cache_len_ignores_linear_attention_layers(self):
+        """`max_cache_len` must skip linear attention layers (which have no such attribute), else the static-cache
+        reuse check in `_prepare_static_cache` raises `AttributeError` on a hybrid model."""
+        config = LlamaConfig(num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, hidden_size=32)
+        config.layer_types = ["full_attention", "linear_attention"]
+        cache = StaticCache(config=config, max_cache_len=8)
+        self.assertEqual(cache.max_cache_len, 8)
+
 
 def _skip_on_failed_cache_prerequisites(test, cache_implementation):
     """Function to skip tests on failed cache prerequisites, given a cache implementation"""
@@ -488,6 +532,79 @@ def test_static_cache_no_cuda_graph_skips(self):
             model.generate(**inputs, max_new_tokens=2, cache_implementation="static")
         self.assertNotIn("cuda", cap.err.lower())
 
+    def test_static_cache_honors_max_cache_len(self):
+        """
+        `generation_config.max_cache_len` must size the auto-allocated `StaticCache` on the static path, so a cache
+        pinned to a worst-case length is reused across calls instead of reallocating (and recompiling). Without this,
+        the cache is sized to the current call's `max_length` only. Regression test for #46424.
+        """
+        # Llama has only full-attention layers, so the cache length is not capped by a sliding window.
+        model_repo = "hf-internal-testing/tiny-random-LlamaForCausalLM"
+        model = AutoModelForCausalLM.from_pretrained(model_repo).to(torch_device)
+        tokenizer = AutoTokenizer.from_pretrained(model_repo)
+        inputs = tokenizer(["The quick brown fox"], return_tensors="pt").to(torch_device)
+
+        max_new_tokens = 5
+        # Deliberately ask for a cache much larger than the natural `max_length - 1` for this call.
+        natural_max_cache_len = inputs.input_ids.shape[-1] + max_new_tokens - 1
+        requested_max_cache_len = natural_max_cache_len + 64
+
+        out = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            cache_implementation="static",
+            max_cache_len=requested_max_cache_len,
+            return_dict_in_generate=True,
+        )
+        self.assertIsInstance(out.past_key_values, StaticCache)
+        self.assertEqual(out.past_key_values.max_cache_len, requested_max_cache_len)
+
+    def test_chunked_prefill_initializes_static_cache_eagerly(self):
+        """
+        With chunked prefill the prefill runs inside a compiled region, where the static cache's lazy initialization
+        cannot tag its tensors via `torch._dynamo.mark_static_address`, so a later call recompiles. `generate` must
+        therefore initialize the static cache eagerly, and only when it compiles the prefill (i.e. when
+        `prefill_chunk_size` is set). Regression test for #46421.
+        """
+        model_repo = "hf-internal-testing/tiny-random-LlamaForCausalLM"
+        model = AutoModelForCausalLM.from_pretrained(model_repo).to(torch_device)
+        tokenizer = AutoTokenizer.from_pretrained(model_repo)
+        inputs = tokenizer(["The quick brown fox jumps over the lazy dog"], return_tensors="pt").to(torch_device)
+        generation_kwargs = {"max_new_tokens": 3, "do_sample": False, "cache_implementation": "static"}
+
+        # Without chunked prefill, the static cache is left to lazily initialize on the (eager) prefill.
+        with patch.object(Cache, "early_initialization", autospec=True) as eager_init:
+            model.generate(**inputs, **generation_kwargs)
+        eager_init.assert_not_called()
+
+        # Drop the (now initialized) cache so the chunked-prefill call re-allocates and initializes a fresh one.
+        model._cache = None
+        # With chunked prefill, `generate` must initialize the static cache itself, before the prefill runs.
+        prefill_chunk_size = max(inputs.input_ids.shape[-1] // 2, 1)
+        with patch.object(Cache, "early_initialization", autospec=True, wraps=Cache.early_initialization) as init:
+            model.generate(**inputs, **generation_kwargs, prefill_chunk_size=prefill_chunk_size)
+        init.assert_called_once()
+
+        # Under TP, eager init is still used but the head count is sharded by `_tp_size` (one device per rank).
+        tc = model.config.get_text_config(decoder=True)
+        num_kv_heads = getattr(tc, "num_key_value_heads", None) or tc.num_attention_heads
+        model._cache = None
+        model._tp_size = num_kv_heads  # divides evenly -> 1 head per rank
+        with patch.object(Cache, "early_initialization", autospec=True) as tp_init:
+            model.generate(**inputs, **generation_kwargs, prefill_chunk_size=prefill_chunk_size)
+        tp_init.assert_called_once()
+        self.assertEqual(tp_init.call_args.kwargs["num_heads"], 1)
+        model._tp_size = None
+
+        # A multi-device `device_map` (single `model.device` can't cover all layers) skips eager init -> lazy.
+        model._cache = None
+        model.hf_device_map = {"a": 0, "b": 1}
+        with patch.object(Cache, "early_initialization", autospec=True) as md_init:
+            model.generate(**inputs, **generation_kwargs, prefill_chunk_size=prefill_chunk_size)
+        md_init.assert_not_called()
+        del model.hf_device_map
+
     @require_torch_multi_accelerator
     @slow
     def test_static_cache_multi_accelerator(self):