Since `transformers` v4.56.0` the dictionary `ALL_STATIC_CACHE_IMPLEMENTATIONS` replaced the existing dictionary

modeling_decilm.py

@@ -27,7 +27,8 @@ import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers import GenerationConfig
-from transformers.generation.utils import NEED_SETUP_CACHE_CLASSES_MAPPING, GenerationMixin, GenerateOutput
+from transformers.generation.utils import GenerationMixin, GenerateOutput
+from transformers.generation.configuration_utils import ALL_STATIC_CACHE_IMPLEMENTATIONS
 from transformers.modeling_utils import PreTrainedModel
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 from transformers.utils import (
@@ -810,7 +811,7 @@ class DeciLMPreTrainedModel(PreTrainedModel):
         # DeciLM-specific code
         generation_config, model_kwargs = super()._prepare_generation_config(generation_config, *args, **kwargs)
         generation_config.cache_implementation = "variable"
-        NEED_SETUP_CACHE_CLASSES_MAPPING["variable"] = VariableCache
+        ALL_STATIC_CACHE_IMPLEMENTATIONS["variable"] = VariableCache
         return generation_config, model_kwargs