System Info

transformers 5.10.1 (latest release)

Cache.early_initialization(batch_size, num_heads, head_dim, dtype, device) pre-allocates all cache layers ahead of time, assuming there is a helper registered to precompute the cache shape available for the layer type.

But there are no registered helpers for linear-attention layers (mamba/conv/linear_attention, LinearAttentionCacheLayerMixin) because they manage their own statically-shaped conv/recurrent states, whose shapes (conv kernel, SSM state size, …) aren't derivable from (num_heads, head_dim). Feeding them a (batch, num_heads, 0, head_dim) tensor silently allocates a wrongly-shaped conv_states (note the spurious 0-length axis) and flags them initialized; the first real update_conv_state(...) then does conv_states.copy_(...) with mismatched shapes and raises RuntimeError.

Relatedly, Cache.is_initialized does all(layer.is_initialized for layer in self.layers), which AttributeErrors on a hybrid cache because linear-attention layers have no is_initialized attribute.

Reproduction

import torch
from transformers import LlamaConfig, StaticCache

cfg = LlamaConfig(num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, hidden_size=32)
cfg.layer_types = ["full_attention", "linear_attention"]
cache = StaticCache(config=cfg, max_cache_len=8)

cache.early_initialization(batch_size=1, num_heads=2, head_dim=8, dtype=torch.float32, device="cpu")
print(cache.layers[1].conv_states.shape)          # torch.Size([1, 2, 0, 8])  <- wrong (0-length axis)
cache.layers[1].update_conv_state(torch.zeros((1, 8, 4)))
# RuntimeError: The size of tensor a (8) must match the size of tensor b (4) at non-singleton dimension 3

Proposed fix

In early_initialization, read the flag defensively (getattr(layer, "is_initialized", None)) and skip layers where it is None (layers that don't use the key/value lazy-init, e.g. linear attention) or already True (idempotent on a reused cache). Make Cache.is_initialized ignore those layers so a hybrid cache reflects its attention layers. (Fix available; comes with a model-free regression test.)

Expected behavior

early_initialization should initialize attention layers and leave linear-attention layers to lazily initialize from their real states on first update; cache.is_initialized should not crash on a hybrid cache.

Who can help?

@ArthurZucker @Cyrilvallez @vasqu