Repository: huggingface/transformers
Pull request: fix: HiggsAudioV2 cached decode inputs in compiled generation

# What does this PR do?

main is currently failing with 

```
FAILED tests/models/higgs_audio_v2/test_modeling_higgs_audio_v2.py::HiggsAudioV2ModelTest::test_generate_compilation_all_outputs - AssertionError: Lists differ: [torch.Size([2, 15, 32])] != [(2, 1, 32)]
```

likely introduced by [1618d44b92](https://github.com/huggingface/transformers/commit/1618d44b92) or a more recent conflictual change (I did not bisect)

Fix: Changed to slice last token when `use_cache=True`

I am not 100% sure my fix is right


Changed files:
filename: src/transformers/models/higgs_audio_v2/generation_higgs_audio_v2.py
patch: @@ -265,18 +265,21 @@ def _sample(
             else self.__call__
         )
 
-        # Assisted generation completes the prefill stage in candidate generator so that
-        # we don't have several `prefill` calls in one generation loop. Skip `_prefill` for assistants
-        if not generation_config.is_assistant:
-            outputs = self._prefill(input_ids, generation_config, model_kwargs)
-            prefill_consumed = False
-        else:
-            model_kwargs = self._get_initial_cache_position(input_ids.shape[1], input_ids.device, model_kwargs)
-            prefill_consumed = True
+        prefill_consumed = False
+        outputs = self._prefill(
+            input_ids,
+            generation_config,
+            model_kwargs,
+            is_first_iteration=not generation_config.is_assistant,
+        )
 
         while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
             if prefill_consumed:
-                model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+                use_cache = model_kwargs.get("use_cache", True)
+                new_input_ids = input_ids[:, -1:] if use_cache else input_ids
+                model_inputs = self.prepare_inputs_for_generation(
+                    new_input_ids, full_input_ids=input_ids, **model_kwargs
+                )
                 with self._optimize_model_for_decode():
                     outputs = model_forward(**model_inputs, return_dict=True)
             prefill_consumed = True

filename: src/transformers/models/higgs_audio_v2/modeling_higgs_audio_v2.py
patch: @@ -633,12 +633,13 @@ def prepare_inputs_for_generation(
         audio_input_ids_mask: torch.LongTensor | None = None,
         **kwargs,
     ):
+        full_input_ids = kwargs.pop("full_input_ids", input_ids)
         model_inputs = super().prepare_inputs_for_generation(input_ids, **kwargs)
 
         if audio_input_ids is not None and model_inputs.get("past_key_values") is not None:
             current_cache_length = model_inputs["cache_position"][0]
-            audio_token_mask = (input_ids == self.config.audio_token_id) | (
-                input_ids == self.config.audio_delay_token_id
+            audio_token_mask = (full_input_ids == self.config.audio_token_id) | (
+                full_input_ids == self.config.audio_delay_token_id
             )
             in_cache_num_audio_input_ids = audio_token_mask[:, :current_cache_length].sum(dim=-1)
 

filename: src/transformers/models/higgs_audio_v2/modular_higgs_audio_v2.py
patch: @@ -515,12 +515,13 @@ def prepare_inputs_for_generation(
         audio_input_ids_mask: torch.LongTensor | None = None,
         **kwargs,
     ):
+        full_input_ids = kwargs.pop("full_input_ids", input_ids)
         model_inputs = super().prepare_inputs_for_generation(input_ids, **kwargs)
 
         if audio_input_ids is not None and model_inputs.get("past_key_values") is not None:
             current_cache_length = model_inputs["cache_position"][0]
-            audio_token_mask = (input_ids == self.config.audio_token_id) | (
-                input_ids == self.config.audio_delay_token_id
+            audio_token_mask = (full_input_ids == self.config.audio_token_id) | (
+                full_input_ids == self.config.audio_delay_token_id
             )
             in_cache_num_audio_input_ids = audio_token_mask[:, :current_cache_length].sum(dim=-1)