4444 PreTrainedModel ,
4545)
4646from transformers .cache_utils import StaticLayer
47- from transformers .integrations .accelerate import init_empty_weights
4847from transformers .integrations .executorch import TorchExportableModuleForDecoderOnlyLM
4948from transformers .integrations .sdpa_attention import sdpa_attention_forward
5049from transformers .masking_utils import (
@@ -64,6 +63,9 @@ def _load(
6463 text_config .layer_types = text_config .layer_types [:layers ]
6564 kwargs = {"torch_dtype" : torch .float32 , "attn_implementation" : "sdpa" }
6665 if empty_weights :
66+ # transformers 5.11 no longer ships init_empty_weights; only --gguf needs it.
67+ from accelerate import init_empty_weights
68+
6769 with init_empty_weights ():
6870 model = AutoModelForCausalLM .from_config (config , ** kwargs )
6971 elif random_weights :
@@ -161,7 +163,9 @@ def _store_kv_as(self, model: PreTrainedModel, dtype: torch.dtype) -> None:
161163 if type (old ) is not StaticLayer :
162164 raise ValueError (f"cannot change the KV dtype of { type (old ).__name__ } " )
163165 layer = _CastingStaticLayer (old .max_cache_len )
164- layer .lazy_initialization (old .keys [:, :, :1 ].to (dtype ))
166+ layer .lazy_initialization (
167+ old .keys [:, :, :1 ].to (dtype ), old .values [:, :, :1 ].to (dtype )
168+ )
165169 self .cache .layers [i ] = layer
166170 self .cached .register_buffer (f"key_cache_{ i } " , layer .keys , persistent = False )
167171 self .cached .register_buffer (
@@ -176,6 +180,9 @@ def _store_kv_as(self, model: PreTrainedModel, dtype: torch.dtype) -> None:
176180 def forward (
177181 self , input_ids : torch .Tensor , cache_position : torch .Tensor
178182 ) -> torch .Tensor :
183+ # StaticLayer writes at cumulative_length and ignores cache_position.
184+ for layer in self .cache .layers :
185+ layer .cumulative_length = cache_position [:1 ].clone ()
179186 return self .cached .model (
180187 input_ids = input_ids ,
181188 cache_position = cache_position ,
0 commit comments