Rename supports_static_cache to can_compile_fullgraph (huggingface#39505)

* update all

* Apply suggestions from code review

Co-authored-by: Joao Gante <[email protected]>

* apply suggestions

* fix copies

---------

Co-authored-by: Joao Gante <[email protected]>

Author: 2 people

examples/modular-transformers/modeling_my_new_model2.py

@@ -294,7 +294,7 @@ class MyNewModel2PreTrainedModel(PreTrainedModel):
     _supports_flex_attn = True
     _supports_cache_class = True
     _supports_quantized_cache = True
-    _supports_static_cache = True
+    _can_compile_fullgraph = True
     _supports_attention_backend = True
     _can_record_outputs = {
         "hidden_states": MyNewModel2DecoderLayer,

examples/modular-transformers/modeling_new_task_model.py

@@ -94,7 +94,7 @@ class NewTaskModelPreTrainedModel(PreTrainedModel):
     _skip_keys_device_placement = "past_key_values"
     _supports_cache_class = True
     _supports_quantized_cache = True
-    _supports_static_cache = True
+    _can_compile_fullgraph = True
     _supports_flash_attn = True
     _supports_sdpa = True
     _supports_flex_attn = True

examples/modular-transformers/modeling_super.py

@@ -293,7 +293,7 @@ class SuperPreTrainedModel(PreTrainedModel):
     _supports_flex_attn = True
     _supports_cache_class = True
     _supports_quantized_cache = True
-    _supports_static_cache = True
+    _can_compile_fullgraph = True
     _supports_attention_backend = True
     _can_record_outputs = {
         "hidden_states": SuperDecoderLayer,

src/transformers/generation/utils.py

@@ -2059,7 +2059,7 @@ def _prepare_cache_for_generation(
         )
         if generation_config.cache_implementation is not None:
             if generation_config.cache_implementation in NEED_SETUP_CACHE_CLASSES_MAPPING:
-                if generation_config.cache_implementation == "static" and not self._supports_static_cache:
+                if generation_config.cache_implementation == "static" and not self._can_compile_fullgraph:
                     raise ValueError(
                         "This model does not support `cache_implementation='static'`. Please check the following "
                         "issue: https://github.com/huggingface/transformers/issues/28981"
@@ -2215,7 +2215,8 @@ def _valid_auto_compile_criteria(self, model_kwargs: dict, generation_config: Ge
         using_compilable_cache = (
             isinstance(model_kwargs.get("past_key_values"), Cache) and model_kwargs["past_key_values"].is_compileable
         )
-        can_compile = valid_hardware and using_compilable_cache and self._supports_static_cache
+        # TODO @raushan `self._can_compile_fullgraph` can be removed and inferred from model arch (e.g. MoE doesn't support compile)
+        can_compile = valid_hardware and using_compilable_cache and self._can_compile_fullgraph
 
         # Exception 1: Some quantization methods do not support compilation
         if getattr(self, "hf_quantizer", None) is not None:

src/transformers/modeling_utils.py

@@ -2063,8 +2063,7 @@ class PreTrainedModel(nn.Module, EmbeddingAccessMixin, ModuleUtilsMixin, PushToH
     # Flex Attention support
     _supports_flex_attn = False
 
-    # Has support `torch.compile(fullgraph=True)`
-    _supports_static_cache = False
+    _can_compile_fullgraph = False
 
     # A tensor parallel plan to be applied to the model when TP is enabled. For
     # top-level models, this attribute is currently defined in respective model

src/transformers/models/arcee/modeling_arcee.py

@@ -313,7 +313,7 @@ class ArceePreTrainedModel(PreTrainedModel):
     _supports_sdpa = True
     _supports_flex_attn = True
 
-    _supports_static_cache = True
+    _can_compile_fullgraph = True
     _supports_attention_backend = True
     _can_record_outputs = {
         "hidden_states": ArceeDecoderLayer,

src/transformers/models/aria/modeling_aria.py

@@ -654,7 +654,7 @@ class AriaPreTrainedModel(PreTrainedModel):
     _supports_flash_attn = True
     _supports_sdpa = True
     _supports_flex_attn = True
-    _supports_static_cache = False  # MoE models don't work with torch.compile (dynamic slicing)
+    _can_compile_fullgraph = False  # MoE models don't work with torch.compile (dynamic slicing)
     _supports_attention_backend = True
     _can_record_outputs = {
         "hidden_states": AriaTextDecoderLayer,

src/transformers/models/aria/modular_aria.py

@@ -1302,7 +1302,7 @@ def _init_weights(self, module):
 class AriaPreTrainedModel(LlamaPreTrainedModel):
     config: AriaConfig
     base_model_prefix = ""
-    _supports_static_cache = False  # MoE models don't work with torch.compile (dynamic slicing)
+    _can_compile_fullgraph = False  # MoE models don't work with torch.compile (dynamic slicing)
     _supports_attention_backend = True
 
     def _init_weights(self, module):

src/transformers/models/aya_vision/modeling_aya_vision.py

@@ -96,7 +96,7 @@ class AyaVisionPreTrainedModel(PreTrainedModel):
 
     _supports_flash_attn = True
     _supports_sdpa = True
-    _supports_static_cache = False
+    _can_compile_fullgraph = False
     _supports_flex_attn = True
     _supports_attention_backend = True
 

src/transformers/models/aya_vision/modular_aya_vision.py

@@ -90,7 +90,7 @@ def pixel_shuffle(self, image_features):  # B, S, D
 
 
 class AyaVisionPreTrainedModel(LlavaPreTrainedModel):
-    _supports_static_cache = False
+    _can_compile_fullgraph = False
 
 
 class AyaVisionCausalLMOutputWithPast(LlavaCausalLMOutputWithPast):