support eager attention

configuration_aria.py

@@ -68,6 +68,8 @@ class AriaConfig(PretrainedConfig):
         self.ignore_index = ignore_index
         self.image_token_index = image_token_index
 
+        attn_implementation = kwargs.pop("attn_implementation", None)
+
         # Convert the keys and values of projector_patch_to_query_dict to integers
         # This ensures consistency even if they were provided as strings
         self.projector_patch_to_query_dict = {
@@ -76,10 +78,21 @@ class AriaConfig(PretrainedConfig):
 
         if isinstance(vision_config, dict) and "model_type" in vision_config:
             vision_config = AriaVisionConfig(**vision_config)
+            vision_attn_implementation = (
+                "flash_attention_2"
+                if attn_implementation is None
+                else attn_implementation
+            )
+            vision_config._attn_implementation = vision_attn_implementation
 
         self.vision_config = vision_config
 
         if isinstance(text_config, dict) and "model_type" in text_config:
+            text_attn_implementation = (
+                "sdpa" if attn_implementation is None else attn_implementation
+            )
             text_config = AriaMoELMConfig(**text_config)
+            text_config._attn_implementation = text_attn_implementation
+            print(text_config._attn_implementation)
 
         self.text_config = text_config

modeling_aria.py

@@ -133,7 +133,6 @@ class AriaForConditionalGeneration(AriaPretrainedModel):
     def __init__(self, config: AriaConfig):
         super().__init__(config)
 
-        config.vision_config._attn_implementation = config._attn_implementation
         self.vision_tower = AriaVisionModel(config.vision_config)
         self.multi_modal_projector = build_mm_projector(config)
         self.vocab_size = config.text_config.vocab_size

vision_encoder.py

@@ -82,6 +82,7 @@ class AriaVisionModel(SiglipVisionModel):
 
     config_class = AriaVisionConfig
     main_input_name = "pixel_values"
+    _supports_sdpa = False
 
     def __init__(self, config: AriaVisionConfig):
         super().__init__(config)