[Ernie 4.5 VL Moe] Post merge adjustments (#43117)

post merge fixes

Author: vasqu

src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py

@@ -1705,6 +1705,8 @@ def prepare_inputs_for_generation(
         past_key_values=None,
         image_grid_thw=None,
         video_grid_thw=None,
+        use_cache=True,
+        is_first_iteration=False,
         # Intentionally ignore position ids to force custom cache logic
         position_ids=None,
         **kwargs,
@@ -1717,6 +1719,8 @@ def prepare_inputs_for_generation(
             past_key_values=past_key_values,
             image_grid_thw=image_grid_thw,
             video_grid_thw=video_grid_thw,
+            use_cache=use_cache,
+            is_first_iteration=is_first_iteration,
             **kwargs,
         )
 
@@ -1732,7 +1736,7 @@ def prepare_inputs_for_generation(
             mm_token_type_ids=model_inputs.get("mm_token_type_ids"),
         )
 
-        if model_inputs["cache_position"][0] != 0:
+        if not is_first_iteration and use_cache:
             model_inputs["pixel_values"] = None
             model_inputs["pixel_values_videos"] = None
             model_inputs["mm_token_type_ids"] = None

src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py

@@ -1393,6 +1393,8 @@ def prepare_inputs_for_generation(
         past_key_values=None,
         image_grid_thw=None,
         video_grid_thw=None,
+        use_cache=True,
+        is_first_iteration=False,
         # Intentionally ignore position ids to force custom cache logic
         position_ids=None,
         **kwargs,
@@ -1405,6 +1407,8 @@ def prepare_inputs_for_generation(
             past_key_values=past_key_values,
             image_grid_thw=image_grid_thw,
             video_grid_thw=video_grid_thw,
+            use_cache=use_cache,
+            is_first_iteration=is_first_iteration,
             **kwargs,
         )
 
@@ -1420,7 +1424,7 @@ def prepare_inputs_for_generation(
             mm_token_type_ids=model_inputs.get("mm_token_type_ids"),
         )
 
-        if model_inputs["cache_position"][0] != 0:
+        if not is_first_iteration and use_cache:
             model_inputs["pixel_values"] = None
             model_inputs["pixel_values_videos"] = None
             model_inputs["mm_token_type_ids"] = None

tests/models/ernie4_5_vl_moe/test_modeling_ernie4_5_vl_moe.py

@@ -313,6 +313,7 @@ def load_model(self, dtype, attn_implementation="sdpa"):
             device_map="auto",
             dtype=dtype,
             attn_implementation=attn_implementation,
+            experts_implementation="eager",
             revision="refs/pr/10",
         )
 
@@ -549,6 +550,7 @@ def load_model(self, dtype, attn_implementation="sdpa"):
             device_map="auto",
             dtype=dtype,
             attn_implementation=attn_implementation,
+            experts_implementation="eager",
         )
 
     def test_small_model_integration_test(self):